import gradio as gr
import torch
from PIL import Image
import numpy as np
import json
from huggingface_hub import hf_hub_download
import timm
from torchvision import transforms

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- Download and Load Classification Model (EfficientNet-B3) ---
CLASS_REPO_ID = "sheikh987/efficientnet-b3-skin"
CLASS_MODEL_FILENAME = "efficientnet_b3_skin_model.pth"
NUM_CLASSES = 7

try:
    class_model_path = hf_hub_download(
        repo_id=CLASS_REPO_ID,
        filename=CLASS_MODEL_FILENAME,
        cache_dir="/tmp"
    )
    checkpoint = torch.load(class_model_path, map_location=DEVICE)

    # Auto-detect state_dict vs full model
    if isinstance(checkpoint, dict):
        classification_model = timm.create_model(
            'efficientnet_b3', pretrained=False, num_classes=NUM_CLASSES
        ).to(DEVICE)
        classification_model.load_state_dict(checkpoint, strict=False)
    else:
        classification_model = checkpoint.to(DEVICE)

    classification_model.eval()
    print("✅ Classification model loaded successfully.")

except Exception as e:
    raise gr.Error(f"Failed to load the classification model: {e}")


# --- Load Knowledge Base ---
try:
    with open('knowledge_base.json', 'r') as f:
        knowledge_base = json.load(f)
except FileNotFoundError:
    raise gr.Error("knowledge_base.json not found. Upload it to the Space.")

idx_to_class_abbr = {0: 'MEL', 1: 'NV', 2: 'BCC', 3: 'AKIEC', 4: 'BKL', 5: 'DF', 6: 'VASC'}

# --- Image Transform ---
transform_classify = transforms.Compose([
    transforms.Resize((300, 300)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


# --- Pipeline Function ---
def classify_image(input_image):
    if input_image is None:
        return None, "Please upload an image."

    class_input_tensor = transform_classify(input_image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        logits = classification_model(class_input_tensor)
        probs = torch.nn.functional.softmax(logits, dim=1)

    confidence, predicted_idx = torch.max(probs, 1)
    confidence_percent = confidence.item() * 100
    predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
    info = knowledge_base.get(predicted_abbr, {})

    # Build output
    info_text = (
        f"**Predicted Condition:** {info.get('full_name', 'N/A')} ({predicted_abbr})\n"
        f"**Confidence:** {confidence_percent:.2f}%\n\n"
        f"**Description:**\n{info.get('description', 'No description available.')}\n\n"
        f"**Common Causes:**\n" + "\n".join([f"• {c}" for c in info.get('causes', ['N/A'])]) + "\n\n"
        f"**Common Treatments:**\n" + "\n".join([f"• {t}" for t in info.get('common_treatments', ['N/A'])]) + "\n\n"
        f"**--- IMPORTANT DISCLAIMER ---**\n{info.get('disclaimer', '')}"
    )

    return input_image, info_text


# --- Gradio Interface ---
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload Skin Image"),
    outputs=[gr.Image(type="pil", label="Input Image"),
             gr.Markdown(label="Analysis Result")],
    title="AI Skin Lesion Classifier",
    description="Upload a skin lesion image and the AI EfficientNet-B3 model will classify it.\n\n"
                "**DISCLAIMER:** This is NOT a diagnosis. Always consult a qualified dermatologist."
)

if __name__ == "__main__":
    iface.launch()