import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from huggingface_hub import hf_hub_download
from tensorflow.keras.preprocessing import image

# === Importar funciones del pipeline ===
from preprocessing.zoom import apply_zoom
from preprocessing.hair_removal import quitar_pelos
from preprocessing.segmentation import segmentar_lesion
from preprocessing.metrics import (
    calcular_area,
    calcular_perimetro,
    calcular_circularidad,
    calcular_simetria
)

# Tamaño de entrada del modelo
ROWS, COLS = 224, 224

# === Cargar modelo desde Hugging Face Hub ===
model_path = hf_hub_download(repo_id="Martinagg/simpleNet", filename="simpleNet.h5")
model = tf.keras.models.load_model(model_path)

# === Cargar modelo ZoomNet para Grad-CAM ===
zoom_path = hf_hub_download(repo_id="Martinagg/ZoomNet", filename="ZoomNet.keras")
model_zoom = tf.keras.models.load_model(zoom_path)

# === Función Grad-CAM ===
def make_gradcam_heatmap(img_array, model, last_conv_layer_name="conv4", pred_index=None):
    grad_model = tf.keras.models.Model(
        [model.inputs],
        [model.get_layer(last_conv_layer_name).output, model.output]
    )
    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(img_array)
        if isinstance(predictions, list):
            predictions = predictions[0]
        if pred_index is None:
            pred_index = tf.argmax(predictions[0])
        class_channel = predictions[:, pred_index]
    grads = tape.gradient(class_channel, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    conv_outputs = conv_outputs[0]
    heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)
    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
    return heatmap.numpy()

# === Función principal ===
def preprocess_and_predict(img_input):
    img = np.array(img_input)[:, :, ::-1]
    zoomed = apply_zoom(img, zoom_factor=0.9)
    rgb_clean = cv2.cvtColor(zoomed, cv2.COLOR_BGR2RGB)
    clean = quitar_pelos(rgb_clean)

    mask, lesion_rgb = segmentar_lesion(clean, size=(ROWS, COLS))

    lesion_resized = cv2.resize(lesion_rgb, (ROWS, COLS))
    img_array = image.img_to_array(lesion_resized) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    probs = model.predict(img_array)[0]
    classes = ["Benign", "Malignant"]
    pred_idx = np.argmax(probs)
    pred_label = classes[pred_idx]
    prob_percent = int(probs[pred_idx] * 100)

    # Color dinámico para la barra
    color = "green" if pred_label == "Benign" else "red"
    result_text_html = f"<h3 style='color:{color}; font-weight:bold;'>Predicción: {pred_label}</h3>"
    result_bar_html = f"""
    <div style='width:100%; background:#eee; border-radius:8px;'>
      <div style='width:{prob_percent}%; background:{color}; padding:6px; 
                  border-radius:8px; color:white; text-align:center; font-weight:bold;'>
        {prob_percent}%
      </div>
    </div>
    """

    # === Métricas geométricas ===
    area = calcular_area(mask)
    perim = calcular_perimetro(mask)
    circ = calcular_circularidad(mask)
    sim_v, sim_h = calcular_simetria(mask)

    metrics_data = [
        ["Área (px²)", area],
        ["Perímetro (px)", perim],
        ["Circularidad", round(circ, 3)],
        ["Simetría Vertical", round(sim_v, 3)],
        ["Simetría Horizontal", round(sim_h, 3)]
    ]

    # === Grad-CAM ===
    raw_resized = cv2.resize(np.array(img_input), (ROWS, COLS))
    raw_array = image.img_to_array(raw_resized) / 255.0
    raw_array = np.expand_dims(raw_array, axis=0)

    heatmap = make_gradcam_heatmap(raw_array, model_zoom, last_conv_layer_name="conv4")
    heatmap = cv2.resize(heatmap, (ROWS, COLS))
    heatmap = np.uint8(255 * heatmap)
    heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    overlay = cv2.addWeighted(raw_resized.astype("uint8"), 0.6, heatmap_color, 0.4, 0)

    return mask, lesion_rgb, result_text_html, result_bar_html, metrics_data, overlay


# === Interfaz con estilo ===
with gr.Blocks(css="""
body, .gradio-container {
    font-family: 'Inter', sans-serif;
    background: #ffffff !important;
    font-weight: bold !important; 
}
h1, h2 { font-weight: 600; color: #111827; margin-bottom: 0.5rem; }
.section { 
    background: #f9fafb;   /* gris muy claro */
    border-radius: 0.75rem; 
    padding: 1.5rem; 
    margin: 1.5rem auto; 
    box-shadow: 0 1px 3px rgba(0,0,0,0.08); 
    max-width: 900px; 
}
.gradio-container { max-width: 900px; margin: auto; }
img { border-radius: 0.5rem; }

/* === Botón Analizar personalizado === */
#analyze-btn {
    display: block;
    margin: 1.5rem auto; 
    width: 90%;                   
    padding: 0.75rem 1rem;        
    background-color: #f97316;    
    color: #7c2d12;               
    border: 2px solid #f97316;    
    border-radius: 0.75rem;       
    font-weight: bold;
    font-size: 1rem;
    cursor: pointer;
    transition: background 0.3s ease, color 0.3s ease;
}
#analyze-btn:hover {
    background-color: #ea580c;   
    color: white;                
}
""") as demo:

    # === Título e introducción ===
    gr.HTML("""
    <section style="text-align:center; padding: 2rem;">
        <h1>DermaScan - Clasificación de Melanomas</h1>
        <!-- Caption breve -->
        <p style="color:#6b7280; font-size:1rem; font-style:italic; margin-top:0.5rem;">
        Una herramienta de apoyo basada en IA para la detección temprana de cáncer de piel.
        </p>
    </section>

    <!-- Texto introductorio alineado a la izquierda -->
    <section style="text-align:justify; padding: 1.5rem; max-width: 800px; margin:auto; line-height:1.6;">
        <p style="color:#111827; font-size:1.05rem;">
                <span style='font-weight:bold; color:#000000;'>El melanoma es un cáncer de piel agresivo</span> que se origina en los melanocitos. 
            Aunque poco frecuente, es el más peligroso por su capacidad de generar metástasis si no se detecta a tiempo.
        </p>
        <p style="color:#111827; font-size:1.05rem; margin-top:1rem;">
            En esta aplicación hemos implementado una <span style='font-weight:bold; color:#000000;'>red neuronal convolucional (CNN)</span>  entrenada con imágenes dermatoscópicas
            para estimar la probabilidad de que una lesión sea  benigna o maligna.
        </p>
        <p style="color:#111827; font-size:1.05rem; margin-top:1rem;">
            Además, incorporamos técnicas de interpretabilidad como <span style='font-weight:bold; color:#000000;'>Grad-CAM </span> y <span style='font-weight:bold; color:#000000;'>métricas geométricas </span> basadas en el criterio clínico ABCDE,
            que sirven como apoyo en la exploración médica.
        </p>
    </section>
    """)

    # === Subir imagen ===
    with gr.Column(elem_classes="section"):
        gr.HTML("<h2>Subir imagen</h2>")
        gr.HTML("<p style='color:#111827;'>Sube una imagen clara de tu lunar Asegúrate de que esté bien iluminada y centrada.</p>")
        img_input = gr.Image(type="pil", label="Imagen de la lesión", elem_id="upload-img")
        run_btn = gr.Button("Analizar", elem_id="analyze-btn", scale=0)

    # === Segmentación ===
    with gr.Column(elem_classes="section"):
        gr.HTML("<h2>Preprocesamiento y Segmentación</h2>")
        gr.HTML("""
        <p style='color:#111827;'>En este apartado verás cómo preparamos tu imagen para que la red neuronal se concentre en la lesión:</p>
        <ul style='color:#111827;'>
                <li style="color:#111827; font-weight:bold;">Conversión de canales de color a HSV.</li>
                <li style="color:#111827; font-weight:bold;">Eliminación de pelos y bordes para reducir ruido.</li>
                <li style="color:#111827; font-weight:bold;">Segmentación del lunar.</li>
        </ul>
        """)
        img_mask = gr.Image(type="numpy", label="Máscara Binaria", elem_id="mask-img")
        img_segmented = gr.Image(type="numpy", label="Lesión Segmentada", elem_id="seg-img")

    # === Grad-CAM ===
    with gr.Column(elem_classes="section"):
        gr.HTML("<h2>Grad-CAM</h2>")
        gr.HTML("""
        <p style='color:#111827;'>
        El Grad-CAM resalta las zonas en las que la red neuronal ha puesto mayor atención para clasificar la lesión. 
        Los colores cálidos indican mayor relevancia: rojo y amarillo muestran las áreas más importantes.
        </p>
        """)
        gradcam_img = gr.Image(type="numpy", label="Mapa de activación", elem_id="gradcam-img")
        gr.HTML("""
        <div style="display:flex; justify-content:center; margin-top:10px; gap:2rem;">
            <div style="display:flex; align-items:center; gap:0.5rem;">
                <div style="width:16px; height:16px; border-radius:50%; background:blue;"></div>
                <span style="color:#111827;">Baja relevancia</span>
            </div>
            <div style="display:flex; align-items:center; gap:0.5rem;">
                <div style="width:16px; height:16px; border-radius:50%; background:gold;"></div>
                <span style="color:#111827;">Media relevancia</span>
            </div>
            <div style="display:flex; align-items:center; gap:0.5rem;">
                <div style="width:16px; height:16px; border-radius:50%; background:red;"></div>
                <span style="color:#111827;">Alta relevancia</span>
            </div>
        </div>
        """)    

   # === Resultados ===
    with gr.Column(elem_classes="section"):
        gr.HTML("<h2>Resultados del modelo</h2>")
        result_text = gr.HTML()
        result_bar = gr.HTML()

    # === Cuadro amarillo estilo ABCDE ===
        gr.HTML("""
        <div style="background-color:#FFFBEB; border-radius:10px; padding:1.5rem; margin:1rem 0; 
                    border:1px solid #FDE68A;">
            <h3 style="color:#111827; font-weight:600; margin-bottom:1rem;">Criterio ABCDE del melanoma</h3>
            <div style="display:flex; justify-content:space-around; text-align:center; margin-bottom:1rem;">
                <div>
                    <div style="background:#FDE68A; border-radius:50%; width:36px; height:36px; 
                                display:flex; align-items:center; justify-content:center; font-weight:bold; color:#92400E; margin:auto;">
                        A
                    </div>
                    <p style="margin-top:0.5rem; font-size:0.9rem; color:#111827;">Asimetría</p>
                </div>
                <div>
                    <div style="background:#FDE68A; border-radius:50%; width:36px; height:36px; 
                                display:flex; align-items:center; justify-content:center; font-weight:bold; color:#92400E; margin:auto;">
                        B
                    </div>
                    <p style="margin-top:0.5rem; font-size:0.9rem; color:#111827;">Bordes irregulares</p>
                </div>
                <div>
                    <div style="background:#FDE68A; border-radius:50%; width:36px; height:36px; 
                                display:flex; align-items:center; justify-content:center; font-weight:bold; color:#92400E; margin:auto;">
                        C
                    </div>
                    <p style="margin-top:0.5rem; font-size:0.9rem; color:#111827;">Color variado</p>
                </div>
                <div>
                    <div style="background:#FDE68A; border-radius:50%; width:36px; height:36px; 
                                display:flex; align-items:center; justify-content:center; font-weight:bold; color:#92400E; margin:auto;">
                        D
                    </div>
                    <p style="margin-top:0.5rem; font-size:0.9rem; color:#111827;">Diámetro &gt; 6mm</p>
                </div>
                <div>
                    <div style="background:#FDE68A; border-radius:50%; width:36px; height:36px; 
                                display:flex; align-items:center; justify-content:center; font-weight:bold; color:#92400E; margin:auto;">
                        E
                    </div>
                    <p style="margin-top:0.5rem; font-size:0.9rem; color:#111827;">Evolución</p>
                </div>
            </div>
            <p style="color:#374151; font-size:0.9rem; text-align:center;">
                Estas características aumentan la probabilidad de melanoma. 
                Consulte a un dermatólogo si observa alguna.
            </p>
        </div>
        """)

        # === Tabla de métricas debajo ===
        gr.HTML("""
        <p style='color:#111827;'>Métricas lunar para consultar en casos</p>
        """)
        metrics_table = gr.Dataframe(
            headers=["Métrica", "Valor"],
            datatype=["str", "number"],
            interactive=False,
            label="Métricas calculadas",
            wrap=True,
            row_count=(5, "fixed"),
            col_count=(2, "fixed")
        )


    # === Aviso final ===
    gr.HTML("""
    <section style="text-align:center; padding: 1rem; margin-top: 2rem;">
        <p style="color:#b91c1c; font-weight:bold; font-size:1rem;">
         Este sistema es solo de apoyo y nunca sustituye la valoración de un experto médico.
        </p>
    </section>
    """)

    # === Conexión botón -> función ===
    run_btn.click(
        fn=preprocess_and_predict,
        inputs=[img_input],
        outputs=[img_mask, img_segmented, result_text, result_bar, metrics_table, gradcam_img]
    )

# === Lanzar en tema claro ===
if __name__ == "__main__":
    demo.launch()