"""
Hugging Face Space entrypoint: head-to-head Gradio baseline benchmark
for the VoxPixel paper (Section V).

This is a SEPARATE Space from the production VoxPixel app. Same repo,
same Dockerfile, same dependencies; the entrypoint is selected by the
APP_FILE environment variable in the Space's "Variables and secrets"
settings (set APP_FILE=app_bench_space.py for this Space).

Purpose: let any reviewer reproduce the V1/V2/V3 vs. VoxPixel head-to-
head comparison from the browser. Upload one or more NIfTI volumes,
choose how many trials per volume, click "Run benchmark"; the app runs
the same protocol used in Table III of the paper, displays a summary
inline, and writes a downloadable CSV.

Failure modes the experiment exposes (paper Section II / V):
    V1: gr.Image rejects .nii / .nii.gz at upload (NIfTI ingestion).
    V2: single-shot full-volume forward pass -> CUDA OOM (no sliding window).
    V3: + manually-wired MONAI sliding window, but the output mask
        is written with affine = identity, biasing liver volume in ml
        (no spatial-metadata round-trip in the Gradio I/O contract).
"""

import csv
import gc
import os
import sys
import tempfile
from pathlib import Path

import gradio as gr
import torch

from app_gradio_baseline import predict_v2, predict_v3
from bench_gradio_baseline import (
    VOXPIXEL_TRUTH_ML,
    run_v1,
    run_with_safety,
)


CSV_FIELDS = [
    "variant", "trial", "volume", "size_mb",
    "status", "runtime_s", "peak_vram_gb",
    "voxel_count", "naive_volume_ml", "true_volume_ml",
    "voxpixel_volume_ml", "abs_err_naive_vs_voxpixel_ml",
    "error",
]


def _file_path(file_obj) -> str:
    """Gradio 4.x returns a list of objects with .name; older versions
    sometimes pass a path string directly. Handle both defensively."""
    return getattr(file_obj, "name", str(file_obj))


def _row(variant: str, trial: int, vol_path: str, res: dict) -> dict:
    vol_name = Path(vol_path).name
    size_mb = round(os.path.getsize(vol_path) / (1024 ** 2), 2)
    naive_ml = res.get("naive_volume_ml")
    voxpixel = VOXPIXEL_TRUTH_ML.get(vol_name)
    if (voxpixel is not None and naive_ml is not None
            and naive_ml == naive_ml):  # NaN check
        abs_err = abs(naive_ml - voxpixel)
    else:
        abs_err = float("nan")
    return {
        "variant": variant,
        "trial": trial,
        "volume": vol_name,
        "size_mb": size_mb,
        "status": res.get("status", "?"),
        "runtime_s": round(res.get("runtime_s", 0.0), 3),
        "peak_vram_gb": round(res.get("peak_vram_gb", 0.0), 2),
        "voxel_count": res.get("voxel_count", 0),
        "naive_volume_ml": naive_ml,
        "true_volume_ml": res.get("true_volume_ml"),
        "voxpixel_volume_ml": voxpixel,
        "abs_err_naive_vs_voxpixel_ml": abs_err,
        "error": res.get("error", ""),
    }


def _summary_markdown(rows: list) -> str:
    by_variant: dict = {}
    for r in rows:
        v = r["variant"]
        d = by_variant.setdefault(v, {
            "n": 0, "ok": 0, "oom": 0, "ui_reject": 0, "err": 0,
            "rt_sum": 0.0, "vram_max": 0.0,
            "naive_err_sum": 0.0, "naive_err_n": 0,
        })
        d["n"] += 1
        s = r["status"]
        if s == "OK":
            d["ok"] += 1
            d["rt_sum"] += r["runtime_s"]
            d["vram_max"] = max(d["vram_max"], r["peak_vram_gb"])
            err = r["abs_err_naive_vs_voxpixel_ml"]
            if err == err:
                d["naive_err_sum"] += err
                d["naive_err_n"] += 1
        elif s == "OOM":
            d["oom"] += 1
        elif s == "UI_REJECT":
            d["ui_reject"] += 1
        else:
            d["err"] += 1

    lines = [
        "### Per-variant summary",
        "",
        "| Variant | Success / Total | OOM | UI reject | Other err | "
        "Mean runtime (s) | Peak VRAM (GB) | "
        "Mean \\|naive − VoxPixel\\| (ml) |",
        "|---|---|---|---|---|---|---|---|",
    ]
    for v, d in by_variant.items():
        mean_rt = (d["rt_sum"] / d["ok"]) if d["ok"] else float("nan")
        mean_err = (
            (d["naive_err_sum"] / d["naive_err_n"])
            if d["naive_err_n"] else float("nan")
        )
        lines.append(
            f"| **{v}** | {d['ok']}/{d['n']} | {d['oom']} | {d['ui_reject']} "
            f"| {d['err']} | {mean_rt:.2f} | {d['vram_max']:.2f} "
            f"| {mean_err:.2f} |"
        )
    return "\n".join(lines)


def _table_data(rows: list) -> list:
    return [
        [r["variant"], r["trial"], r["volume"], r["size_mb"],
         r["status"], r["runtime_s"], r["peak_vram_gb"],
         r["voxel_count"], r["naive_volume_ml"], r["true_volume_ml"],
         r["voxpixel_volume_ml"], r["abs_err_naive_vs_voxpixel_ml"]]
        for r in rows
    ]


def run_benchmark(files, trials: int, modality: str,
                  progress=gr.Progress()):
    if not files:
        return ("**Please upload at least one .nii / .nii.gz volume.**",
                None, None, "")

    paths = [_file_path(f) for f in files]
    log_lines: list = []

    progress(0.0, desc="Loading model...")
    try:
        from app_gradio_baseline import _get_model
        _get_model(modality)
    except Exception as e:
        return (f"**Failed to load model: {e}**\n\n"
                "Check that the Space hardware has a GPU and that "
                "mamba_ssm / selective_scan_cuda_oflex built successfully.",
                None, None, "")

    log_lines.append(
        f"Model loaded ({modality}). Hardware: "
        f"{torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}"
    )

    rows: list = []
    runners = [
        ("V1", lambda p, m: run_v1(p, m)),
        ("V2", lambda p, m: run_with_safety(predict_v2, p, m)),
        ("V3", lambda p, m: run_with_safety(predict_v3, p, m)),
    ]
    total_calls = len(paths) * trials * len(runners)
    done = 0

    for vol_path in paths:
        vol_name = Path(vol_path).name
        for variant, runner in runners:
            for trial in range(1, int(trials) + 1):
                msg = (f"[{variant}] trial {trial}/{int(trials)} on "
                       f"{vol_name}")
                log_lines.append(msg)
                progress(done / total_calls, desc=msg)
                res = runner(vol_path, modality)
                rows.append(_row(variant, trial, vol_path, res))
                done += 1
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                gc.collect()

    csv_path = os.path.join(
        tempfile.gettempdir(), "gradio_baseline_results.csv"
    )
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=CSV_FIELDS)
        w.writeheader()
        w.writerows(rows)

    summary_md = _summary_markdown(rows)
    return summary_md, _table_data(rows), csv_path, "\n".join(log_lines)


with gr.Blocks(title="VoxPixel — Head-to-Head Gradio Baseline") as demo:
    gr.Markdown("""
# VoxPixel — Head-to-Head Gradio Baseline (paper §V)

This Space reproduces the head-to-head experiment between three idiomatic
Gradio configurations and the production VoxPixel deployment, using the
**same SRMA-Mamba checkpoints** and the **same protocol** as Table III of
the paper.

| Variant | What it implements | Failure mode demonstrated |
| --- | --- | --- |
| **V1** | `gr.Image` interface | NIfTI ingestion not supported by Gradio's image components |
| **V2** | `gr.File` + nibabel + single-shot full-volume forward pass | CUDA OOM on full 3D volumes (no sliding window) |
| **V3** | `gr.File` + nibabel + manually-wired MONAI `SlidingWindowInferer` | Spatial metadata loss — output mask written with `affine = identity`, biasing liver volume in ml |

Upload one or more `.nii / .nii.gz` volumes, choose trials per volume, and
click **Run benchmark**. Three trials per volume across three volumes
reproduces the 9-call protocol used for VoxPixel.

Reference values for the affine-loss bias column come from VoxPixel
runs in the paper:
""")
    gr.Markdown(
        "\n".join(
            f"- `{k}` → **{v:.2f} ml** (VoxPixel)"
            for k, v in VOXPIXEL_TRUTH_ML.items()
        )
    )

    with gr.Row():
        with gr.Column(scale=2):
            files = gr.File(
                label="NIfTI volumes (.nii / .nii.gz)",
                file_count="multiple",
                file_types=[".nii", ".gz"],
            )
            trials = gr.Slider(
                label="Trials per volume",
                minimum=1, maximum=5, value=3, step=1,
            )
            modality = gr.Radio(
                label="Modality",
                choices=["T1", "T2"], value="T1",
            )
            run_btn = gr.Button(
                "Run head-to-head benchmark", variant="primary",
            )
            gr.Markdown(
                "First run takes ~30–60 s extra to load the model. "
                "Each V2/V3 trial typically takes 2–5 s on an L40S."
            )

        with gr.Column(scale=3):
            summary_md = gr.Markdown(label="Summary")
            results_df = gr.Dataframe(
                headers=[
                    "variant", "trial", "volume", "size_mb",
                    "status", "runtime_s", "peak_vram_gb",
                    "voxel_count", "naive_volume_ml", "true_volume_ml",
                    "voxpixel_volume_ml", "abs_err_naive_vs_voxpixel_ml",
                ],
                label="Per-run results",
                wrap=True,
            )
            csv_file = gr.File(label="Download results CSV")
            log_box = gr.Textbox(label="Log", lines=10)

    run_btn.click(
        fn=run_benchmark,
        inputs=[files, trials, modality],
        outputs=[summary_md, results_df, csv_file, log_box],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
    )