"""
🏥 Skin Lesion Classification — HuggingFace Space
Models are in the same Space directory.
Compatible with Gradio 6.0+
"""

import os
import json
import numpy as np
import cv2
import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from datetime import datetime

# ════════════════════════════════════════════
# 📥 LOAD MODEL (Local — same Space directory)
# ════════════════════════════════════════════

# HuggingFace Spaces store files in the root directory
MODEL_PATH = "model_b3.keras"
CONFIG_PATH = "final_config.json"

# Fallback: check if files are in subdirectory
if not os.path.exists(MODEL_PATH):
    MODEL_PATH = os.path.join(os.path.dirname(__file__), "model_b3.keras")
if not os.path.exists(CONFIG_PATH):
    CONFIG_PATH = os.path.join(os.path.dirname(__file__), "final_config.json")

print(f"📂 Model path: {MODEL_PATH}")
print(f"📂 Config path: {CONFIG_PATH}")

with open(CONFIG_PATH, "r") as f:
    config = json.load(f)

IMG_SIZE = config["img_size"]
THRESHOLD = config["threshold"]

model = load_model(MODEL_PATH)
print(f"✅ Model loaded | Threshold: {THRESHOLD:.4f} | IMG_SIZE: {IMG_SIZE}")

# ════════════════════════════════════════════
# 🧠 CORE FUNCTIONS
# ════════════════════════════════════════════

def generate_gradcam(mdl, img_arr):
    grad_model = tf.keras.models.Model(
        inputs=mdl.input,
        outputs=[mdl.get_layer('top_activation').output, mdl.output]
    )
    with tf.GradientTape() as tape:
        conv_out, preds = grad_model(img_arr)
        loss = preds[0]
    grads = tape.gradient(loss, conv_out)
    pooled = tf.reduce_mean(grads, axis=(0, 1, 2))
    heatmap = tf.squeeze(conv_out[0] @ pooled[..., tf.newaxis])
    heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-8)
    return heatmap.numpy()

def predict_fast(mdl, img_batch):
    return float(mdl.predict(img_batch, verbose=0)[0][0])

def predict_tta(mdl, img_batch, n_augments=10):
    preds = [mdl.predict(img_batch, verbose=0)[0][0]]
    for _ in range(n_augments):
        aug = img_batch.copy()
        if np.random.random() > 0.5: aug = np.flip(aug, axis=2)
        if np.random.random() > 0.5: aug = np.flip(aug, axis=1)
        aug = np.clip(aug * np.random.uniform(0.85, 1.15), 0, 255)
        preds.append(float(mdl.predict(aug.copy(), verbose=0)[0][0]))
    return np.mean(preds)

PRECAUTIONS = {
    "benign": {
        "title": "✅ BENIGN — Low Risk",
        "summary": "Lesion appears benign. Regular monitoring recommended.",
        "items": [
            "📋 Monthly self-skin examinations",
            "🔍 Monitor for changes in size, shape, color",
            "📸 Photograph periodically to track changes",
            "🧴 SPF 30+ sunscreen daily",
            "🩺 Dermatologist check-up every 6-12 months",
            "⚠️ See doctor immediately if lesion changes"
        ],
        "extra_title": "⚠️ WHEN TO WORRY:",
        "extra": ["🔴 Rapid growth", "🔴 Color changes", "🔴 Irregular borders",
                   "🔴 Bleeding or itching", "🔴 New lesions nearby"]
    },
    "malignant": {
        "title": "⚠️ MALIGNANT — High Risk",
        "summary": "IMMEDIATE medical consultation strongly recommended.",
        "items": [
            "🚨 Consult dermatologist IMMEDIATELY",
            "🏥 Schedule biopsy for definitive diagnosis",
            "📋 Do NOT self-treat",
            "📸 Document with photos + size reference",
            "🧬 Ask about genetic testing if family history",
            "☀️ Avoid sun exposure on affected area",
            "🩺 Request full-body skin examination"
        ],
        "extra_title": "🔍 ABCDE RULE — Signs of Melanoma:",
        "extra": ["A — Asymmetry", "B — Irregular Border", "C — Multiple Colors",
                   "D — Diameter > 6mm", "E — Evolving shape/size"]
    }
}

# ════════════════════════════════════════════
# 🎯 MAIN PREDICTION
# ════════════════════════════════════════════

def predict_skin_lesion(input_image, patient_name, patient_age, patient_gender, mode):
    if input_image is None:
        return None, None, "⚠️ Please upload an image."

    patient_name = patient_name if patient_name and patient_name.strip() else "Anonymous"
    patient_age = patient_age if patient_age and patient_age > 0 else "N/A"
    patient_gender = patient_gender or "Not specified"

    img_resized = input_image.resize((IMG_SIZE, IMG_SIZE))
    img_array = np.array(img_resized, dtype=np.float32)
    img_batch = np.expand_dims(img_array, axis=0)

    if mode == "🎯 Best Mode (TTA — Higher Accuracy, Slower)":
        prediction = predict_tta(model, img_batch, n_augments=10)
        mode_label = "🎯 Best Mode (TTA x10)"
        mode_stats = "AUC: 0.935 | Recall: 94.7% | Accuracy: 83.8%"
    else:
        prediction = predict_fast(model, img_batch)
        mode_label = "⚡ Fast Mode (Single)"
        mode_stats = "AUC: 0.911 | Recall: 90.0% | Accuracy: 81.7%"

    is_malignant = prediction >= THRESHOLD
    label = "MALIGNANT" if is_malignant else "BENIGN"
    confidence = float(prediction if is_malignant else 1 - prediction)

    # Grad-CAM
    heatmap = generate_gradcam(model, img_batch)
    hm = cv2.resize(heatmap, (IMG_SIZE, IMG_SIZE))
    hm_color = cv2.applyColorMap(np.uint8(255 * hm), cv2.COLORMAP_JET)
    hm_color = cv2.cvtColor(hm_color, cv2.COLOR_BGR2RGB)
    img_np = cv2.resize(np.array(img_resized), (IMG_SIZE, IMG_SIZE))
    overlay = np.uint8(hm_color * 0.4 + img_np * 0.6)

    # Output images
    border_color = (255, 0, 0) if is_malignant else (0, 180, 0)
    out = cv2.copyMakeBorder(img_np, 8, 8, 8, 8, cv2.BORDER_CONSTANT, value=border_color)
    bar = np.zeros((50, out.shape[1], 3), dtype=np.uint8); bar[:] = border_color
    cv2.putText(bar, f"{label} ({confidence:.1%})", (10, 35),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 2)
    output_img = np.vstack([bar, out])

    gc = cv2.copyMakeBorder(overlay, 8, 8, 8, 8, cv2.BORDER_CONSTANT, value=(255, 165, 0))
    bar_gc = np.zeros((50, gc.shape[1], 3), dtype=np.uint8); bar_gc[:] = (255, 165, 0)
    cv2.putText(bar_gc, "GRAD-CAM: Model Focus", (10, 35),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    gradcam_img = np.vstack([bar_gc, gc])

    # Report
    key = "malignant" if is_malignant else "benign"
    p = PRECAUTIONS[key]
    report = [
        "=" * 50, "🏥 SKIN LESION ANALYSIS REPORT", "=" * 50, "",
        f"👤 {patient_name} | Age: {patient_age} | {patient_gender}",
        f"📅 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        f"⚙️ {mode_label}", f"   Performance: {mode_stats}", "",
        "─" * 50,
        f"🔬 Diagnosis: {label}",
        f"   Confidence: {confidence:.2%}",
        f"   Risk: {'HIGH ⚠️' if is_malignant else 'LOW ✅'}",
        f"   Score: {prediction:.4f} | Threshold: {THRESHOLD:.4f}", "",
        "─" * 50,
        f"{'🔴' if is_malignant else '🟢'} {p['title']}",
        f"📝 {p['summary']}", "",
        "🛡️ PRECAUTIONS:",
        *[f"  {i+1}. {item}" for i, item in enumerate(p['items'])], "",
        f"{p['extra_title']}",
        *[f"  {item}" for item in p['extra']], "",
        "─" * 50,
        "🔥 GRAD-CAM: 🔴 Red=Focus | 🟡 Yellow=Moderate | 🔵 Blue=Ignore", "",
        "═" * 50,
        "⚕️ DISCLAIMER: Educational/screening only.",
        "   Always consult a qualified dermatologist.",
        "═" * 50
    ]

    return output_img, gradcam_img, "\n".join(report)

# ════════════════════════════════════════════
# 🚀 INTERFACE (Gradio 6.0 Compatible)
# ════════════════════════════════════════════

with gr.Blocks(title="🏥 Skin Lesion Classifier") as demo:

    gr.Markdown("""
    # 🏥 Skin Lesion Classification — AI Screening Tool
    ### EfficientNetB3 + Grad-CAM | AUC: 0.935 | Recall: 94.7%
    > ⚠️ Educational/screening only. Consult a dermatologist.
    ---
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 📤 Upload & Patient Info")
            input_image = gr.Image(type="pil", label="📷 Skin Lesion Image", height=300)
            gr.Markdown("### 👤 Patient Details")
            patient_name = gr.Textbox(label="Name", placeholder="Patient name...", max_lines=1)
            patient_age = gr.Number(label="Age", minimum=0, maximum=120, precision=0)
            patient_gender = gr.Dropdown(label="Gender",
                                          choices=["Male", "Female", "Other", "Prefer not to say"])
            gr.Markdown("### ⚙️ Prediction Mode")
            mode = gr.Radio(label="Select Mode",
                             choices=["⚡ Fast Mode (Single Prediction — Quick)",
                                      "🎯 Best Mode (TTA — Higher Accuracy, Slower)"],
                             value="⚡ Fast Mode (Single Prediction — Quick)")
            gr.Markdown("""
            | Mode | AUC | Recall | Speed |
            |---|---|---|---|
            | ⚡ Fast | 0.911 | 90.0% | ~2s |
            | 🎯 Best | 0.935 | 94.7% | ~20s |
            """)
            with gr.Row():
                predict_btn = gr.Button("🔬 Analyze", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="lg")

        with gr.Column(scale=1):
            gr.Markdown("## 🔍 Results")
            with gr.Row():
                output_image = gr.Image(label="📊 Prediction", height=250)
                gradcam_image = gr.Image(label="🔥 Grad-CAM", height=250)
            gr.Markdown("### 📋 Report & Precautions")
            report_output = gr.Textbox(label="📝 Report", lines=22, max_lines=50)

    gr.Markdown("""
    ---
    | 🔴 Red = High Focus | 🟡 Yellow = Moderate | 🔵 Blue = Low |
    |---|---|---|
    ---
    > 🏥 Always consult a medical professional.
    """)

    predict_btn.click(predict_skin_lesion,
                       [input_image, patient_name, patient_age, patient_gender, mode],
                       [output_image, gradcam_image, report_output])
    clear_btn.click(lambda: (None, "", None, None, None, None, "",
                              "⚡ Fast Mode (Single Prediction — Quick)"), [],
                     [input_image, patient_name, patient_age, patient_gender,
                      output_image, gradcam_image, report_output, mode])

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