import gradio as gr
import torch
from PIL import Image
import torchvision.transforms as transforms
from ultralytics import YOLO

# Load models and map them to the CPU
skin_tone_model = torch.load("SkinTone.pth", map_location=torch.device('cpu'))
disease_model = YOLO("last.pt")

CONFIDENCE_THRESHOLD = 0.5

# Define image transformations for the skin tone model
transform_skin_tone = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Define prediction function for skin tone model
def predict_skin_tone(image):
    image_tensor = transform_skin_tone(image).unsqueeze(0)
    with torch.no_grad():
        prediction = skin_tone_model(image_tensor)
        _, predicted_class = torch.max(prediction, 1)
        skin_tone_classes = ['dark', 'light']  
        return skin_tone_classes[predicted_class.item()]

def predict_disease(image):
    results = disease_model(image)
    detected_classes = []

    for result in results:
        boxes = result.boxes
        if boxes is not None:
            for box in boxes:
                if box.conf > CONFIDENCE_THRESHOLD:
                    class_id = int(box.cls)
                    detected_classes.append(result.names[class_id])

    return ', '.join(detected_classes) if detected_classes else "No disease detected"

# cream recommendation logic
def recommend_creams(skin_tone, disease):
    recommendations = {
        "No disease detected": {
            "light": [
                "Face10.jpg",
                "Face08.jpg",
                "Face07.jpg",
                "Face06.jpg",
                "Face05.jpg"
            ],

            "dark": [
                "Face04.jpg",
                "Face03.jpg",
                "Face02.jpg",
                "Face01.jpg",
                "Face09.jpg"
            ]
        },
        
        "acne": {
            "light": [
                "Face01.jpg",
                "Face02.jpg",
                "Face03.jpg",
                "Face04.jpg",
                "Face05.jpg"
            ],
            "dark": [
                "Face06.jpg",
                "Face07.jpg",
                "Face08.jpg",
                "Face09.jpg",
                "Face10.jpg"
                
            ],

        },
        "redness": {
            "light": [
                "Face03.jpg",
                "Face04.jpg",
                "Face08.jpg",
                "Face10.jpg",
                "Face06.jpg"
            ],
            "dark": [
                "Face07.jpg",
                "Face02.jpg",
                "Face01.jpg",
                "Face09.jpg"
            ],
        },

    }

    # Return a list of cream images based on skin tone and disease
    disease = disease.split(', ')[0] 
    if disease in recommendations:
        return recommendations[disease].get(skin_tone, ["No specific cream recommended."])

    return []

# Define message generation based on predictions
def generate_message(skin_tone, disease):
    messages = {
        ("dark", "acne"): "Hello, ඔබගේ සම අදුරු පැහැති සමක්. ඒ වගේම ඔබගේ මුහුණේ කුරුලෑ දකින්න ලැබෙනවා. ඔබගේ සම තවත් පැහැපත් කරගන්න අතරතුර කරදරකාරී කුරුලෑ වලිනුත් මිදීමට පහත ඇති ඔබ වෙනුවෙන්ම නිෂ්පාදනය කළ අපගේ Face ක්‍රීම් භාවිතා කරන්න. ඔබට ස්තූතීයි! .",
        ("light", "acne"): "Hello, ඔබට ඇත්තේ පැහැපත් සමක්. ඒ වගේම ඔබගේ මුහුණේ කුරුලෑ දකින්න ලැබෙනවා. ඔබගේ සම තවත් පැහැපත් කරගන්න අතරතුර කරදරකාරී කුරුලෑ වලිනුත් මිදීමට පහත ඇති ඔබ වෙනුවෙන්ම නිෂ්පාදනය කළ අපගේ Face ක්‍රීම් භාවිතා කරන්න. ඔබට ස්තූතීයි!  .",
        ("dark", "redness"): "Your skin is dark colored, with redness detected. You might want to use a cream that soothes redness.",
        ("light", "redness"): "Your skin is light colored, with redness detected. A soothing cream can help alleviate this.",
        ("dark", "No disease detected"): "Your skin is dark colored, and there are no visible issues. Keep up the good skincare!",
        ("light", "No disease detected"): "Your skin is light colored, and there are no visible issues. Maintain your skincare routine!",
    }
    disease = disease.split(', ')[0]
    message = messages.get((skin_tone, disease), "No specific recommendations available.")
    
    return f"<span style='font-size: 18px; color: #4A90E2; font-weight: bold;'>{message}</span>"

def analyze_image(image):
    if image is None:
        return "No image captured. Please capture an image.", "No disease detected", [], ""

    skin_tone = predict_skin_tone(image)
    disease = predict_disease(image)
    recommended_creams = recommend_creams(skin_tone, disease)

#   return messages.get((skin_tone, disease), "No specific recommendations available.")
    message = generate_message(skin_tone, disease)

    return f"Skin Tone: {skin_tone}", f"Disease: {disease}", recommended_creams, message

def clear_interface():
    return "", "", [], ""

# Gradio interface for capturing an image
with gr.Blocks(theme='earneleh/paris') as interface:
    gr.Markdown(
        """
        <h1 style='text-align: center; color: #4A90E2;'>Skin Analysis Application</h1>
        <p style='text-align: center; color: #333;'>Hello!</p>
        """,
        elem_id="header"
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Capture Image", type="pil")
            analyze_button = gr.Button("Analyze", elem_id="analyze-btn")
            clear_button = gr.Button("Clear", elem_id="clear-btn")

        with gr.Column():
            skin_tone_display = gr.Textbox(label="Skin Tone", lines=1, interactive=False)
            disease_display = gr.Textbox(label="identified Special Charateristics", lines=1, interactive=False)
    #       message_display = gr.Textbox(label="Message", lines=2, interactive=False)
            message_display = gr.HTML(label="Message")

    # Using a gallery to display recommended creams
    recommended_creams_display = gr.Gallery(
        label="Recommended Creams",
        show_label=True,
        elem_id="cream-gallery"
    )

    analyze_button.click(
        analyze_image,
        inputs=image_input,
        outputs=[skin_tone_display, disease_display, recommended_creams_display, message_display]
    )

    clear_button.click(
        clear_interface,
        outputs=[skin_tone_display, disease_display, recommended_creams_display, message_display]
    )

# Launch the app
interface.launch(debug=True)