import torch
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr
import io
import base64

MODEL_NAME = "ahishamm/vit-base-HAM-10000-sharpened-patch-32"
feature_extractor = ViTFeatureExtractor.from_pretrained(MODEL_NAME)
model = ViTForImageClassification.from_pretrained(MODEL_NAME)
model.eval()

CLASSES = [
    "Queratosis actínica / Bowen",      # 0
    "Carcinoma células basales",        # 1
    "Lesión queratósica benigna",       # 2
    "Dermatofibroma",                   # 3
    "Melanoma maligno",                 # 4
    "Nevus melanocítico",               # 5
    "Lesión vascular"                   # 6
]
RISK_LEVELS = {
    0: {'level': 'Moderado', 'color': '#ffaa00', 'weight': 0.6},
    1: {'level': 'Alto', 'color': '#ff4444', 'weight': 0.8},
    2: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
    3: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
    4: {'level': 'Crítico', 'color': '#cc0000', 'weight': 1.0},
    5: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
    6: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1}
}

def analizar_lesion_vit_web(img):
    inputs = feature_extractor(img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = outputs.logits.softmax(dim=-1).cpu().numpy()[0]
    pred_idx = int(np.argmax(probs))
    pred_clase = CLASSES[pred_idx]
    confianza = probs[pred_idx]
    cancer_risk_score = sum(probs[i] * RISK_LEVELS[i]['weight'] for i in range(7))
    melanoma_risk = probs[4]
    bcc_risk = probs[1]
    precancer_risk = probs[0]
    benign_total = sum(probs[i] for i in [2,3,5,6])

    colors_bars = [RISK_LEVELS[i]['color'] for i in range(7)]
    fig, ax = plt.subplots(figsize=(8,3))
    ax.bar(CLASSES, probs*100, color=colors_bars)
    ax.set_title("Probabilidad por tipo de lesión")
    ax.set_ylabel("Probabilidad (%)")
    ax.set_xticklabels(CLASSES, rotation=45, ha='right')
    ax.grid(axis='y', alpha=0.2)
    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)
    img_bytes = buf.getvalue()
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
    html_chart = f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'

    urgency = ""
    recommendation = ""
    timeframe = ""
    if cancer_risk_score > 0.6:
        urgency = "🚨 <b>CRÍTICO</b>"
        recommendation = "Derivación INMEDIATA a oncología dermatológica"
        timeframe = "En 24-48 horas máximo"
    elif cancer_risk_score > 0.4:
        urgency = "⚠️ <b>ALTO RIESGO</b>"
        recommendation = "Consulta urgente con dermatólogo especialista"
        timeframe = "En 1 semana máximo"
    elif cancer_risk_score > 0.2:
        urgency = "📋 <b>RIESGO MODERADO</b>"
        recommendation = "Evaluación dermatológica programada"
        timeframe = "En 2-4 semanas"
    else:
        urgency = "✅ <b>BAJO RIESGO</b>"
        recommendation = "Seguimiento de rutina"
        timeframe = "En 3-6 meses"

    informe = f"""
    <div style="font-family:sans-serif; max-width:700px; margin:auto">
    <h2>🏥 INFORME DE ANÁLISIS DERMATOLÓGICO CON IA</h2>
    <b>📊 Diagnóstico principal:</b> {pred_clase}<br>
    <b>🎯 Confianza del modelo:</b> {confianza:.1%}<br>
    <b>📈 Score de riesgo oncológico:</b> {cancer_risk_score:.1%}<br>
    <br>
    <b>⚠️ Análisis de riesgo detallado:</b><br>
    &emsp;🔴 Melanoma maligno: {melanoma_risk:.1%}<br>
    &emsp;🟠 Carcinoma células basales: {bcc_risk:.1%}<br>
    &emsp;🟡 Lesión pre-cancerosa: {precancer_risk:.1%}<br>
    &emsp;🟢 Lesiones benignas: {benign_total:.1%}<br>
    <br>
    <b>🩺 Evaluación clínica:</b><br>
    &emsp;{urgency}<br>
    &emsp;💡 Recomendación: {recommendation}<br>
    &emsp;⏰ Plazo: {timeframe}<br>
    <br>
    <b>🔍 Características analizadas:</b><br>
    """
    if melanoma_risk > 0.3:
        informe += "• ⚠️  Posibles irregularidades sugestivas de melanoma<br>• 🔍 Asimetría o variación de color detectada<br>"
    if bcc_risk > 0.25:
        informe += "• 🔍 Características compatibles con carcinoma basocelular<br>"
    if precancer_risk > 0.25:
        informe += "• 🔍 Posible queratosis actínica (lesión pre-maligna)<br>"
    if benign_total > 0.6:
        informe += "• ✅ Características predominantemente benignas<br>"

    informe += "<br><b>📊 Diagnósticos diferenciales (ordenados por probabilidad):</b><ul>"
    sorted_indices = np.argsort(probs)[::-1]
    for i, idx in enumerate(sorted_indices):
        marker = "🎯" if i==0 else f"{i+1}."
        prob_bars = int(probs[idx] * 30)
        bar_visual = "█" * prob_bars + "░" * (30 - prob_bars)
        risk_indicator = ("🔴" if idx in [1,4] else "🟡" if idx==0 else "🟢")
        informe += f"<li>{marker} {risk_indicator} {CLASSES[idx]} — {probs[idx]:.1%} <code>{bar_visual}</code></li>"
    informe += "</ul>"

    informe += """
    <br>
    <b>📋 Información técnica:</b><br>
    • Modelo: ViT-Base (HAM10000, sharpened)<br>
    • Precisión validada: 83.7%<br>
    • Preprocesamiento: Optimizado para dermatoscopia<br>
    • Clases detectadas: 7 tipos de lesiones cutáneas<br>
    <br>
    <b>📝 Notas clínicas:</b><br>
    """
    if confianza < 0.7:
        informe += "• ⚠️  Confianza moderada - considerar segunda opinión<br>"
    if cancer_risk_score > 0.3:
        informe += "• 🩺 Documentar evolución con fotografías seriadas<br>"
    informe += "• 📏 Regla ABCDE: Evaluar Asimetría, Bordes, Color, Diámetro, Evolución<br>"

    informe += """
    <br><hr>
    <b>⚕️ ADVERTENCIA MÉDICA</b><br>
    • Este análisis es una HERRAMIENTA DE APOYO diagnóstico.<br>
    • NO sustituye la evaluación clínica de un dermatólogo.<br>
    • Ante cualquier duda, consulte con un profesional médico.<br>
    • La decisión terapéutica final corresponde al médico tratante.<br>
    </div>
    """
    return informe, html_chart

demo = gr.Interface(
    fn=analizar_lesion_vit_web,
    inputs=gr.Image(type="pil", label="Sube una imagen de la lesión"),
    outputs=[gr.HTML(label="Informe detallado"), gr.HTML(label="Gráfico de barras")],
    title="Detector de Cáncer de Piel IA (HAM10000, ViT)",
    description="Sube una imagen dermatológica y obtén un informe detallado generado por IA (ViT-HAM10000, accuracy ~83.7%)",
    flagging_mode="never"
)

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