"""Benchmark the 4 rigid CT->T1c registration pipelines on N M6 test cases.

Pipelines (see src/registration_pipelines.py for the shared transform contract):
  simpleitk | elastix | antspyx | deepreg

Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
the T1c native grid.  Every pipeline is scored with the SAME crop-independent
geometric QA used in production (scripts/02_build_m6_dataset.py):

  - skull_ring_gap_mm  : concentricity of the warped-CT skull ring vs the MR
                         head silhouette (mid/low level + angular std).
                         Aligned = 0 <= mid <= 20 mm and std <= 12 mm.
  - tv_in_ct_frac      : fraction of the (dilated) warped GTV that falls inside
                         the warped-CT head tissue (the slab must straddle it).
  - tv_nvox            : warped GTV volume (>0.5 voxels).
  - accepted           : production rule (nvox>0 and tv_in_ct and concentric).
  - reg_wall_s         : wall-clock registration time (excludes I/O + scoring).

Because three of the four toolkits are not pip-installable into the main
`longitudinal` env (elastix needs conda-forge; antspyx has no cp314 wheel;
PyPI deepreg is a broken stub), the benchmark is env-portable: run it under
whichever conda env has a given toolkit installed, and it records only the
pipelines that are importable there.  The `--merge` mode then combines the
per-pipeline JSON files (accumulated across envs) into one summary + plot.

Usage (repo root):
  # main env (SimpleITK baseline; elastix/deepreg report as unavailable)
  source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
  python scripts/bench_registration.py --pipelines simpleitk elastix deepreg

  # isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
  source /opt/conda/etc/profile.d/conda.sh && conda activate <ants312>
  python scripts/bench_registration.py --pipelines antspyx simpleitk

  # combine everything produced so far into one table + plot
  python scripts/bench_registration.py --merge
"""
import argparse
import importlib.util
import json
import os
import sys
import time
from datetime import datetime
from concurrent.futures import ProcessPoolExecutor, as_completed

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import numpy as np
import SimpleITK as sitk
from scipy import ndimage

from src.common import DATA, d, read_nii_img
from src import registration_pipelines as RP


def _m6():
    return RP._load_m6()


# --------------------------------------------------------------------------- #
# case selection
# --------------------------------------------------------------------------- #

def ct_complete(ct):
    """Production guard: skip corrupt CTs (upper brain missing)."""
    m6 = _m6()
    vz = m6._ct_vertex_slice(ct)
    if vz is None:
        return False
    sp = ct.GetSpacing()[2]
    per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2))
    iz = vz - int(round(100 / sp))
    return bool(0 <= iz < len(per) and per[iz] > 5000)


def select_cases(args):
    """Pick N labeled M6 cases deterministically.  No CT I/O here -- each file
    is read exactly once, inside the registration worker (NFS is the bottleneck)."""
    m6 = _m6()
    labeled, _unlabeled = m6.discover()
    print(f"[cases] m6 labeled candidates (CT+T1c+TV): {len(labeled)}")
    pool = list(labeled)
    if args.sample == "random":
        idx = np.random.default_rng(args.seed).choice(
            len(pool), min(args.n_cases, len(pool)), replace=False)
        sel = [pool[i] for i in sorted(idx.tolist())]
    else:
        sel = pool[:args.n_cases]
    print(f"[cases] selected {len(sel)} test cases (sample={args.sample} seed={args.seed})")
    for r in sel:
        print(f"   - {r['key']}")
    return sel


# --------------------------------------------------------------------------- #
# shared evaluation (identical for every pipeline)
# --------------------------------------------------------------------------- #

def evaluate(ct, t1c, tv, T):
    """Warp CT + GTV onto the T1c grid with T and score them (production QA)."""
    m6 = _m6()
    ha = sitk.GetArrayFromImage(t1c)
    warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
    warped_tv = sitk.GetArrayFromImage(sitk.Resample(tv, t1c, T, sitk.sitkNearestNeighbor, 0.0))
    gap = m6._skull_ring_gap(warped_ct, ha)
    posw = warped_ct[warped_ct > 0]
    tissue_w = (warped_ct > max(2.0, float(np.percentile(posw, 2)))) if posw.size \
        else np.zeros_like(warped_ct, bool)
    tissue_w = ndimage.binary_opening(tissue_w, iterations=2)
    nvox = int((warped_tv > 0.5).sum())
    tv_in_ct_frac = 1.0
    if nvox > 0:
        tvd = ndimage.binary_dilation(warped_tv > 0.5, iterations=3)
        tv_in_ct_frac = float((tvd & tissue_w).sum() / max(tvd.sum(), 1))
    ok_gap = bool(gap) and 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0
    accepted = nvox > 0 and tv_in_ct_frac >= 0.5 and ok_gap
    return {
        "skull_ring_gap_mm": gap,
        "tv_nvox": nvox,
        "tv_in_ct_frac": round(tv_in_ct_frac, 3),
        "alignment": "concentric" if ok_gap else "not_concentric",
        "accepted": bool(accepted),
    }


def _call(pipeline, ct, t1c, tv, args):
    if pipeline == "simpleitk":
        return RP.register_simpleitk(ct, t1c, threads=args.threads, mode=args.simpleitk_mode)
    if pipeline == "elastix":
        return RP.register_elastix(ct, t1c, workdir=args.elastix_dir, threads=args.threads)
    if pipeline == "antspyx":
        return RP.register_antspyx(ct, t1c)
    if pipeline == "deepreg":
        return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
                                   checkpoint=args.deepreg_ckpt)
    raise KeyError(pipeline)


def _run_case(pipeline, case, args):
    """Worker body: read once + register + score one case.  Runs in a forked
    process.  The CT-completeness guard is applied here (on the already-read CT)
    so each file is read exactly once."""
    key = case["key"]
    ct = read_nii_img(case["ct"])
    t1c = read_nii_img(case["img"])
    tv = read_nii_img(case["label"])
    out = {"key": key, "subject": case["subject"], "date": case["date"]}
    if not ct_complete(ct):
        out.update({"skull_ring_gap_mm": None, "tv_nvox": 0, "tv_in_ct_frac": None,
                    "alignment": "ct_incomplete", "accepted": False,
                    "note": "CT head incomplete (data issue, not a pipeline failure)"})
        return out
    t0 = time.time()
    T, meta = _call(pipeline, ct, t1c, tv, args)
    out["reg_wall_s"] = round(time.time() - t0, 1)
    out.update(meta)
    out.update(evaluate(ct, t1c, tv, T))
    return out


def run_pipeline(pipeline, cases, args):
    avail, note = RP.availability().get(pipeline, (False, "unknown"))
    if not avail:
        row = {"pipeline": pipeline, "available": False, "reason": note,
               "n_cases": len(cases), "cases": []}
        print(f"[{pipeline}] UNAVAILABLE: {note}")
        return row
    print(f"[{pipeline}] available ({note}); registering {len(cases)} cases "
          f"with {args.workers} workers")
    results = []
    errors = []
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(_run_case, pipeline, c, args): c for c in cases}
        for fu in as_completed(futs):
            c = futs[fu]
            try:
                results.append(fu.result())
                r = results[-1]
                g = r.get("skull_ring_gap_mm")
                reg = f"reg={r.get('reg_wall_s')}s" if r.get("reg_wall_s") is not None else "incomplete-CT"
                print(f"   {c['key']}: accepted={r['accepted']} "
                      f"gap={g['mid_mm'] if g else 'None'}mm "
                      f"std={g['std_mm'] if g else 'None'}mm "
                      f"tv_in_ct={r.get('tv_in_ct_frac')} {reg}", flush=True)
            except Exception as e:  # noqa
                errors.append({"key": c["key"], "error": repr(e)})
                print(f"   {c['key']}: ERROR {e!r}", flush=True)
    results.sort(key=lambda x: x["key"])
    # aggregate
    registered = [r for r in results if r.get("reg_wall_s") is not None]
    incomplete = [r for r in results if r.get("alignment") == "ct_incomplete"]
    accepted = [r for r in registered if r["accepted"]]
    gaps_mid = [r["skull_ring_gap_mm"]["mid_mm"] for r in registered if r.get("skull_ring_gap_mm")]
    gaps_std = [r["skull_ring_gap_mm"]["std_mm"] for r in registered if r.get("skull_ring_gap_mm")]
    regt = [r["reg_wall_s"] for r in registered]
    tvfr = [r["tv_in_ct_frac"] for r in registered if r.get("tv_in_ct_frac") is not None]
    first = (registered or results)[0] if (registered or results) else {}
    row = {
        "pipeline": pipeline, "available": True, "toolkit": first.get("toolkit"),
        "metric": first.get("metric"), "transform": first.get("transform"),
        "n_cases": len(cases), "n_registered": len(registered),
        "n_ct_incomplete": len(incomplete), "n_error": len(errors),
        "n_accepted": len(accepted),
        "accepted_rate": round(len(accepted) / max(len(cases), 1), 3),
        "reg_wall_s": {"mean": round(float(np.mean(regt)), 1), "median": round(float(np.median(regt)), 1),
                       "max": round(float(np.max(regt)), 1)} if regt else None,
        "skull_ring_gap_mid_mm": {"mean": round(float(np.mean(gaps_mid)), 1),
                                  "median": round(float(np.median(gaps_mid)), 1)} if gaps_mid else None,
        "skull_ring_gap_std_mm": {"mean": round(float(np.mean(gaps_std)), 1),
                                  "median": round(float(np.median(gaps_std)), 1)} if gaps_std else None,
        "tv_in_ct_frac": {"mean": round(float(np.mean(tvfr)), 3)} if tvfr else None,
        "cases": results, "errors": errors,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
    }
    return row


# --------------------------------------------------------------------------- #
# reporting / merge
# --------------------------------------------------------------------------- #

def _write(row, out_dir):
    p = os.path.join(out_dir, row["pipeline"] + ".json")
    os.makedirs(out_dir, exist_ok=True)
    with open(p, "w") as f:
        json.dump(row, f, indent=1)
    return p


def _fmt(v, spec="{}", default="-"):
    return default if v is None else spec.format(v)


def print_table(rows):
    order = [p for p in RP.ALL_PIPELINES if any(r["pipeline"] == p for r in rows)]
    print("\n" + "=" * 100)
    hdr = (f"{'pipeline':10s} {'avail':6s} {'toolkit':22s} {'n':>3s} {'ok':>3s} {'err':>4s} "
           f"{'accept':>7s} {'gapMid':>7s} {'gapStd':>7s} {'s/reg':>7s} {'tvInCT':>7s}")
    print(hdr)
    print("-" * 100)
    for p in order:
        r = next(x for x in rows if x["pipeline"] == p)
        if not r.get("available"):
            print(f"{p:10s} {('no'):6s} {'-':22s} {r['n_cases']:>3d} {0:>3d} {'-':>4s} "
                  f"{'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s}   [{r['reason']}]")
            continue
        gm = (r.get("skull_ring_gap_mid_mm") or {}).get("median")
        gs = (r.get("skull_ring_gap_std_mm") or {}).get("median")
        rt = (r.get("reg_wall_s") or {}).get("median")
        tv = (r.get("tv_in_ct_frac") or {}).get("mean")
        print(f"{p:10s} {'yes':6s} {(r.get('toolkit') or '')[:22]:22s} {r['n_cases']:>3d} "
              f"{r.get('n_registered', 0):>3d} {r['n_error']:>4d} {r['accepted_rate']:>7.2f} "
              f"{_fmt(gm, '{:.1f}'):>7s} {_fmt(gs, '{:.1f}'):>7s} {_fmt(rt, '{:.1f}'):>7s} "
              f"{_fmt(tv, '{:.2f}'):>7s}")
    print("=" * 100)


def write_plot(rows, out_png):
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except Exception as e:  # noqa
        print("plot skipped:", repr(e))
        return
    ok_rows = {r["pipeline"]: r for r in rows if r.get("available")}
    if not ok_rows:
        print("no available pipelines to plot")
        return
    labels = list(ok_rows)
    gap_mid = [ok_rows[p].get("skull_ring_gap_mid_mm", {}).get("median") or 0 for p in labels]
    gap_std = [ok_rows[p].get("skull_ring_gap_std_mm", {}).get("median") or 0 for p in labels]
    regt = [ok_rows[p].get("reg_wall_s", {}).get("median") or 0 for p in labels]
    acc = [ok_rows[p].get("accepted_rate") or 0 for p in labels]
    fig, axes = plt.subplots(1, 4, figsize=(18, 4))
    colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD"]
    for ax, (title, vals, fmt, ylab, inv) in zip(axes, [
        ("Skull-ring gap mid (mm, median)", gap_mid, "{:.1f}", "mm", True),
        ("Skull-ring gap std (mm, median)", gap_std, "{:.1f}", "mm", True),
        ("Registration time (s, median)", regt, "{:.0f}", "s", True),
        ("Acceptance rate", acc, "{:.2f}", "frac", False)]):
        bars = ax.bar(range(len(vals)), vals, color=colors[:len(vals)])
        for b, v in zip(bars, vals):
            ax.annotate(fmt.format(v), (b.get_x() + b.get_width() / 2, b.get_height()),
                        ha="center", va="bottom", fontsize=8)
        ax.set_title(title, fontsize=9)
        ax.set_ylabel(ylab)
        ax.set_xticks(range(len(vals)))
        ax.set_xticklabels(labels, fontsize=9)
    plt.suptitle("CT->T1c registration benchmark (M6, cross-modality rigid)", fontsize=11)
    fig.tight_layout()
    fig.savefig(out_png, dpi=120)
    plt.close(fig)
    print("plot:", out_png)


def merge(args):
    out_dir = os.path.join(args.out_dir)
    rows = []
    for p in RP.ALL_PIPELINES:
        fp = os.path.join(out_dir, p + ".json")
        if os.path.exists(fp):
            with open(fp) as f:
                rows.append(json.load(f))
    if not rows:
        raise SystemExit(f"no pipeline results found in {out_dir}")
    print(f"[merge] {len(rows)} pipeline result files in {out_dir}")
    print_table(rows)
    summary = {
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "protocol": ("rigid CT->T1c; cross-modality; scored by production skull-ring "
                     "concentricity (aligned: 0<=mid<=20mm, std<=12mm) + GTV-in-CT fraction "
                     "(>=0.5); reg time excludes I/O+scoring; per-pipeline best-of-T vs T^-1 "
                     "orientation gate for elastix/antspyx"),
        "n_cases": next((r["n_cases"] for r in rows if r.get("available")), None),
        "case_keys": sorted({c["key"] for r in rows if r.get("available") for c in r.get("cases", [])}),
        "results": rows,
    }
    sf = os.path.join(out_dir, "summary.json")
    with open(sf, "w") as f:
        json.dump(summary, f, indent=1)
    print("summary:", sf)
    write_plot(rows, os.path.join(out_dir, "bench_registration.png"))
    # markdown table for the report
    md = ["| pipeline | available | toolkit | cases | ok | err | accepted | gap mid (mm) | gap std (mm) | s/reg (median) | tv-in-CT |",
          "|---|---|---|---|---|---|---|---|---|---|---|"]
    for p in [x["pipeline"] for x in rows]:
        r = next(x for x in rows if x["pipeline"] == p)
        if not r.get("available"):
            md.append(f"| {p} | NO | - | {r['n_cases']} | 0 | - | - | - | - | - | - | ({r['reason']}) |")
            continue
        md.append(
            f"| {p} | yes | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | "
            f"{r['n_error']} | {r['accepted_rate']:.2f} | "
            f"{(r.get('skull_ring_gap_mid_mm') or {}).get('median')} | "
            f"{(r.get('skull_ring_gap_std_mm') or {}).get('median')} | "
            f"{(r.get('reg_wall_s') or {}).get('median')} | "
            f"{(r.get('tv_in_ct_frac') or {}).get('mean')} |")
    mf = os.path.join(out_dir, "summary.md")
    with open(mf, "w") as f:
        f.write("\n".join(md) + "\n")
    print("markdown:", mf)


def parse_args():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pipelines", default=" ".join(RP.ALL_PIPELINES),
                    help="space-separated subset of: " + " ".join(RP.ALL_PIPELINES))
    ap.add_argument("--n-cases", type=int, default=10, help="number of M6 test cases (default 10)")
    ap.add_argument("--sample", choices=["first", "random"], default="first",
                    help="first N by discover order, or a seeded random sample of N")
    ap.add_argument("--seed", type=int, default=0, help="seed for --sample random")
    ap.add_argument("--keys", default=None, help="comma-separated explicit case keys (overrides selection)")
    ap.add_argument("--workers", type=int, default=4, help="cases registered in parallel")
    ap.add_argument("--threads", type=int, default=12, help="per-registration CPU threads")
    ap.add_argument("--simpleitk-mode", choices=["benchmark", "production"], default="benchmark",
                    help="benchmark = single pyramid (comparable budget to other toolkits); "
                         "production = full multi-start tilt grid (slow)")
    ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration"))
    ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)")
    ap.add_argument("--deepreg-dir", default=None, help="deepreg scratch dir")
    ap.add_argument("--deepreg-ckpt", default=None, help="path to a DeepReg 3D pretrained checkpoint")
    ap.add_argument("--merge", action="store_true", help="combine existing per-pipeline JSONs")
    args = ap.parse_args()
    return args


def main():
    args = parse_args()
    if args.merge:
        merge(args)
        return
    args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix")
    args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg")
    pipelines = [p.lower() for p in args.pipelines.split()]
    unknown = set(pipelines) - set(RP.ALL_PIPELINES)
    if unknown:
        raise SystemExit(f"unknown pipeline(s): {unknown}")

    if args.keys:
        m6 = _m6()
        labeled, _ = m6.discover()
        bykey = {r["key"]: r for r in labeled}
        missing = [k for k in args.keys.split(",") if k not in bykey]
        if missing:
            raise SystemExit(f"keys not found in m6 labeled: {missing}")
        cases = [bykey[k.strip()] for k in args.keys.split(",")]
        print(f"[cases] using {len(cases)} explicit keys")
    else:
        cases = select_cases(args)

    print(f"\n[bench] pipelines={pipelines} n_cases={len(cases)} workers={args.workers} "
          f"threads={args.threads} out_dir={args.out_dir}")
    rows = []
    for p in pipelines:
        try:
            rows.append(run_pipeline(p, cases, args))
        except RP.PipelineUnavailable as e:
            rows.append({"pipeline": p, "available": False, "reason": str(e),
                         "n_cases": len(cases), "cases": []})
            print(f"[{p}] UNAVAILABLE: {e}")
    for r in rows:
        fp = _write(r, args.out_dir)
        print(f"[{r['pipeline']}] wrote {fp}")
    print_table(rows)
    # run-level availability snapshot (which env saw what)
    avail = RP.availability()
    af = os.path.join(args.out_dir, "_availability.json")
    os.makedirs(args.out_dir, exist_ok=True)
    with open(af, "w") as f:
        json.dump({"timestamp": datetime.now().isoformat(timespec="seconds"),
                   "env_python": sys.version.split()[0],
                   "exec": sys.executable, "availability": avail}, f, indent=1)
    print(f"[bench] availability in this env: " +
          ", ".join(f"{k}={'yes' if v[0] else 'no'}" for k, v in avail.items()))
    print(f"[bench] done. Re-run with --merge to combine across envs.")


if __name__ == "__main__":
    main()