# app.py

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

# Import necessary model libraries
import segmentation_models_pytorch as smp
import timm
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torchvision import transforms

# --- 1. SETUP: Download and Load all models and data ---

print("--> Initializing application and downloading models...")
DEVICE = "cpu"

# --- Download and Load Segmentation Model (UNet) ---
try:
    SEG_REPO_ID = "sheikh987/unet-isic2018"
    SEG_MODEL_FILENAME = "unet_full_data_best_model.pth"
    print(f"--> Downloading segmentation model from: {SEG_REPO_ID}")
    seg_model_path = hf_hub_download(repo_id=SEG_REPO_ID, filename=SEG_MODEL_FILENAME)
    
    segmentation_model = smp.Unet(encoder_name="resnet34", encoder_weights=None, in_channels=3, classes=1).to(DEVICE)
    segmentation_model.load_state_dict(torch.load(seg_model_path, map_location=DEVICE))
    segmentation_model.eval()
    print("    Segmentation model loaded successfully.")

except Exception as e:
    print(f"!!! ERROR loading segmentation model: {e}")
    raise gr.Error("Failed to load the segmentation model. Check repository name and file paths.")


# --- Download and Load Classification Model (EfficientNet) ---
try:
    CLASS_REPO_ID = "sheikh987/efficientnet-isic"
    CLASS_MODEL_FILENAME = "efficientnet_augmented_best.pth"
    print(f"--> Downloading classification model from: {CLASS_REPO_ID}")
    class_model_path = hf_hub_download(repo_id=CLASS_REPO_ID, filename=CLASS_MODEL_FILENAME)

    NUM_CLASSES = 7
    classification_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=NUM_CLASSES).to(DEVICE)
    classification_model.load_state_dict(torch.load(class_model_path, map_location=DEVICE))
    classification_model.eval()
    print("    Classification model loaded successfully.")

except Exception as e:
    print(f"!!! ERROR loading classification model: {e}")
    raise gr.Error("Failed to load the classification model. Check repository name and file paths.")


# --- Load Knowledge Base and Labels ---
try:
    with open('knowledge_base.json', 'r') as f:
        knowledge_base = json.load(f)
    print("--> Knowledge base loaded.")
except FileNotFoundError:
    raise gr.Error("knowledge_base.json not found. Make sure it has been uploaded to the Space.")

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

# --- Define Image Transforms ---
transform_segment = A.Compose([
    A.Resize(height=256, width=256),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0),
    ToTensorV2(),
])
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]),
])

print("\n--> Application ready to accept requests.")


# --- 2. DEFINE THE FULL PIPELINE FUNCTION (UPDATED) ---

def full_pipeline(input_image):
    if input_image is None:
        return None, None, "Please upload an image."

    image_np = np.array(input_image.convert("RGB"))

    # STAGE 1: SEGMENTATION
    augmented_seg = transform_segment(image=image_np)
    seg_input_tensor = augmented_seg['image'].unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        seg_logits = segmentation_model(seg_input_tensor)
    seg_mask = (torch.sigmoid(seg_logits) > 0.5).float().squeeze().cpu().numpy()
    
    if seg_mask.sum() < 200:
        return None, None, "Analysis Failed: No lesion could be clearly identified."

    # STAGE 2: CROP and CLASSIFY
    rows = np.any(seg_mask, axis=1)
    cols = np.any(seg_mask, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    
    padding = 15
    rmin, rmax = max(0, rmin - padding), min(image_np.shape[0], rmax + padding)
    cmin, cmax = max(0, cmin - padding), min(image_np.shape[1], cmax + padding)

    cropped_image_pil = Image.fromarray(image_np[rmin:rmax, cmin:cmax])

    class_input_tensor = transform_classify(cropped_image_pil).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        class_logits = classification_model(class_input_tensor)
        probabilities = torch.nn.functional.softmax(class_logits, dim=1)
    
    confidence, predicted_idx = torch.max(probabilities, 1)
    confidence_percent = confidence.item() * 100
    
    # SAFETY NET
    CONFIDENCE_THRESHOLD = 60.0
    if confidence_percent < CONFIDENCE_THRESHOLD:
        inconclusive_text = (
            f"**Analysis Inconclusive**\n\n"
            f"The AI model's confidence ({confidence_percent:.2f}%) is below the required threshold of {CONFIDENCE_THRESHOLD}%.\n\n"
            "This can happen if the image is blurry, has poor lighting, or shows a condition the model was not trained on.\n\n"
            "**--- IMPORTANT DISCLAIMER ---**\n"
            "This is NOT a diagnosis. Please consult a qualified dermatologist for an accurate assessment."
        )
        mask_display = Image.fromarray((seg_mask * 255).astype(np.uint8))
        return mask_display, cropped_image_pil, inconclusive_text

    # --- STAGE 3: LOOKUP and FORMAT (UPDATED) ---
    predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
    info = knowledge_base.get(predicted_abbr, {})
    
    # Format the 'causes' and 'treatments' lists into clean, bulleted strings
    causes_list = info.get('causes', ['Specific causes not listed.'])
    causes_text = "\n".join([f"• {c}" for c in causes_list])

    treatments_list = info.get('common_treatments', ['No specific treatments listed.'])
    treatments_text = "\n".join([f"• {t}" for t in treatments_list])

    # Build the final output text using all the information
    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{causes_text}\n\n"
        f"**Common Treatments:**\n{treatments_text}\n\n"
        f"**--- IMPORTANT DISCLAIMER ---**\n{info.get('disclaimer', '')}"
    )
    
    mask_display = Image.fromarray((seg_mask * 255).astype(np.uint8))
    
    return mask_display, cropped_image_pil, info_text


# --- 3. CREATE AND LAUNCH THE GRADIO INTERFACE ---
iface = gr.Interface(
    fn=full_pipeline,
    inputs=gr.Image(type="pil", label="Upload Skin Image"),
    outputs=[
        gr.Image(type="pil", label="Segmentation Mask"),
        gr.Image(type="pil", label="Cropped Lesion"),
        gr.Markdown(label="Analysis Result")
    ],
    title="AI Skin Lesion Analyzer",
    description="This tool performs a two-stage analysis on a skin lesion image. **Stage 1:** A UNet model segments the lesion. **Stage 2:** An EfficientNet model classifies the segmented lesion. \n\n**DISCLAIMER:** This is an educational tool and is NOT a substitute for professional medical advice. Always consult a qualified dermatologist for any health concerns.",
    allow_flagging="never"
)

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