"""
Headless benchmark for the head-to-head comparison in the VoxPixel paper
(Section V). Runs the same nine-call protocol used for VoxPixel
(Table III in the paper) for each of V1, V2, V3 and writes one CSV row
per call, plus a printed summary suitable for direct inclusion in the
paper.

Usage (single command):

    python bench_gradio_baseline.py \
        --volumes ./test_volumes/56.nii.gz \
                  ./test_volumes/58.nii.gz \
                  ./test_volumes/60.nii.gz \
        --trials 3 \
        --modality T1 \
        --out gradio_baseline_results.csv

The 9 runs match the 9 VoxPixel runs in Table III. Use the same
GPU and same checkpoints when reporting numbers.
"""

import argparse
import csv
import gc
import os
import time
from pathlib import Path

import torch

# Keep this import local so the script can still print V1's deterministic
# UI_REJECT row even if the model dependencies fail to load.
def _load_predict_fns():
    from app_gradio_baseline import predict_v2, predict_v3
    return predict_v2, predict_v3


# Reference numbers from the existing VoxPixel runs (Table III in paper).
# Used to compute the affine-loss bias ("Δ vs. VoxPixel") in the report.
VOXPIXEL_TRUTH_ML = {
    "56.nii.gz": 1156.13,
    "58.nii.gz": 896.77,
    "60.nii.gz": 904.30,
}


def _classify_error(exc: BaseException) -> tuple[str, str]:
    msg = str(exc)
    if isinstance(exc, getattr(torch.cuda, "OutOfMemoryError", RuntimeError)):
        if "out of memory" in msg.lower():
            return "OOM", msg[:200]
    if isinstance(exc, RuntimeError) and "out of memory" in msg.lower():
        return "OOM", msg[:200]
    return "ERROR", msg[:200]


def run_v1(volume_path: str, modality: str) -> dict:
    """gr.Image rejects .nii.gz before predict() is ever called.
    Recorded deterministically without launching the UI so the row
    is reproducible in CI."""
    return {
        "status": "UI_REJECT",
        "runtime_s": 0.0,
        "peak_vram_gb": 0.0,
        "voxel_count": 0,
        "naive_volume_ml": float("nan"),
        "true_volume_ml": float("nan"),
        "error": "gr.Image does not accept .nii / .nii.gz uploads",
    }


def run_with_safety(fn, volume_path: str, modality: str) -> dict:
    try:
        return fn(volume_path, modality=modality)
    except BaseException as exc:
        status, msg = _classify_error(exc)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        return {"status": status, "error": msg}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--volumes", nargs="+", required=True,
                        help="Paths to NIfTI test volumes.")
    parser.add_argument("--trials", type=int, default=3,
                        help="Trials per (variant, volume). 3 reproduces "
                             "the 9-call protocol in Table III.")
    parser.add_argument("--modality", default="T1")
    parser.add_argument("--out", default="gradio_baseline_results.csv")
    parser.add_argument("--variants", nargs="+", default=["V1", "V2", "V3"],
                        choices=["V1", "V2", "V3"])
    args = parser.parse_args()

    for v in args.volumes:
        if not os.path.exists(v):
            raise FileNotFoundError(f"Volume not found: {v}")

    predict_v2 = predict_v3 = None
    if any(v in ("V2", "V3") for v in args.variants):
        predict_v2, predict_v3 = _load_predict_fns()

    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),
    }

    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",
        "abs_err_true_vs_voxpixel_ml",
        "error",
    ]
    rows: list[dict] = []

    for variant in args.variants:
        runner = runners[variant]
        for vol_path in args.volumes:
            vol_name = Path(vol_path).name
            size_mb = round(os.path.getsize(vol_path) / (1024 ** 2), 2)
            for trial in range(1, args.trials + 1):
                print(f"[{variant}] trial {trial}/{args.trials} on {vol_name} "
                      f"({size_mb} MB)...")
                t0 = time.time()
                result = runner(vol_path, args.modality)
                elapsed = time.time() - t0

                naive_ml = result.get("naive_volume_ml")
                true_ml = result.get("true_volume_ml")
                voxpixel = VOXPIXEL_TRUTH_ML.get(vol_name)

                def _abs_err(estimate):
                    if (voxpixel is None or estimate is None
                            or estimate != estimate):  # NaN check
                        return float("nan")
                    return abs(estimate - voxpixel)

                row = {
                    "variant": variant,
                    "trial": trial,
                    "volume": vol_name,
                    "size_mb": size_mb,
                    "status": result.get("status", "?"),
                    "runtime_s": round(result.get("runtime_s", elapsed), 3),
                    "peak_vram_gb": round(result.get("peak_vram_gb", 0.0), 2),
                    "voxel_count": result.get("voxel_count", 0),
                    "naive_volume_ml": naive_ml,
                    "true_volume_ml": true_ml,
                    "voxpixel_volume_ml": voxpixel,
                    "abs_err_naive_vs_voxpixel_ml": _abs_err(naive_ml),
                    "abs_err_true_vs_voxpixel_ml": _abs_err(true_ml),
                    "error": result.get("error", ""),
                }
                rows.append(row)
                print(f"    -> status={row['status']}  "
                      f"runtime={row['runtime_s']}s  "
                      f"peak_vram={row['peak_vram_gb']}GB")

                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                gc.collect()

    with open(args.out, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)
    print(f"\nWrote {len(rows)} rows to {args.out}")

    print("\n=== Per-variant summary (paper Table) ===")
    by_variant: dict[str, dict] = {}
    for r in rows:
        v = r["variant"]
        d = by_variant.setdefault(v, {
            "n": 0, "ok": 0, "oom": 0, "ui_reject": 0, "err": 0,
            "runtime_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["runtime_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:  # not NaN
                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

    header = (
        f"{'Variant':<6} {'Success':>10} {'OOM':>5} {'UI_REJ':>7} "
        f"{'ERR':>5} {'mean_rt(s)':>11} {'peak_vram(GB)':>14} "
        f"{'mean_|naive-voxpixel|(ml)':>26}"
    )
    print(header)
    print("-" * len(header))
    for v, d in by_variant.items():
        mean_rt = (d["runtime_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")
        )
        print(
            f"{v:<6} {d['ok']:>4}/{d['n']:<5} "
            f"{d['oom']:>5} {d['ui_reject']:>7} {d['err']:>5} "
            f"{mean_rt:>11.2f} {d['vram_max']:>14.2f} "
            f"{mean_err:>26.2f}"
        )


if __name__ == "__main__":
    main()