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

# Load your 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"

# Define cream recommendation logic
def recommend_creams(skin_tone, disease):
    recommendations = {
        "acne": {
            "light": [
                "/content/acne_cream.jpg",
                "/content/acne_cream.jpg",
                "/content/leek-1.jpg"
            ],
            "dark": [
                "/content/fairandlovely.jpg",
                "/content/acne_cream.jpg",
                "/content/leek-1.jpg"
            ],
        },
        "redness": {
            "light": [
                "/content/redness_cream_light.jpg",
                "/content/redness_cream_generic.jpg"
            ],
            "dark": [
                "/content/redness_cream_dark.jpg",
                "/content/redness_cream_generic.jpg"
            ],
        },
    }

    # Return a list of cream images based on skin tone and disease
    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, your skin is dark colored, and it shows signs of acne. Consider using a targeted acne treatment.",
        ("light", "acne"): "Hello, your skin is light colored, and it shows signs of acne. A gentle cream for acne would be beneficial.",
        ("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!",
    }

    return messages.get((skin_tone, disease), "No specific recommendations available.")

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)
    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="Disease", lines=1, interactive=False)
            message_display = gr.Textbox(label="Message", lines=2, interactive=False)

    # 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)