"""
BRACS WSI Explainability — Hugging Face Space (Gradio).
Two input modes:
  • Upload a pre-extracted .h5 feature file  → fast, always works.
  • Upload a whole slide (.svs/.tiff/.ndpi/.mrxs) → app extracts features
    with TITAN+CONCH (slow on CPU; needs HF_TOKEN secret).
Outputs: prediction + a downloadable PDF explainability report.

Research use only — NOT for clinical diagnosis.
"""
import os, tempfile, traceback
import pickle
import gradio as gr

import engine

# ---- Model source: private HF model repo (downloaded with HF_TOKEN) ----------
# Set these as Space env vars / secrets:
#   HF_TOKEN       : token with read access to your private model repo (required)
#   MODEL_REPO_ID  : e.g. "your-username/bracs-models"
#   MODEL_FILENAME : default "bracs_v2_model.pkl"
#   ABMIL_FILENAME : default "abmil_model.pkl" (optional)
#   MODEL_PATH     : optional local override (skips download if file exists)
MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "jehadcheyi/bc_models")
MODEL_FILENAME = os.environ.get("bc_models", "bracs_v2_model.pkl")
ABMIL_FILENAME = os.environ.get("bc_models", "abmil_model.pkl")
MODEL_PATH = os.environ.get("MODEL_PATH", "")  # local fallback
HF_TOKEN = os.environ.get("jj1", "").strip() or None
THUMB_MAX = int(os.environ.get("THUMB_MAX", "4000"))
ZOOM_CONTEXT = int(os.environ.get("ZOOM_CONTEXT", "4"))

WSI_EXTS = (".svs", ".tif", ".tiff", ".ndpi", ".mrxs", ".scn", ".bif")


def _resolve_model_path():
    """Return a local path to the classifier .pkl, downloading from the
    private HF model repo if needed."""
    # 1) explicit local override
    if MODEL_PATH and os.path.isfile(MODEL_PATH):
        return MODEL_PATH
    # 2) already sitting next to app.py (if user uploaded it into the Space)
    if os.path.isfile(MODEL_FILENAME):
        return MODEL_FILENAME
    # 3) download from the private model repo
    if not MODEL_REPO_ID:
        raise RuntimeError(
            "No model found. Set MODEL_REPO_ID (and HF_TOKEN) to your private "
            "HF model repo, or upload the .pkl into the Space."
        )
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME,
                           repo_type="model", token=HF_TOKEN)
 
 
def _resolve_abmil_path():
    """Optional ABMIL bundle; returns None if unavailable."""
    if os.path.isfile(ABMIL_FILENAME):
        return ABMIL_FILENAME
    if not MODEL_REPO_ID:
        return None
    try:
        from huggingface_hub import hf_hub_download
        return hf_hub_download(repo_id=MODEL_REPO_ID, filename=ABMIL_FILENAME,
                               repo_type="model", token=HF_TOKEN)
    except Exception:
        return None
 
 
# Load the classifier bundle once.
_BUNDLE = None
def get_bundle():
    global _BUNDLE
    if _BUNDLE is None:
        path = _resolve_model_path()
        with open(path, "rb") as f:
            _BUNDLE = pickle.load(f)
    return _BUNDLE
 
 
def _result_markdown(info):
    badge = "🔴" if info["prediction"].lower().startswith("malig") else "🟢"
    return (
        f"## {badge} {info['prediction']}\n"
        f"- **Confidence:** {info['confidence']*100:.1f}%\n"
        f"- **P(malignant):** {info['p_malignant']:.3f}\n"
        f"- **Stability:** {info['stability']*100:.0f}%\n"
        f"- **Patches analyzed:** {info['n_patches']}\n"
        f"- **Operating threshold:** {info['threshold']:.2f}\n"
        f"- **WSI image in report:** {'yes' if info['had_wsi'] else 'no (.h5 only)'}\n\n"
        f"_Research use only — not for clinical diagnosis._"
    )
 
 
def run_h5(h5_file, wsi_file, progress=gr.Progress()):
    """Path A: user uploaded an .h5 (optionally a WSI for the images)."""
    if h5_file is None:
        return "Please upload an .h5 feature file.", None
    try:
        progress(0.1, desc="Loading features…")
        slide = engine.load_slide_h5(h5_file.name)
        bundle = get_bundle()
        wsi_path = wsi_file.name if wsi_file is not None else None
        out_pdf = os.path.join(tempfile.gettempdir(), f"{slide['slide_id']}_report.pdf")
        progress(0.5, desc="Predicting & building report…")
        info = engine.build_report(bundle, slide, wsi_path, out_pdf,
                                   thumb_max=THUMB_MAX, zoom_context=ZOOM_CONTEXT)
        progress(1.0, desc="Done.")
        return _result_markdown(info), out_pdf
    except Exception as e:
        return f"**Error:** {e}\n\n```\n{traceback.format_exc()}\n```", None
 
 
def run_wsi(wsi_file, progress=gr.Progress()):
    """Path B: user uploaded a whole slide; we extract features then report."""
    if wsi_file is None:
        return "Please upload a whole-slide image.", None
    if not wsi_file.name.lower().endswith(WSI_EXTS):
        return f"Unsupported file type. Supported: {', '.join(WSI_EXTS)}", None
    try:
        import extract
        msgs = {"t": 0.0}
        def prog(msg):
            msgs["t"] = min(0.95, msgs["t"] + 0.05)
            progress(msgs["t"], desc=msg)
 
        out_h5 = os.path.join(tempfile.gettempdir(), "extracted_features.h5")
        progress(0.05, desc="Starting extraction (slow on CPU)…")
        extract.extract_to_h5(wsi_file.name, out_h5, progress=prog)
 
        slide = engine.load_slide_h5(out_h5)
        bundle = get_bundle()
        out_pdf = os.path.join(tempfile.gettempdir(), f"{slide['slide_id']}_report.pdf")
        progress(0.97, desc="Building report…")
        info = engine.build_report(bundle, slide, wsi_file.name, out_pdf,
                                   thumb_max=THUMB_MAX, zoom_context=ZOOM_CONTEXT)
        progress(1.0, desc="Done.")
        return _result_markdown(info), out_pdf
    except Exception as e:
        return f"**Error during extraction/prediction:** {e}\n\n```\n{traceback.format_exc()}\n```", None
 
 
INTRO = """
# 🔬 BRACS Breast Histopathology Classifier
Upload a whole-slide image **or** a pre-extracted `.h5` feature file. The app
predicts **Benign vs Malignant** and generates a downloadable **PDF explainability
report** with an attention heatmap, high-resolution tissue zooms, and a patch-level
relevance map.
 
⚠️ **Research use only — not for clinical diagnosis.**
"""
 
with gr.Blocks(title="BRACS WSI Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)
 
    with gr.Tab("⚡ Upload .h5 features (fast)"):
        gr.Markdown(
            "Upload an `.h5` produced by the BRACS extraction pipeline. "
            "Optionally add the matching slide image so the report includes the "
            "tissue heatmap and zoom panels."
        )
        with gr.Row():
            h5_in = gr.File(label="Feature file (.h5)", file_types=[".h5"])
            wsi_opt = gr.File(label="Slide image (optional, for visuals)",
                              file_types=list(WSI_EXTS) + [".png", ".jpg", ".jpeg"])
        h5_btn = gr.Button("Predict & build report", variant="primary")
        h5_out = gr.Markdown()
        h5_pdf = gr.File(label="📄 Download PDF report")
        h5_btn.click(run_h5, inputs=[h5_in, wsi_opt], outputs=[h5_out, h5_pdf])
 
    with gr.Tab("🧫 Upload whole slide (extract on CPU — slow)"):
        gr.Markdown(
            "Upload a whole-slide image. The app extracts TITAN+CONCH features, then "
            "predicts. **This is slow on free CPU and may time out for large slides.** "
            "For `.mrxs`, the companion data folder must be present.\n\n"
            "Requires the `HF_TOKEN` secret (TITAN is a gated model)."
        )
        wsi_in = gr.File(label="Whole slide", file_types=list(WSI_EXTS))
        wsi_btn = gr.Button("Extract, predict & build report", variant="primary")
        wsi_out = gr.Markdown()
        wsi_pdf = gr.File(label="📄 Download PDF report")
        wsi_btn.click(run_wsi, inputs=[wsi_in], outputs=[wsi_out, wsi_pdf])
 
    gr.Markdown(
        "---\nModel: LR_concat (TITAN + pooled CONCH). "
        "Built from the BRACS patient-grouped pipeline."
    )
 
if __name__ == "__main__":
    demo.queue(max_size=8).launch(show_api=False, ssr_mode=False)