"""LeukoAI 1.1 web application (research use only).

Run locally:   python app/app.py          then open http://127.0.0.1:7860
Hugging Face:  see docs/DEPLOY.md
"""
from __future__ import annotations

import csv
import glob
import html
import os
import sys
import tempfile

import gradio as gr
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from leukoai import load_models  # noqa: E402
from leukoai.explain import gradcam, overlay  # noqa: E402
from leukoai.models import DX1, DX2, DX3  # noqa: E402
from leukoai.pipeline import analyse_paths  # noqa: E402



torch.set_num_threads(max(1, os.cpu_count() or 1))
MODELS = load_models(device="cpu")
MAX_IMAGES = int(os.environ.get("LEUKOAI_MAX_IMAGES", "100"))
USERS = dict(u.strip().split(":", 1) for u in os.environ.get("APP_USERS", "").split(",") if ":" in u)


def check_login(username, password):
    ok = USERS.get(username) == password
    print(f"[ACCESS] user={username} login={'ok' if ok else 'FAILED'}", flush=True)
    return ok
    
ASSETS = os.path.join(ROOT, "assets")

PAPER = "https://doi.org/10.21037/jmai-24-74"
CODE = "https://github.com/synav/LeukoAI"
DATA = "https://doi.org/10.5281/zenodo.17814776"
CARD = "https://github.com/synav/LeukoAI/blob/main/MODEL_CARD.md"

EXAMPLES = {
    key: sorted(glob.glob(os.path.join(ASSETS, "examples", key, "*.jpg")))
    for key in ("aml", "cll", "cml", "normal")
}


def _logo_svg() -> str:
    path = os.path.join(ASSETS, "logo-mark.svg")
    if not os.path.exists(path):
        return ""
    with open(path) as f:
        return f.read().replace("<svg ", '<svg aria-hidden="true" focusable="false" ', 1)


LOGO_SVG = _logo_svg()
FAVICON = os.path.join(ASSETS, "logo-mark.svg")

# ---------------------------------------------------------------- page text

HEAD = """
<meta name="description" content="LeukoAI: research tool that predicts leukemia and its subtype from peripheral blood smear images. Based on Naveed S et al., J Med Artif Intell 2025.">
<meta property="og:title" content="LeukoAI: leukemia prediction from blood smear images">
<meta property="og:description" content="Upload the microscope fields from one smear and get a case-level result with heatmaps. Research use only.">
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@500;600&family=IBM+Plex+Sans:wght@400;500;600;700&display=swap" rel="stylesheet">
<script>
  (function () {
    var u = new URL(window.location.href);
    if (u.searchParams.get('__theme') !== 'light') { u.searchParams.set('__theme', 'light'); window.location.replace(u.toString()); }
  })();
</script>
"""

HEADER = f"""
<header class="lk-nav">
  <a class="lk-brand" href="#top" aria-label="LeukoAI home">
    <span class="lk-mark">{LOGO_SVG}</span><span class="lk-wordmark">Leuko<span>AI</span></span>
  </a>
  <nav class="lk-links" aria-label="Main">
    <a href="#tool">Try it</a>
    <a href="#about">About</a>
    <a href="{PAPER}" target="_blank" rel="noopener">Paper</a>
    <a href="{CODE}" target="_blank" rel="noopener">Code</a>
    <a href="{DATA}" target="_blank" rel="noopener">Data</a>
  </nav>
</header>
"""

HERO = """
<section class="lk-hero" id="top">
  <p class="lk-eyebrow">LeukoAI 1.1 · Research tool</p>
  <h1>Leukemia prediction from peripheral blood smear images</h1>
  <p class="lk-lead">Upload the microscope fields from one patient's smear. LeukoAI classifies every field and
  combines them into one case result: acute or chronic leukemia, the likely subtype (ALL, AML, CLL, CML) or
  no leukemia, with heatmaps showing where the model looked.</p>
  <div class="lk-stats">
    <div><span class="lk-num">417/417</span><span>images give identical results to the published desktop tool</span></div>
    <div><span class="lk-num">0/417</span><span>real smear fields falsely flagged by the image check</span></div>
    <div><span class="lk-num">No scanner</span><span>works on ordinary photos taken at the microscope</span></div>
  </div>
</section>
"""

STEPS = """
<section class="lk-steps" aria-label="How it works">
  <div><span class="lk-step">1</span><h3>Capture</h3><p>Photograph 20 or more fields of one smear at 1000× (oil immersion, May-Grünwald-Giemsa stain).</p></div>
  <div><span class="lk-step">2</span><h3>Upload</h3><p>Select all the images at once (Ctrl+A in the folder) and tick the de-identification box.</p></div>
  <div><span class="lk-step">3</span><h3>Review</h3><p>Read the case result, the warnings and the heatmaps on the most influential fields.</p></div>
</section>
"""

PLACEHOLDER = """
<div class="lk-empty">
  <p class="lk-empty-title">Your result will appear here</p>
  <p>Upload one patient's smear fields, or load an example case, then select <b>Analyse case</b>.
  Analysis takes about 1–2 seconds per image.</p>
</div>
"""

ABOUT = f"""
<section class="lk-about" id="about">
  <h2>About LeukoAI</h2>
  <div class="lk-grid">
    <div>
      <h3>How it works</h3>
      <p>Three DenseNet-121 models work in sequence: the first screens each field for leukemia, the second
      separates acute from chronic, and the third suggests the subtype. A rule then combines the fields into one
      case result: leukemia is reported when more than 3 fields are classed leukemic.</p>
    </div>
    <div>
      <h3>How it was validated</h3>
      <p>The published study used 7,255 images from 162 adult patients. On 64 external patients the
      acute / chronic / no-leukemia result reached an F1 score of 0.94 (accuracy 0.92); the five-class subtype
      result reached a macro-F1 of 0.77. This web version reproduces the original desktop tool exactly.</p>
    </div>
    <div>
      <h3>Limitations</h3>
      <ul>
        <li>Single centre, adult patients, 1000× MGG-stained peripheral smears only.</li>
        <li>ALL vs AML separation is weak (F1 0.36 for ALL); treat the subtype as unconfirmed.</li>
        <li>APML is not reported separately. Scores are not calibrated probabilities.</li>
      </ul>
    </div>
    <div>
      <h3>Privacy</h3>
      <p>Images are used only for your analysis and are deleted from the server automatically within about an
      hour. Nothing is kept for training. Upload de-identified images only, with no names or numbers burned in.</p>
    </div>
  </div>
  <div class="lk-cite">
    <h3>Cite</h3>
    <p>Naveed S, et al. Novel hierarchical deep learning models predict type of leukemia from whole slide
    microscopic images of peripheral blood. <i>Journal of Medical Artificial Intelligence</i>. 2025.
    <a href="{PAPER}" target="_blank" rel="noopener">doi:10.21037/jmai-24-74</a></p>
    <p class="lk-small">Model card: <a href="{CARD}" target="_blank" rel="noopener">github.com/synav/LeukoAI</a> ·
    Public image dataset: <a href="{DATA}" target="_blank" rel="noopener">Zenodo 10.5281/zenodo.17814776</a></p>
  </div>
</section>
"""

FOOTER = """
<footer class="lk-footer">
  <p><b>Research use only.</b> LeukoAI is not a medical device and must not be used for diagnosis or treatment
  decisions. Confirm every case with standard laboratory work-up.</p>
  <p>Developed by Dr Syed Naveed, FRCPath, and team at Sheikh Shakhbout Medical City, Abu Dhabi, with the
  Mohamed bin Zayed University of Artificial Intelligence.</p>
</footer>
"""

CSS = """
:root, .gradio-container { --lk-ink:#14212E; --lk-muted:#55606E; --lk-line:#DCE3EA; --lk-ground:#F4F6F9;
  --lk-brand:#2B6698; --lk-brand-dark:#1D4A70; --lk-card:#FFFFFF; }
html, body, gradio-app, .gradio-container { background: var(--lk-ground) !important; }
.gradio-container { max-width: min(1180px, 100vw) !important; margin: 0 auto !important; padding: 0 clamp(10px, 2.5vw, 16px) !important; overflow-x: hidden;
  font-family: 'IBM Plex Sans', system-ui, sans-serif !important; color: var(--lk-ink); }
.lk-nav { display:flex; align-items:center; justify-content:space-between; flex-wrap:wrap; gap:12px;
  padding: 18px 0 14px; border-bottom: 1px solid var(--lk-line); }
.lk-brand { display:inline-flex; align-items:center; gap:10px; text-decoration:none; }
.lk-mark { width: clamp(34px, 4.5vw, 42px); height: clamp(34px, 4.5vw, 42px); display:inline-block; }
.lk-mark svg { width:100%; height:100%; display:block; }
.lk-wordmark { font-weight:700; font-size: clamp(22px, 2.8vw, 27px); letter-spacing:-0.01em; color:var(--lk-ink); }
.lk-wordmark span { color: var(--lk-brand); }
.lk-links { display:flex; flex-wrap:wrap; gap: 4px 22px; }
.lk-links a { color: var(--lk-ink); text-decoration:none; font-weight:500; font-size:15px; padding:8px 0; }
.lk-links a:hover { color: var(--lk-brand); }
.lk-hero { padding: clamp(28px, 5vw, 44px) 0 clamp(20px, 3vw, 28px); }
.lk-eyebrow { font-family:'IBM Plex Mono', monospace; font-size:13px; letter-spacing:.08em; text-transform:uppercase;
  color: var(--lk-brand); font-weight:600; margin:0 0 12px; }
.lk-hero h1 { font-size: clamp(30px, 4.4vw, 48px); line-height:1.1; letter-spacing:-.02em; margin:0 0 16px;
  font-weight:700; color:var(--lk-ink); max-width: 860px; }
.lk-lead { font-size: clamp(16px, 2vw, 18px); line-height:1.55; color:var(--lk-muted); max-width: 760px; margin:0 0 28px; }
.lk-stats { display:grid; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); gap:16px;
  border-top: 2px solid var(--lk-ink); padding-top: 20px; }
.lk-stats div { display:flex; flex-direction:column; gap:6px; }
.lk-stats span { font-size:15px; line-height:1.45; color:var(--lk-muted); }
.lk-stats .lk-num { font-family:'IBM Plex Mono', monospace; font-size:28px; font-weight:600; color:var(--lk-brand); }
.lk-steps { display:grid; grid-template-columns: repeat(auto-fit, minmax(240px, 1fr)); gap:16px; margin: 8px 0 28px; }
.lk-steps > div { background: var(--lk-card); border:1px solid var(--lk-line); border-radius:10px; padding:18px 20px; }
.lk-steps h3 { margin: 10px 0 6px; font-size:17px; color:var(--lk-ink); }
.lk-steps p { margin:0; color:var(--lk-muted); font-size:15px; line-height:1.5; }
.lk-step { display:inline-flex; align-items:center; justify-content:center; width:30px; height:30px; border-radius:50%;
  background: var(--lk-brand); color:#fff; font-weight:600; font-size:15px; }
#tool { gap: 20px; scroll-margin-top: 16px; }
.lk-panel { background: var(--lk-card) !important; border:1px solid var(--lk-line) !important; border-radius:12px !important;
  padding: clamp(12px, 2vw, 18px) !important; max-width: 100% !important; }
.gradio-container .main, .gradio-container .app { max-width: 100% !important; min-width: 0 !important; }
#tool { flex-wrap: wrap !important; max-width: 100% !important; }
#tool > .lk-panel:first-child { flex: 5 1 340px !important; min-width: min(340px, 100%) !important; }
#tool > .lk-panel:last-child { flex: 7 1 340px !important; min-width: min(340px, 100%) !important; }
.lk-panel-title { margin:0 0 4px; font-size:18px; font-weight:600; color:var(--lk-ink); }
.lk-panel-sub { margin:0 0 10px; font-size:14px; color:var(--lk-muted); line-height:1.5; }
.lk-ex-label { font-size:14px; color:var(--lk-muted); margin: 6px 0 0; }
.lk-empty { border: 2px dashed var(--lk-line); border-radius:10px; padding: 36px 24px; text-align:center; color:var(--lk-muted); }
.lk-empty-title { font-size:18px; font-weight:600; color:var(--lk-ink); margin:0 0 6px; }
.lk-result { border-radius:10px; padding:18px 20px; margin-bottom:12px; border:2px solid; background:#fff; }
.lk-result .lk-label { font-family:'IBM Plex Mono', monospace; font-size:12px; letter-spacing:.06em; text-transform:uppercase; color:var(--lk-muted); }
.lk-result .lk-dx { font-size:30px; font-weight:700; margin:4px 0; letter-spacing:-.01em; }
.lk-result .lk-sub { font-size:18px; color:var(--lk-ink); }
.lk-result .lk-meta { font-size:14px; color:var(--lk-muted); margin-top:10px; line-height:1.5; }
.lk-flag { padding:10px 14px; margin:8px 0; border-radius:8px; font-size:14px; line-height:1.5; color:#101828; }
.lk-about { padding: 40px 0 8px; scroll-margin-top: 16px; }
.lk-about h2 { font-size: 30px; margin:0 0 20px; color:var(--lk-ink); letter-spacing:-.01em; }
.lk-grid { display:grid; grid-template-columns: repeat(auto-fit, minmax(260px, 1fr)); gap: 20px 32px; }
.lk-about h3 { font-size:17px; margin: 0 0 8px; color:var(--lk-ink); }
.lk-about p, .lk-about li { font-size:15px; line-height:1.6; color:var(--lk-muted); margin:0 0 8px; }
.lk-about ul { padding-left: 18px; margin:0; }
.lk-about a, .lk-footer a { color: var(--lk-brand); }
.lk-cite { margin-top: 20px; padding: 18px 20px; background: var(--lk-card); border:1px solid var(--lk-line); border-radius:10px; }
.lk-small { font-size:14px !important; }
.lk-footer { border-top:1px solid var(--lk-line); margin-top: 32px; padding: 20px 0 8px; }
.lk-footer p { font-size:14px; color:var(--lk-muted); margin: 0 0 8px; line-height:1.5; }
"""

THEME = gr.themes.Base(
    primary_hue=gr.themes.Color(c50="#EEF4FA", c100="#D6E5F2", c200="#AECBE5", c300="#7FAAD3", c400="#4F88BE",
                                c500="#2B6698", c600="#245A87", c700="#1D4A70", c800="#173B59", c900="#112C43",
                                c950="#0B1E2E"),
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
    radius_size="md",
).set(button_primary_background_fill="#2B6698", button_primary_background_fill_hover="#1D4A70",
      button_primary_text_color="#FFFFFF")

# ---------------------------------------------------------------- analysis

SEV_STYLE = {
    "warning": ("#B42318", "#FEF3F2", "Warning"),
    "caution": ("#93370D", "#FFFAEB", "Caution"),
    "info": ("#175CD3", "#EFF8FF", "Note"),
}
DX_COLOUR = {"NO LEUKEMIA": "#027A48", "ACUTE LEUKEMIA": "#B42318", "CHRONIC LEUKEMIA": "#93370D"}


def _result_html(rep) -> str:
    r = rep.result
    colour = DX_COLOUR[r.three_class]
    title = r.three_class.capitalize().replace("leukemia", "leukemia")
    subtype = "" if r.three_class == "NO LEUKEMIA" else (
        f"<div class='lk-sub'>Likely subtype: <b>{html.escape(rep.subtype_display)}</b></div>")
    votes = " · ".join(f"{k.title() if k == 'NO LEUKEMIA' else k} {v}" for k, v in r.votes.items())
    parts = [
        f"<div class='lk-result' style='border-color:{colour}'>",
        "<div class='lk-label'>Case result</div>",
        f"<div class='lk-dx' style='color:{colour}'>{html.escape(title)}</div>{subtype}",
        f"<div class='lk-meta'>{r.n_images} fields analysed · {r.n_leukemic} classed leukemic "
        f"(needs more than {r.threshold}) · acute {r.n_acute}, chronic {r.n_chronic}<br>Subtype votes: {votes}</div>",
        "</div>",
    ]
    for fl in sorted(rep.flags, key=lambda f: ["warning", "caution", "info"].index(f.severity)):
        fg, bg, label = SEV_STYLE[fl.severity]
        parts.append(f"<div class='lk-flag' style='background:{bg};border-left:4px solid {fg}'>"
                     f"<b style='color:{fg}'>{label}:</b> {html.escape(fl.message)}</div>")
    return "\n".join(parts)


def _per_image_rows(rep):
    rows = []
    for im in rep.images:
        q = im.quality or {}
        rows.append([im.name, DX1[im.k1].capitalize(), round(im.p1[1], 3), DX2[im.k2].capitalize(), DX3[im.k3],
                     round(max(im.p3), 3), "OK" if q.get("ok", True) else "; ".join(q.get("reasons", []))])
    return rows


def _write_csv(rep) -> str:
    fd, path = tempfile.mkstemp(prefix="leukoai_", suffix=".csv")
    with os.fdopen(fd, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["image"] + [f"m1_{c}" for c in DX1] + [f"m2_{c}" for c in DX2] + [f"m3_{c}" for c in DX3]
                   + ["quality_ok", "quality_reasons"])
        for im in rep.images:
            q = im.quality or {}
            w.writerow([im.name] + [round(x, 5) for x in im.p1 + im.p2 + im.p3]
                       + [q.get("ok", ""), "; ".join(q.get("reasons", []))])
        r = rep.result
        w.writerow([])
        w.writerow(["case_result", r.three_class, "subtype", rep.subtype_display, "threshold", r.threshold])
        for fl in rep.flags:
            w.writerow(["flag", fl.code, fl.severity, fl.message])
    return path


def analyse(files, threshold, show_heatmaps, consent):
    if not consent:
        raise gr.Error("Please confirm the images are de-identified and for research use.")
    if not files:
        raise gr.Error("Please upload the field images from one patient's smear, or load an example case.")
    paths = [f if isinstance(f, str) else f.name for f in files]
    if len(paths) > MAX_IMAGES:
        raise gr.Error(f"Please upload at most {MAX_IMAGES} images per case.")

    rep = analyse_paths(paths, MODELS, threshold=int(threshold), mode="1.1")
    by_name = {os.path.basename(p): p for p in paths}

    gallery = []
    if rep.result.images_of_interest:
        chronic = rep.result.three_class == "CHRONIC LEUKEMIA"
        # Model 2 (acute/chronic) heatmaps localise leukocytes well; Model 1 is usually
        # saturated near 100% on leukemic fields, which makes its gradients uninformative.
        model, cls = MODELS["m2"], (1 if chronic else 0)
        from PIL import Image
        for name in rep.result.images_of_interest:
            img = Image.open(by_name[name]).convert("RGB")
            if show_heatmaps:
                img = overlay(img, gradcam(model, img, cls))
            gallery.append((img, name))

    headers = ["Image", "Screening", "Leukemia score", "Category", "Subtype vote", "Subtype score", "Image check"]
    return (_result_html(rep), gr.update(value=gallery, visible=bool(gallery)),
            gr.update(visible=True), gr.update(value=_per_image_rows(rep), headers=headers),
            gr.update(value=_write_csv(rep), visible=True))


_HIDDEN = (PLACEHOLDER, gr.update(value=[], visible=False), gr.update(visible=False), None,
           gr.update(value=None, visible=False))


def load_example(key):
    return (EXAMPLES[key], True) + _HIDDEN


def clear():
    return (None, False) + _HIDDEN


# ---------------------------------------------------------------- layout

with gr.Blocks(title="LeukoAI · Leukemia prediction from blood smears", delete_cache=(1800, 3600),
               analytics_enabled=False) as demo:
    gr.HTML(HEADER)
    gr.HTML(HERO)
    gr.HTML(STEPS)
    with gr.Row(elem_id="tool", equal_height=False):
        with gr.Column(scale=5, min_width=320, elem_classes="lk-panel"):
            gr.HTML("<p class='lk-panel-title'>Upload one patient's smear</p>"
                    "<p class='lk-panel-sub'>All fields from a single smear, 1000× MGG. JPEG, PNG or TIFF; up to "
                    f"{MAX_IMAGES} images.</p>")
            files = gr.File(label="Field images", file_count="multiple",
                            file_types=["image", ".tif", ".tiff"], height=220)
            gr.HTML("<p class='lk-ex-label'>No images to hand? Load a public example case "
                    "(from the LeukoAI Zenodo dataset):</p>")
            with gr.Row():
                ex_buttons = {k: gr.Button(label, size="sm", variant="secondary", min_width=70)
                              for k, label in (("aml", "AML"), ("cll", "CLL"), ("cml", "CML"), ("normal", "Normal"))}
            with gr.Accordion("Settings", open=False):
                threshold = gr.Slider(0, 10, value=3, step=1, label="Image count threshold",
                                      info="Leukemia is reported when more than this many fields are classed "
                                           "leukemic. Published setting: 3.")
                show_heatmaps = gr.Checkbox(value=True, label="Show heatmaps on the most influential fields")
            consent = gr.Checkbox(label="I confirm these images are de-identified and used for research only")
            with gr.Row():
                run = gr.Button("Analyse case", variant="primary", scale=3)
                reset = gr.Button("Clear", variant="secondary", scale=1)
        with gr.Column(scale=7, min_width=320, elem_classes="lk-panel"):
            result = gr.HTML(PLACEHOLDER)
            gallery = gr.Gallery(label="Most influential fields (warm colours = where the model looked)",
                                 columns=3, rows=1, height=260, object_fit="cover", visible=False)
            with gr.Accordion("Per-field results", open=False, visible=False) as table_box:
                table = gr.Dataframe(interactive=False, wrap=True)
            download = gr.File(label="Download all results (CSV)", visible=False)
    gr.HTML(ABOUT)
    gr.HTML(FOOTER)

    outputs = [result, gallery, table_box, table, download]
    run.click(analyse, [files, threshold, show_heatmaps, consent], outputs)
    for k, b in ex_buttons.items():
        b.click(lambda k=k: load_example(k), None, [files, consent] + outputs)
    reset.click(clear, None, [files, consent] + outputs)

if __name__ == "__main__":
    demo.queue(max_size=20).launch(
        server_name=os.environ.get("HOST") or os.environ.get("GRADIO_SERVER_NAME") or "127.0.0.1",
        server_port=int(os.environ.get("PORT") or os.environ.get("GRADIO_SERVER_PORT") or 7860),
        theme=THEME, css=CSS, head=HEAD, allowed_paths=[ASSETS],
        favicon_path=FAVICON if os.path.exists(FAVICON) else None,
        auth=check_login,
        auth_message=(
            "LeukoAI — research use only. Access is restricted to registered users; credentials are issued "
            "by the study team. LeukoAI is not a medical device and must not be used for diagnosis or "
            "treatment decisions. By logging in you agree to upload de-identified images only and to use "
            "results for research purposes only."
        ),
    )