# ============================================================
#  app.py — Interface Gradio · Agentic AI v5 FIXED
#  HuggingFace Space : mayoula/RAMTUNET-AgenticAI
#  Déployer avec hardware: t4-small (GPU requis pour pipeline)
# ============================================================

import os
import json
import time
import warnings
import requests
from io import BytesIO

import numpy as np
import torch
import gradio as gr
from PIL import Image

warnings.filterwarnings("ignore")

# ─── Secrets depuis les variables d'environnement du Space ──────────────────
HF_TOKEN   = os.environ.get("HF_TOKEN", "")
GROQ_KEY   = os.environ.get("GROQ_API_KEY", "")
OUTPUT_DIR = "/tmp/agentic_results"
DEVICE     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ─── Import du pipeline (toutes les classes agents du notebook) ──────────────
try:
    from pipeline import AgentConfig, AgenticNeuroOncologySystemV5Fixed
    AgentConfig.HF_TOKEN     = HF_TOKEN
    AgentConfig.GROQ_API_KEY = GROQ_KEY
    AgentConfig.OUTPUT_DIR   = OUTPUT_DIR
    PIPELINE_OK = True
    print("✅ pipeline.py chargé")
except ImportError as e:
    PIPELINE_OK = False
    print(f"⚠️  pipeline.py introuvable : {e} — seul le mode démo sera disponible")

# ─── Singleton : le système est initialisé une seule fois ───────────────────
_system = None

def get_system():
    """Charge et initialise les 3 modèles HF (une seule fois en mémoire)."""
    global _system
    if _system is None:
        if not PIPELINE_OK:
            raise RuntimeError(
                "pipeline.py manquant. Copiez le code de votre notebook "
                "dans pipeline.py (voir README)."
            )
        _system = AgenticNeuroOncologySystemV5Fixed(AgentConfig)
        _system.initialize_all()
    return _system


# ─── Images pré-calculées depuis HF repo (mode démo sans GPU) ───────────────
HF_DEMO_REPO = "mayoula/RAMTUNET-AgenticAI"

DEMO_FILES = {
    "xai":  "xai_maps.png",
    "dt":   "dt_simulation_fixed.png",
    "vlm":  "vlm_input_slice.png",
    "pub":  "publication_figure_fixed.png",
    "mem":  "patient_memory_graph.png",
    "json": "agentic_results_v5_fixed.json",
}

def _hf_url(filename: str) -> str:
    token_param = f"?token={HF_TOKEN}" if HF_TOKEN else ""
    return (f"https://huggingface.co/{HF_DEMO_REPO}"
            f"/resolve/main/{filename}{token_param}")

def _safe_img(path) -> Image.Image | None:
    """Ouvre une image PIL depuis un chemin local."""
    if path and os.path.exists(str(path)):
        try:
            return Image.open(path).convert("RGB")
        except Exception:
            return None
    return None


# ─── Construction du rapport texte ──────────────────────────────────────────
def _build_report(results: dict, patient_info: dict) -> str:
    seg    = results.get("segmentation", {})
    dt_s   = results.get("digital_twin", {}).get("summary", {})
    bio    = results.get("biomarkers",   {}).get("summary", {})
    vlm_r  = results.get("vlm",          {}).get("report",       "N/A")
    chief  = results.get("chief",        {}).get("chief_report",  "N/A")
    con_s  = results.get("consensus",    {}).get("summary",       {})
    crit_s = results.get("critic",       {}).get("summary",       {})
    vols   = seg.get("volumes_cm3",   {})
    eval_m = seg.get("eval_metrics",  {})
    risk   = con_s.get("final_risk", "N/A").upper()
    icon   = {"HIGH": "🔴", "INTERMEDIATE": "🟡", "LOW": "🟢"}.get(risk, "⚪")

    sep = "═" * 60
    return f"""{icon}  RISQUE : {risk}  — Consensus Dempster-Shafer [FIX-3]
{sep}
Patient  : {patient_info.get('id','?')}  |  Âge : {patient_info.get('age','?')}  |  KPS : {patient_info.get('kps','?')}
{sep}

📊 VOLUMES TUMORAUX
   Whole Tumor  (WT) = {vols.get('WT','?')} cm³
   Tumor Core   (TC) = {vols.get('TC','?')} cm³
   Enhancing    (ET) = {vols.get('ET','?')} cm³
   NCR               = {round(float(vols.get('TC',0)) - float(vols.get('ET',0)), 2)} cm³
   Oedème       (ED) = {round(float(vols.get('WT',0)) - float(vols.get('TC',0)), 2)} cm³

🎯 MÉTRIQUES SEGMENTATION  [FIX-4 depuis HF evaluation_metrics.json]
   Dice WT  = {eval_m.get('dice_WT', 0):.3f}
   Dice TC  = {eval_m.get('dice_TC', 0):.3f}
   Dice ET  = {eval_m.get('dice_ET', 0):.3f}
   HD95 WT  = {eval_m.get('hd95_WT', 0):.1f} mm
   HD95 ET  = {eval_m.get('hd95_ET', 0):.1f} mm

🧬 BIOMARQUEURS RADIOMIQUES
   Sphéricité ET    = {bio.get('ET_sphericity', 0):.3f}
   Hétérogénéité    = {bio.get('heterogeneity', 0):.4f}
   Ratio ET/TC      = {bio.get('et_tc_ratio', 0):.4f}
   Ratio NCR/TC     = {bio.get('ncr_tc_ratio', 0):.4f}
   Diamètre max     = {bio.get('max_diameter_mm', 0)} mm

🔮 DIGITAL TWIN  [FIX-2 PDE Fisher-KPP réel]
   ρ (croissance)        = {dt_s.get('rho', '?')}
   Temps de doublement   = {dt_s.get('doubling_time_days', '?')} jours
   TTP (avec Stupp RT)   = {dt_s.get('TTP_months', '?')} mois
   OS prédit             = {dt_s.get('OS_months', '?')} mois
   Source OS             = {dt_s.get('OS_source', '?')}
   Infiltration PDE J90  = {float(dt_s.get('infiltration_score_pde', 0)):.4f}
   Catégorie risque DT   = {dt_s.get('risk_category', '?').upper()}
   Score risque          = {dt_s.get('risk_score', '?')}

🗳️  CONSENSUS DEMPSTER-SHAFER  [FIX-3]
   Risque final      = {risk}
   Accord agents DS  = {con_s.get('agreement', '?')}
   Confiance globale = {con_s.get('global_confidence', '?')}
   Incertitude DS    = {con_s.get('uncertainty', '?')}
   OS IC 90%         = {con_s.get('os_ci_90', 'N/A')}
   TTP IC 90%        = {con_s.get('ttp_ci_90', 'N/A')}

🔍 CRITIC AGENT  [Bayésien FIX-3]
   Score             = {crit_s.get('critic_score', '?')}
   Issues critiques  = {crit_s.get('n_critical_issues', 0)}
   Avertissements    = {crit_s.get('n_warnings', 0)}
   Formule confiance = {crit_s.get('rigid_formula', 'N/A')}

{sep}
📋 RAPPORT RADIOLOGIQUE VLM  [FIX-1 — image PNG 3 plans transmise à LLaVA-Med]
{sep}
{vlm_r[:4000]}

{sep}
🎯 RAPPORT CHIEF AGENT
{sep}
{chief[:4000]}
"""


def _empty() -> tuple:
    return None, None, None, None, None, "", "En attente…"


# ─── Fonction analyse complète ───────────────────────────────────────────────
def run_analysis(
    flair_file, t1_file, t1ce_file, t2_file,
    patient_id, age, kps,
    progress=gr.Progress(),
):
    """Lance le pipeline 12-agents complet sur les 4 fichiers NIfTI uploadés."""

    # Validation des fichiers
    if not all([flair_file, t1_file, t1ce_file, t2_file]):
        gr.Warning("⚠️  Veuillez uploader les 4 modalités IRM avant de lancer l'analyse.")
        return _empty()[:-1] + ("❌ Fichiers manquants — uploadez FLAIR · T1 · T1CE · T2",)

    progress(0.02, desc="⚙️  Initialisation des modèles (HuggingFace)…")

    try:
        sys_inst = get_system()
    except RuntimeError as e:
        return _empty()[:-1] + (f"❌ {e}",)
    except Exception as e:
        return _empty()[:-1] + (f"❌ Erreur initialisation : {e}",)

    mri_paths = {
        "flair": flair_file,
        "t1":    t1_file,
        "t1ce":  t1ce_file,
        "t2":    t2_file,
    }
    patient_info = {
        "id":  (patient_id or "PATIENT_001").strip(),
        "age": int(age),
        "kps": int(kps),
        "sex": "M",
    }

    progress(0.10, desc="🔬 Segmentation RCMTUNetV4  (TTA × 8)…")

    try:
        results = sys_inst.run_pipeline(mri_paths, patient_info)
    except Exception as e:
        return _empty()[:-1] + (f"❌ Erreur pipeline : {str(e)[:500]}",)

    progress(0.90, desc="📊 Collecte des visualisations…")

    xai = _safe_img(results.get("explainability", {}).get("xai_map_path"))
    dt  = _safe_img(results.get("explainability", {}).get("dt_sim_path"))
    vlm = _safe_img(results.get("vlm",            {}).get("vis_path"))
    pub = _safe_img(results.get("survival_benchmark", {}).get("publication_figure"))
    mem = _safe_img(results.get("memory_graph_path"))

    report  = _build_report(results, patient_info)
    elapsed = results.get("total_time_s", "?")
    risk    = results.get("consensus",{}).get("summary",{}).get("final_risk","N/A").upper()
    icon    = {"HIGH":"🔴","INTERMEDIATE":"🟡","LOW":"🟢"}.get(risk,"⚪")

    progress(1.0, desc="✅ Analyse terminée !")
    status = (f"✅ Terminé en {elapsed}s  |  {icon} Risque : {risk}  |"
              f"  Device : {DEVICE}  |  LLM : {sys_inst.llm.mode.upper()}")

    return xai, dt, vlm, pub, mem, report, status


# ─── Fonction mode démo ──────────────────────────────────────────────────────
def load_demo(progress=gr.Progress()):
    """Charge les résultats pré-calculés depuis HF repo (aucun GPU requis)."""

    progress(0.05, desc="📥 Connexion à HuggingFace…")
    imgs = {}
    file_keys = [k for k in DEMO_FILES if k != "json"]

    for i, key in enumerate(file_keys):
        fname = DEMO_FILES[key]
        try:
            resp = requests.get(_hf_url(fname), timeout=25)
            resp.raise_for_status()
            imgs[key] = Image.open(BytesIO(resp.content)).convert("RGB")
            print(f"  ✅ {fname} ({len(resp.content)//1024} KB)")
        except Exception as e:
            imgs[key] = None
            print(f"  ⚠️  {fname} : {e}")
        progress(0.05 + 0.75 * (i + 1) / len(file_keys),
                 desc=f"📥 {fname}…")

    # Rapport depuis le JSON pré-calculé
    report = "═══ DEMO MODE — Résultats pré-calculés (BraTS2021_00000) ═══\n\n"
    try:
        resp_j = requests.get(_hf_url("agentic_results_v5_fixed.json"), timeout=25)
        rj     = resp_j.json()
        report += _build_report(rj, {"id": "BraTS2021_00000", "age": 55, "kps": 80})
    except Exception as e:
        report += f"(JSON non disponible : {e})\n"

    progress(1.0, desc="✅ Démo chargée depuis HuggingFace !")
    status = ("✅ Mode démo — résultats pré-calculés depuis "
              f"huggingface.co/{HF_DEMO_REPO}  (aucun GPU requis)")

    return (imgs.get("xai"), imgs.get("dt"),  imgs.get("vlm"),
            imgs.get("pub"), imgs.get("mem"), report, status)


# ─── Interface Gradio ────────────────────────────────────────────────────────
CSS = """
.gradio-container { max-width: 1400px !important; }
.header-block { text-align: center; padding: 6px 0 2px; }
footer { display: none !important; }
"""

GPU_INFO = (
    f"🚀 GPU — {torch.cuda.get_device_name(0)}"
    if torch.cuda.is_available()
    else "⚠️  CPU (pipeline non disponible — utilisez le mode démo)"
)

with gr.Blocks(
    title="🧠 GBM Agentic AI v5",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"),
    css=CSS,
) as demo:

    # ── En-tête ───────────────────────────────────────────────────────────────
    gr.Markdown(f"""
    <div class="header-block">

    # 🧠 Agentic AI v5 FIXED — GBM Neuro-Oncology Analysis

    **RCMTUNetV4** (Segmentation) · **Digital Twin** (Gompertz + Fisher-KPP PDE) · **Dempster-Shafer** (Consensus)

    Uploadez 4 modalités IRM au format NIfTI — analyse complète par **12 agents IA**.
    Sans GPU → **👁 Charger la démo** pour voir les résultats pré-calculés instantanément.

    `{GPU_INFO}` &nbsp;|&nbsp; `Pipeline: {"✅ prêt" if PIPELINE_OK else "❌ pipeline.py manquant"}`

    </div>
    """)

    # ── Onglet 1 — Upload & Analyser ─────────────────────────────────────────
    with gr.Tab("📤 Upload & Analyser"):

        with gr.Row(equal_height=False):

            # Colonne patient
            with gr.Column(scale=1, min_width=240):
                gr.Markdown("### 👤 Informations Patient")
                pid_in = gr.Textbox(
                    label="ID Patient",
                    value="PATIENT_001",
                    placeholder="ex: BraTS2021_00000",
                )
                age_in = gr.Slider(18, 90,  value=55, step=1,  label="Âge (ans)")
                kps_in = gr.Slider(40, 100, value=80, step=10, label="KPS Score")
                gr.Markdown("---")
                gr.Markdown("""
**Format requis :**
- NIfTI `.nii` ou `.nii.gz`
- 1 mm isotropique
- Crâne stripé + co-registré
- Convention BraTS 2021
                """)

            # Colonne uploads
            with gr.Column(scale=3):
                gr.Markdown("### 📂 4 Modalités IRM (NIfTI)")
                with gr.Row():
                    flair_in = gr.File(
                        label="🟡 FLAIR",
                        file_types=[".nii", ".gz"],
                        type="filepath",
                    )
                    t1_in = gr.File(
                        label="⚪ T1",
                        file_types=[".nii", ".gz"],
                        type="filepath",
                    )
                with gr.Row():
                    t1ce_in = gr.File(
                        label="🔴 T1CE (Gadolinium)",
                        file_types=[".nii", ".gz"],
                        type="filepath",
                    )
                    t2_in = gr.File(
                        label="🟢 T2",
                        file_types=[".nii", ".gz"],
                        type="filepath",
                    )

        with gr.Row():
            run_btn  = gr.Button(
                "🚀 Lancer l'analyse complète (GPU requis)",
                variant="primary", size="lg", scale=3,
            )
            demo_btn = gr.Button(
                "👁️ Charger la démo (sans GPU)",
                variant="secondary", size="lg", scale=2,
            )
            clr_btn  = gr.Button(
                "🔄 Effacer", variant="stop", size="sm", scale=1,
            )

        status_out = gr.Textbox(
            label="Statut",
            interactive=False,
            lines=1,
            placeholder="Uploadez les 4 fichiers puis cliquez sur Analyser…",
        )

    # ── Onglet 2 — Segmentation & XAI ────────────────────────────────────────
    with gr.Tab("🔬 Segmentation & XAI"):
        gr.Markdown("""
**XAI Maps** : superposition segmentation · attention gradient · risk zone map  
**VLM Input** : vue 3-plans axiale/coronale/sagittale transmise à LLaVA-Med `[FIX-1]`
        """)
        with gr.Row():
            xai_out = gr.Image(
                label="XAI Maps — Seg · Attention · Risk Zones",
                type="pil", height=480,
            )
            vlm_out = gr.Image(
                label="Vue 3-Plans — Input VLM [FIX-1]",
                type="pil", height=480,
            )

    # ── Onglet 3 — Digital Twin ───────────────────────────────────────────────
    with gr.Tab("🔮 Digital Twin"):
        gr.Markdown("""
Dashboard **4 × 4** :  
- Ligne 1 : vues anatomiques (T1CE axiale · segmentation · coronale · sagittale)  
- Ligne 2 : PDE Fisher-KPP infiltration J0 / J30 / J90 / J180 `[FIX-2]`  
- Ligne 3 : overlay FLAIR + multi-coupes axiales  
- Ligne 4 : courbes ODE Gompertz · TTP avec IC Bootstrap · résumé patient
        """)
        dt_out = gr.Image(
            label="Digital Twin Dashboard 4×4 (PDE + ODE)",
            type="pil", height=680,
        )

    # ── Onglet 4 — Figure Publication ─────────────────────────────────────────
    with gr.Tab("📊 Figure Publication"):
        gr.Markdown("""
Figure **Q1 journal** — 10 agents :  
`A` Kaplan-Meier &nbsp;·&nbsp; `B` Forest plot HR &nbsp;·&nbsp; `C` Radar confiance `[FIX-4]`  
`D` Heatmap DS `[FIX-3]` &nbsp;·&nbsp; `E` Timeline traitement &nbsp;·&nbsp; `F` Barres agents
        """)
        pub_out = gr.Image(
            label="Figure Publication Q1",
            type="pil", height=680,
        )

    # ── Onglet 5 — Graphe Mémoire ─────────────────────────────────────────────
    with gr.Tab("🧬 Graphe Mémoire"):
        gr.Markdown("""
Knowledge Graph longitudinal **NetworkX DiGraph** :  
sessions IRM · DT snapshots · traitements · récidives · provenance HF `[FIX-3]`
        """)
        mem_out = gr.Image(
            label="Patient Knowledge Graph",
            type="pil", height=580,
        )

    # ── Onglet 6 — Rapport Clinique ────────────────────────────────────────────
    with gr.Tab("📝 Rapport Clinique"):
        gr.Markdown("""
Rapport complet généré par : **VLM** `[FIX-1]` · **Digital Twin** `[FIX-2]`
· **Consensus DS** `[FIX-3]` · **Chief Agent**
        """)
        report_out = gr.Textbox(
            label="Rapport Clinique Complet",
            lines=55,
            max_lines=150,
            placeholder="Lancez l'analyse ou chargez la démo pour voir le rapport…",
        )

    # ─── Câblage ──────────────────────────────────────────────────────────────
    ALL_OUTPUTS = [xai_out, dt_out, vlm_out, pub_out, mem_out, report_out, status_out]
    INPUTS      = [flair_in, t1_in, t1ce_in, t2_in, pid_in, age_in, kps_in]

    run_btn.click(
        fn=run_analysis,
        inputs=INPUTS,
        outputs=ALL_OUTPUTS,
        api_name="analyze",
    )

    demo_btn.click(
        fn=load_demo,
        inputs=[],
        outputs=ALL_OUTPUTS,
        api_name="demo",
    )

    clr_btn.click(
        fn=lambda: (None, None, None, None, None, "", "🔄 Effacé"),
        outputs=ALL_OUTPUTS,
    )


# ─── Point d'entrée ──────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        max_file_size="2gb",
    )