#!/usr/bin/env python3
"""CellQuant Gradio app (Hugging Face Space entry point).

Upload one or more fluorescence images (TIF/PNG/JPG; blue nuclei, green antibody,
red puncta) and get per-image counts of the six annotation categories, colour
overlays, and downloadable CSV/JSON of every detected object.
"""
from __future__ import annotations

import csv
import io
import json
import os
import sys
import tempfile
import zipfile

os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")

# ZeroGPU (Hugging Face Spaces): `spaces` must be imported before torch; its GPU decorator
# borrows a GPU for the duration of a call and is a no-op everywhere else.
try:
    import spaces  # noqa: E402
except ImportError:  # local runs and tests without the package
    spaces = None
ZERO_GPU = spaces is not None and str(os.environ.get("SPACES_ZERO_GPU", "")).lower() in ("1", "true")

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import time  # noqa: E402
from PIL import Image  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))  # run from any cwd / python -I

from cellquant import __version__  # noqa: E402
from cellquant.annotations import load_image
from cellquant.gemini_summary import enabled as gemini_enabled
from cellquant.gemini_summary import summarize as gemini_summarize
from cellquant.gemini_summary import summarize_comparison as gemini_summarize_comparison
from cellquant.infer import Predictor
from cellquant.labels import COUNT_COLUMNS, LABEL_COLORS, TOTAL_NUCLEI, TOTAL_PUNCTA, count_row
from cellquant.model import load_checkpoint
from cellquant.postprocess import PostConfig, analyze, render_overlay, result_to_json
from cellquant.validation import COMPARISON_HEADERS, COMPARISON_NOTE, DIFF_HEADERS, ValidationSet, compare, write_bundle

WEIGHTS = os.environ.get("CELLQUANT_WEIGHTS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights", "cellquant.pt"))
# On ZeroGPU the model must sit on cuda from start-up (CUDA is emulated outside GPU calls); elsewhere CPU unless told otherwise.
DEVICE = os.environ.get("CELLQUANT_DEVICE", "cuda" if ZERO_GPU else "cpu")
TTA = int(os.environ.get("CELLQUANT_TTA", "8"))     # 1, 4 or 8 passes per image; 8 scored best in cross-validation
TTA_LABEL = {1: "a single pass", 4: "4 flips", 8: "8 rotations and flips"}.get(TTA, f"{TTA} passes")

EX_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
VAL_DIR = os.environ.get("CELLQUANT_VALIDATION_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "validation"))
VALIDATION = ValidationSet(VAL_DIR)

MODEL = None
if os.path.exists(WEIGHTS):
    MODEL, _META = load_checkpoint(WEIGHTS)
    MODEL.to(DEVICE).eval()
_predictors: dict[int, Predictor] = {}
_val_probs: dict[tuple[str, int], dict] = {}   # (image key, tta) -> network probabilities, a few entries


def get_predictor(tta: int) -> Predictor:
    if tta not in _predictors:
        if MODEL is None:
            raise gr.Error(f"Model weights not found at {WEIGHTS}. Train with train.py or set CELLQUANT_WEIGHTS.")
        _predictors[tta] = Predictor(MODEL, device=DEVICE, tta=tta)
    return _predictors[tta]


def gpu(duration):
    """`@spaces.GPU(duration=...)` on ZeroGPU, identity elsewhere. ``duration`` may be a
    callable of the decorated function's arguments (seconds of GPU time to reserve;
    shorter reservations get a better place in the queue)."""
    if spaces is None:
        return lambda f: f
    return spaces.GPU(duration=duration)


def _device_name() -> str:
    if DEVICE.startswith("cuda"):
        return "GPU (ZeroGPU)" if ZERO_GPU else "GPU"
    return "Apple GPU" if DEVICE == "mps" else "CPU"


def _legend_html() -> str:
    items = "".join(
        f'<span style="display:inline-block;margin:2px 10px 2px 0"><span style="display:inline-block;width:14px;height:14px;'
        f'background:rgb{c};border:1px solid #888;vertical-align:middle;margin-right:4px"></span>{lab}</span>'
        for lab, c in LABEL_COLORS.items()
    )
    return f"<div>{items}</div>"


def _gpu_seconds(n_images: int) -> int:
    """GPU time to reserve on ZeroGPU: a fixed start-up share plus a per-image share that grows with TTA."""
    return int(min(600, 15 + {1: 4, 4: 6, 8: 10}.get(TTA, 10) * max(1, n_images)))


def _safe_progress(progress, *args, **kwargs) -> None:
    try:
        progress(*args, **kwargs)
    except Exception:  # noqa: BLE001 - a progress bar must never break a run (e.g. inside a GPU worker)
        pass


def _gemini_section(text: str | None) -> str:
    if not gemini_enabled():
        return "*Gemini summary is off: add a GEMINI_API_KEY secret to this Space to enable it.*"
    return "### Summary (Gemini)\n" + (text or "")


@gpu(duration=lambda files, *a, **k: _gpu_seconds(len(files or [])))
def _analyze_gpu(files, antibody_thr, cellnos_thr, min_cell_area, nostril_rule, cellnos_rule, progress=gr.Progress()):
    """The GPU part of analyze_files: predictions, overlays and the download bundle."""
    t_start = time.time()
    tta = TTA
    cfg = PostConfig(antibody_thr=float(antibody_thr), cellnos_thr=float(cellnos_thr),
                     min_cell_area=int(min_cell_area), nostril_rule=nostril_rule, cellnos_rule=cellnos_rule)
    pred = get_predictor(tta)
    rows, gallery, per_image_json, csv_objects = [], [], {}, []
    tmpdir = tempfile.mkdtemp(prefix="cellquant_")
    overlays = {}
    warnings_: list[str] = []
    seen: dict[str, int] = {}
    for i, f in enumerate(files):
        path = f if isinstance(f, str) else f.name
        name = os.path.basename(path)
        k = seen.get(name, 0)
        seen[name] = k + 1
        if k:  # two uploads with the same file name: keep both
            stem, ext = os.path.splitext(name)
            name = f"{stem}_{k + 1}{ext}"
        _safe_progress(progress, (i, len(files)), desc=f"analysing {name}")
        try:
            img = load_image(path)
        except Exception as e:  # noqa: BLE001
            raise gr.Error(f"{name}: cannot read image ({e})")
        if np.array_equal(img[..., 0], img[..., 1]) and np.array_equal(img[..., 1], img[..., 2]):
            warnings_.append(f"{name}: all three channels are identical (greyscale input); the model needs blue nuclei, green antibody and red puncta.")
        res, _ = pred.predict(img, cfg)
        ov = render_overlay(img, res, thickness=1)
        overlays[name] = ov
        out_png = os.path.join(tmpdir, os.path.splitext(name)[0] + "_overlay.png")
        Image.fromarray(ov).save(out_png)
        row = count_row(res.counts)
        n_nuclei, n_puncta = row[COUNT_COLUMNS.index(TOTAL_NUCLEI)], row[COUNT_COLUMNS.index(TOTAL_PUNCTA)]
        gallery.append((out_png, f"{name}: {n_nuclei} nuclei, {n_puncta} puncta"))
        rows.append([name] + row)
        per_image_json[name] = {"image_size_hw": list(img.shape[:2]), **result_to_json(res),
                                "totals": {TOTAL_NUCLEI: n_nuclei, TOTAL_PUNCTA: n_puncta}}
        for c in res.cells:
            csv_objects.append([name, "nucleus", c["id"], c["label"], round(c["area_px"], 1), round(c["centroid_xy"][0], 1), round(c["centroid_xy"][1], 1), round(c["p_antibody"], 3), round(c["p_nostril"], 3), c["n_puncta_inside"]])
        for n_ in res.nostrils:
            csv_objects.append([name, "punctum", n_["id"], n_["label"], round(n_["area_px"], 1), round(n_["centroid_xy"][0], 1), round(n_["centroid_xy"][1], 1), "", round(n_["p_in_nucleus"], 3), n_["cell_id"]])
    # downloads
    counts_csv = os.path.join(tmpdir, "counts.csv")
    with open(counts_csv, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["image"] + list(COUNT_COLUMNS))
        w.writerows(rows)
    objects_csv = os.path.join(tmpdir, "objects.csv")
    with open(objects_csv, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["image", "kind", "id", "label", "area_px", "centroid_x", "centroid_y", "p_antibody", "p_nostril_or_in_nucleus", "n_puncta_inside_or_cell_id"])
        w.writerows(csv_objects)
    json_path = os.path.join(tmpdir, "results.json")
    with open(json_path, "w") as fh:
        json.dump({"cellquant_version": __version__, "weights": os.path.basename(WEIGHTS), "tta": tta,
                   "postprocess": cfg.__dict__, "images": per_image_json}, fh, indent=1)
    zip_path = os.path.join(tmpdir, "cellquant_results.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in (counts_csv, objects_csv, json_path):
            zf.write(p, os.path.basename(p))
        for p, _ in gallery:
            zf.write(p, os.path.basename(p))
    summary = (f"Analysed {len(rows)} image{'s' if len(rows) != 1 else ''} in {time.time() - t_start:.1f} s "
               f"on the {_device_name()}, each averaged over {TTA_LABEL}.")
    summary += "".join(f"\n\nWarning: {w}" for w in warnings_)
    return rows, gallery, zip_path, summary


def analyze_files(files, antibody_thr, cellnos_thr, min_cell_area, nostril_rule, cellnos_rule, progress=gr.Progress()):
    """Section 1: counts, overlays and downloads for uploaded images, then the Gemini write-up
    (outside the GPU reservation, so the wait for Gemini costs no GPU time)."""
    if not files:
        raise gr.Error("Upload at least one image.")
    rows, gallery, zip_path, summary = _analyze_gpu(files, antibody_thr, cellnos_thr, min_cell_area, nostril_rule, cellnos_rule)
    _safe_progress(progress, 0.95, desc="writing the summary")
    text = None
    if gemini_enabled():
        text = gemini_summarize([{"image": r[0], "counts": dict(zip(COUNT_COLUMNS, r[1:]))} for r in rows],
                                overlay=Image.open(gallery[0][0]))
    return rows, gallery, zip_path, summary + "\n\n" + _gemini_section(text)


@gpu(duration=lambda key, *a, **k: _gpu_seconds(1))
def _validate_gpu(key) -> dict:
    """The GPU part of validate_image: prediction, matching, overlays and bundle (plain data only)."""
    t_start = time.time()
    img = VALIDATION.image(key)
    R = VALIDATION.raster(key)
    ck = (key, TTA)
    if ck not in _val_probs:
        if len(_val_probs) >= 4:
            _val_probs.pop(next(iter(_val_probs)))
        _val_probs[ck] = get_predictor(TTA).probs(img)
    res = analyze(_val_probs[ck], PostConfig())
    cmp = compare(key, img, R, res)
    tmpdir = tempfile.mkdtemp(prefix="cellquant_val_")
    zip_path = write_bundle(cmp, tmpdir)
    return {
        "summary": cmp.summary, "rows": cmp.rows, "diffs": cmp.diffs, "zip": zip_path,
        "manual": os.path.join(tmpdir, f"{key}_manual.png"), "predicted": os.path.join(tmpdir, f"{key}_predicted.png"),
        "missed": os.path.join(tmpdir, f"{key}_missed_only.png"),
        "n_gt": (R.n_cells, R.n_nostrils), "n_pred": (len(res.cells), len(res.nostrils)),
        "n_missed": (sum(1 for d in cmp.diffs if d[1] == "missed" and d[2] == "nucleus"),
                     sum(1 for d in cmp.diffs if d[1] == "missed" and d[2] == "punctum")),
        "seconds": time.time() - t_start,
    }


def validate_image(key, progress=gr.Progress()):
    """Section 2: run the model on one bundled annotated image and compare it with the annotation.

    Returns the summary (with the Gemini write-up), the three panels (manual / predicted /
    missed only, with their captions), the per-label table, the per-object table and the download."""
    if not key or key not in VALIDATION.keys:
        raise gr.Error("Pick one of the annotated images.")
    _safe_progress(progress, 0.1, desc=f"running the model on {key}")
    r = _validate_gpu(key)
    manual = gr.Image(value=r["manual"], label=f"1. Manual annotation of {key}: {r['n_gt'][0]} nuclei, {r['n_gt'][1]} puncta")
    predicted = gr.Image(value=r["predicted"], label=f"2. Model prediction: {r['n_pred'][0]} nuclei, {r['n_pred'][1]} puncta")
    missed = gr.Image(value=r["missed"], label=f"3. Missed only: {r['n_missed'][0]} nuclei and {r['n_missed'][1]} puncta annotated "
                                                f"but not predicted (red, numbered as in the list below)")
    _safe_progress(progress, 0.9, desc="writing the summary")
    text = None
    if gemini_enabled():
        text = gemini_summarize_comparison(key, r["rows"], COMPARISON_HEADERS, overlay=Image.open(r["missed"]))
    summary = (r["summary"] + f"\n* **Run time:** {r['seconds']:.1f} s on the {_device_name()}, averaged over {TTA_LABEL}."
               + "\n\n" + _gemini_section(text))
    return summary, manual, predicted, missed, r["rows"], r["diffs"], r["zip"]


def pick_validation_image(evt: gr.SelectData) -> str:
    """Gallery click -> key of the chosen annotated image (validate_image runs next)."""
    return VALIDATION.keys[evt.index]


FORCE_LIGHT_JS = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.href = url.href;
    }
}
"""

CSS = """
body, .gradio-container { background: #fff !important; }
.gradio-container .prose, .gradio-container .prose * { color: #000; }
.gradio-container .prose a { color: #000; text-decoration: underline; }
.gradio-container table, .gradio-container td, .gradio-container th { color: #000; }
footer { display: none !important; }
"""


def build_demo() -> gr.Blocks:
    theme = gr.themes.Monochrome(primary_hue="neutral", neutral_hue="neutral",
                                 font=["Helvetica Neue", "Helvetica", "Arial", "sans-serif"],
                                 font_mono=["Menlo", "Consolas", "monospace"])
    with gr.Blocks(title="AI Cell Quantification", theme=theme, css=CSS, js=FORCE_LIGHT_JS) as demo:
        gr.Markdown(
            f"# AI Cell Quantification\n"
            f"CellQuant {__version__}. "
            "Counts nuclei (antibody positive or negative), nuclei containing a red punctum (\"Nostril\") and puncta "
            "inside or outside nuclei in RGB fluorescence images (blue = DAPI, green = antibody, red = puncta). "
            "Trained on CVAT outlines of 35-day spinal-cord sections stained for Iba1 or TMEM119."
        )
        gr.HTML(_legend_html())
        # ---- section 1: the user's own images
        gr.Markdown(f"## 1. Analyse your images\nEvery image is analysed as {TTA_LABEL} and the outputs averaged "
                    "(the setting that scored best in cross-validation). A written summary follows the counts.")
        with gr.Row():
            with gr.Column(scale=1):
                files = gr.File(label="Images (TIF / PNG / JPG)", file_count="multiple", file_types=["image", ".tif", ".tiff"])
                with gr.Accordion("Advanced", open=False):
                    antibody_thr = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Antibody + threshold")
                    cellnos_thr = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Nucleus 'Nostril' threshold")
                    min_cell_area = gr.Slider(10, 300, value=50, step=5, label="Minimum nucleus area (px)")
                    nostril_rule = gr.Radio(["head", "geometric"], value="head", label="Punctum type from: network head / position relative to nuclei")
                    cellnos_rule = gr.Radio(["contains", "either", "head"], value="contains", label="Nucleus gets 'Nostril' if: contains an in-nucleus punctum / that OR the network head / head only")
                run = gr.Button("Analyse", variant="primary")
                examples = sorted(os.path.join(EX_DIR, f) for f in os.listdir(EX_DIR)
                                  if f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff"))) if os.path.isdir(EX_DIR) else []
                if examples:
                    gr.Examples(examples=[[[e]] for e in examples], inputs=[files], label="Example (512 px crop of a training image)")
            with gr.Column(scale=2):
                table = gr.Dataframe(headers=["image"] + list(COUNT_COLUMNS), label="Counts", interactive=False, wrap=True)
                gallery = gr.Gallery(label="Overlays", columns=2, height="auto")
                zip_out = gr.File(label="Download results (CSV + JSON + overlays)")
                summary = gr.Markdown()
        run.click(analyze_files, [files, antibody_thr, cellnos_thr, min_cell_area, nostril_rule, cellnos_rule],
                  [table, gallery, zip_out, summary], api_name="analyze_files")
        # ---- section 2: the annotated images, manual vs predicted vs missed
        if VALIDATION.keys:
            n = len(VALIDATION.keys)
            gr.Markdown(
                f"## 2. Check the model against the {n} annotated images\n"
                f"Click one of the {n} images. The model runs on it ({TTA_LABEL}, averaged) and three "
                "panels appear: the **manual annotation**, "
                "the **model prediction** (both in the label colours above) and **missed only**, the annotated objects the "
                "model did not find, outlined in red with their row number from the list at the bottom. "
                "The final model was trained on these images, so this shows its fit, not its accuracy on new images "
                "(those numbers are in the README). A written summary of the comparison follows the panels."
            )
            val_key = gr.Textbox(visible=False)
            val_gallery = gr.Gallery(value=VALIDATION.gallery(), label="Annotated images (click one)", columns=5,
                                     height="auto", allow_preview=False, object_fit="contain", interactive=False)
            val_summary = gr.Markdown("*No image selected yet.*")
            with gr.Row():
                img_manual = gr.Image(label="1. Manual annotation", type="filepath", interactive=False, height=460)
                img_pred = gr.Image(label="2. Model prediction", type="filepath", interactive=False, height=460)
                img_missed = gr.Image(label="3. Missed only (annotated, not predicted)", type="filepath", interactive=False, height=460)
            gr.Markdown("Open a panel full screen with the button in its corner. " + COMPARISON_NOTE)
            val_table = gr.Dataframe(headers=COMPARISON_HEADERS, label="Annotated vs predicted, per label", interactive=False,
                                     wrap=True, column_widths=["25%", "12%", "12%", "12%", "11%", "12%", "8%", "8%"])
            val_diffs = gr.Dataframe(headers=DIFF_HEADERS, interactive=False, wrap=True, max_height=600,
                                     label="Every missed, extra or wrongly labelled object, numbered as on the panels (x, y in pixels from the top-left corner)",
                                     column_widths=["4%", "9%", "8%", "15%", "15%", "6%", "6%", "6%", "31%"])
            val_zip = gr.File(label="Download the five overlays (manual, predicted, missed only, extra only, all differences) and both tables")
            val_gallery.select(pick_validation_image, None, val_key, api_name=False).then(
                validate_image, [val_key], [val_summary, img_manual, img_pred, img_missed, val_table, val_diffs, val_zip],
                api_name="validate_image")
        gr.Markdown(
            "Nucleus outlines come from a watershed over the network's distance and class maps. "
            "A nucleus is labelled Nostril when a punctum classified as in-nucleus lies inside it. "
            "Images larger than 1536 px are processed in overlapping tiles. Nothing leaves the Space except, "
            "the counts and one overlay thumbnail sent to Gemini for the written summary (only when the Space has a Gemini key). "
            + ("Inference runs on a shared GPU (Hugging Face ZeroGPU): no account is needed, each visitor has a daily "
               "allowance of GPU minutes and one image takes a few seconds." if ZERO_GPU else f"Inference runs on the {_device_name()}.")
        )
    return demo


if __name__ == "__main__":
    if MODEL is None:
        raise SystemExit(f"CellQuant weights not found at {WEIGHTS}: run train.py --out weights/cellquant.pt or set CELLQUANT_WEIGHTS")
    build_demo().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
                        allowed_paths=[EX_DIR], show_error=True)