"""Benchmark the 6 DEFORMABLE CT->T1c registration pipelines on N M6 cases.

Companion to ``scripts/bench_registration.py`` (rigid).  Pipelines:
  simpleitk | elastix | antspyx | synthmorph | impact | niftyreg
(each = its toolkit's native deformable refinement initialized from the
toolkit's own rigid result; see src/registration_pipelines_deformable.py
and REGISTRATION_BENCHMARK_DEFORMABLE.md).

Headline scores (identical for every pipeline):
  - mi_ct_mr     : achieved Mattes MI (nats), warped-CT ~ MR over the MR head
                   (higher = better).  Also computed for the pipeline's own
                   rigid base (mi_base) -> d_mi = the deformable gain.
  - impact_loss  : IMPACT semantic loss (TS/M730 features), warped-CT ~ MR
                   over the full volumes (lower = better).  Also mi/IL for
                   the base (il_base).
  - jac_*        : Jacobian-determinant statistics of the FINAL deformable
                   transform over the MR head (finite differences on the
                   T1c grid): p1/p50/mean + fraction < 0 (folded), < 0.3,
                   < 0.5.  A healthy diffeomorphic refinement has J ~ 1.
  - jac_gtv_*    : same Jacobian statistics restricted to the warped-GTV
                   region (does the deformation fold the tumor?).
  - disp_base_*  : |T_final(x) - T_base(x)| over the MR head (p50/p95/max,
                   mm) -- the deformable increment.
  - tv_vol_ratio : warped-GTV volume final / base (a diffeomorphism
                   preserves volume, ~1).
  - reg_wall_s   : wall-clock of the deformable stage (excludes I/O, the
                   rigid base and scoring).
Production QA (skull-ring gap, accepted) is still computed per case and kept
in the JSON.

Usage (repo root):
  source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
  python scripts/bench_registration_deformable.py --pipelines "simpleitk elastix antspyx niftyreg"
  python scripts/bench_registration_deformable.py --pipelines "impact synthmorph"
  python scripts/bench_registration_deformable.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 src.common import DATA, d, read_nii_img, head_mask_from_image
from src import registration_pipelines as RP
from src import registration_pipelines_deformable as RPD

_spec = importlib.util.spec_from_file_location("br", os.path.join(ROOT, "scripts", "bench_registration.py"))
BR = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(BR)

JAC_STEP_MM = 1.0  # finite-difference half-step for the Jacobian estimate


# --------------------------------------------------------------------------- #
# deformable-specific scoring
# --------------------------------------------------------------------------- #

def _dense_H(T, ref):
    """Absolute mapping H(x) = T(x) at every voxel of ``ref`` (z,y,x,3)."""
    return RPD.densify_sitk(T, ref)[2]


def jacobian_of_T(T, ref, step_mm=JAC_STEP_MM):
    """det dT/dx of an arbitrary contract transform over ``ref``'s grid.

    Central finite differences with world-axis steps of ~``step_mm`` on the
    densified absolute mapping (linear resampling of coordinate volumes);
    border voxels (no both-side neighbour) are NaN.  Returns J (z,y,x)."""
    H = RPD.densify_sitk(T, ref)[2]
    sp = np.array(ref.GetSpacing())  # (x, y, z); array axis i has spacing sp[2-i]
    off = np.array([int(np.clip(np.round(step_mm / sp[2 - i]), 1, None))
                    for i in range(3)])
    H = np.ascontiguousarray(H, dtype='float64')
    dH = []
    for i in range(3):
        d = np.full(H.shape, np.nan)
        sl_p = [slice(None)] * 3; sl_m = [slice(None)] * 3; sl_mid = [slice(None)] * 3
        sl_p[i] = slice(2 * off[i], None)
        sl_m[i] = slice(None, -2 * off[i])
        sl_mid[i] = slice(off[i], -off[i])
        d[tuple(sl_mid)] = (H[tuple(sl_p)] - H[tuple(sl_m)]) / (2.0 * off[i] * sp[2 - i])
        dH.append(d)
    # array axis i <-> world axis (2-i): axis 2 = x, axis 1 = y, axis 0 = z.
    # J = det[dH/dx, dH/dy, dH/dz]
    A = dH[2]; B = dH[1]; C = dH[0]
    J = (A[..., 0] * (B[..., 1] * C[..., 2] - B[..., 2] * C[..., 1])
         - A[..., 1] * (B[..., 0] * C[..., 2] - B[..., 2] * C[..., 0])
         + A[..., 2] * (B[..., 0] * C[..., 1] - B[..., 1] * C[..., 0]))
    return J


def _stats(J, mask):
    v = J[mask]
    v = v[np.isfinite(v)]
    if v.size < 50:
        return None
    return {
        "p1": round(float(np.percentile(v, 1)), 4),
        "p50": round(float(np.median(v)), 4),
        "mean": round(float(v.mean()), 4),
        "frac_neg": round(float((v < 0).mean()), 5),
        "frac_lt03": round(float((v < 0.3).mean()), 5),
        "frac_lt05": round(float((v < 0.5).mean()), 5),
        "nvox": int(v.size),
    }


def deformable_metrics(ct, t1c, tv, T_final, T_base):
    """Jacobian / displacement-increment / GTV-preservation metrics of a
    deformable result (final transform vs its rigid base)."""
    ha = sitk.GetArrayFromImage(t1c).astype('float32')
    hm = head_mask_from_image(ha)
    J = jacobian_of_T(T_final, t1c)
    jac_head = _stats(J, hm)
    # warped GTV mask (final)
    wtv = sitk.GetArrayFromImage(
        sitk.Resample(tv, t1c, T_final, sitk.sitkNearestNeighbor, 0.0))
    gtv = wtv > 0.5
    if gtv.sum() > 50:
        gtv = gtv & hm
        jac_gtv = _stats(J, gtv)
    else:
        jac_gtv = None
    # displacement increment |T_final - T_base| over the head
    Hf = _dense_H(T_final, t1c)
    Hb = _dense_H(T_base, t1c)
    inc = np.linalg.norm(Hf - Hb, axis=-1)
    inc_stats = {
        "p50": round(float(np.median(inc[hm])), 2),
        "p95": round(float(np.percentile(inc[hm], 95)), 2),
        "max": round(float(inc[hm].max()), 2),
    }
    # GTV volume preservation (final vs base)
    wtv_b = sitk.GetArrayFromImage(
        sitk.Resample(tv, t1c, T_base, sitk.sitkNearestNeighbor, 0.0))
    nv_f = int((wtv > 0.5).sum()); nv_b = int((wtv_b > 0.5).sum())
    tv_vol_ratio = round(nv_f / nv_b, 3) if nv_b > 0 else None
    return {"jac_head": jac_head, "jac_gtv": jac_gtv,
            "disp_base_mm": inc_stats, "tv_vol_ratio": tv_vol_ratio,
            "gtv_nvox_final": nv_f, "gtv_nvox_base": nv_b}


# --------------------------------------------------------------------------- #
# case worker
# --------------------------------------------------------------------------- #

def _call(pipeline, ct, t1c, args):
    if pipeline == "simpleitk":
        return RPD.register_simpleitk_deformable(ct, t1c, threads=args.threads)
    if pipeline == "elastix":
        return RPD.register_elastix_deformable(ct, t1c, workdir=args.elastix_dir,
                                               threads=args.threads)
    if pipeline == "antspyx":
        return RPD.register_antspyx_deformable(ct, t1c, threads=args.threads)
    if pipeline == "niftyreg":
        return RPD.register_niftyreg_deformable(ct, t1c, threads=args.threads,
                                                gpu=args.niftyreg_gpu,
                                                backend=args.niftyreg_backend)
    if pipeline == "synthmorph":
        return RPD.register_synthmorph_deformable(ct, t1c, gpu=args.synthmorph_gpu,
                                                  extent=args.synthmorph_extent)
    if pipeline == "impact":
        return RPD.register_impact_deformable(
            ct, t1c, gpu=args.impact_gpu,
            iters=args.impact_iters or RP.IMPACT_DEFAULT_ITERS,
            samples=args.impact_samples or RP.IMPACT_DEFAULT_SAMPLES,
            threads=args.threads)
    raise KeyError(pipeline)


def _run_case(pipeline, case, args):
    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 BR.ct_complete(ct):
        out.update({"skull_ring_gap_mm": None, "tv_nvox": 0, "tv_in_ct_frac": None,
                    "impact_loss": None, "alignment": "ct_incomplete",
                    "accepted": False,
                    "note": "CT head incomplete (data issue, not a pipeline failure)"})
        return out
    t0 = time.time()
    T_final, T_base, meta = _call(pipeline, ct, t1c, args)
    out["reg_wall_s"] = round(time.time() - t0, 1)
    out.update(meta)
    # final + base alignment scores
    out.update(BR.evaluate(ct, t1c, tv, T_final))
    out["base"] = BR.evaluate(ct, t1c, tv, T_base)
    out["d_mi"] = (round(out["mi_ct_mr"] - out["base"]["mi_ct_mr"], 4)
                   if out["mi_ct_mr"] is not None and out["base"]["mi_ct_mr"] is not None
                   else None)
    out["d_il"] = (round(out["impact_loss"] - out["base"]["impact_loss"], 4)
                   if out["impact_loss"] is not None and out["base"]["impact_loss"] is not None
                   else None)
    out.update(deformable_metrics(ct, t1c, tv, T_final, T_base))
    return out


def run_pipeline(pipeline, cases, args):
    avail, note = RPD.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]
                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']}: mi={r.get('mi_ct_mr')} (base {r.get('base', {}).get('mi_ct_mr')}) "
                      f"il={r.get('impact_loss')} "
                      f"jac_p1={(r.get('jac_head') or {}).get('p1')} {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"])
    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.get("accepted")]

    def _med(vals, nd=4):
        vals = [v for v in vals if v is not None]
        return round(float(np.median(vals)), nd) if vals else None

    row = {
        "pipeline": pipeline, "available": True,
        "toolkit": (registered or [{}])[0].get("toolkit"),
        "metric": (registered or [{}])[0].get("metric"),
        "transform": (registered or [{}])[0].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": {
            "median": _med([r["reg_wall_s"] for r in registered], 1),
            "mean": round(float(np.mean([r["reg_wall_s"] for r in registered])), 1) if registered else None,
            "max": round(float(np.max([r["reg_wall_s"] for r in registered])), 1) if registered else None,
        },
        "mi_ct_mr": {"median": _med([r.get("mi_ct_mr") for r in registered]),
                     "mean": round(float(np.mean([r["mi_ct_mr"] for r in registered if r.get("mi_ct_mr") is not None])), 4) if any(r.get("mi_ct_mr") is not None for r in registered) else None},
        "mi_base": {"median": _med([r.get("base", {}).get("mi_ct_mr") for r in registered])},
        "d_mi": {"median": _med([r.get("d_mi") for r in registered])},
        "impact_loss": {"median": _med([r.get("impact_loss") for r in registered]),
                        "mean": round(float(np.mean([r["impact_loss"] for r in registered if r.get("impact_loss") is not None])), 4) if any(r.get("impact_loss") is not None for r in registered) else None},
        "il_base": {"median": _med([r.get("base", {}).get("impact_loss") for r in registered])},
        "d_il": {"median": _med([r.get("d_il") for r in registered])},
        "jac_p1": {"median": _med([(r.get("jac_head") or {}).get("p1") for r in registered])},
        "jac_frac_neg": {"median": _med([(r.get("jac_head") or {}).get("frac_neg") for r in registered], 5)},
        "jac_frac_lt05": {"median": _med([(r.get("jac_head") or {}).get("frac_lt05") for r in registered], 5)},
        "jac_gtv_p1": {"median": _med([(r.get("jac_gtv") or {}).get("p1") for r in registered])},
        "disp_base_p95_mm": {"median": _med([(r.get("disp_base_mm") or {}).get("p95") for r in registered], 2)},
        "tv_vol_ratio": {"median": _med([r.get("tv_vol_ratio") for r in registered])},
        "cases": results, "errors": errors,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
    }
    return row


# --------------------------------------------------------------------------- #
# reporting
# --------------------------------------------------------------------------- #

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


def print_table(rows):
    order = [p for p in RPD.ALL_PIPELINES if any(r["pipeline"] == p for r in rows)]
    print("\n" + "=" * 118)
    print(f"{'pipeline':10s} {'avail':5s} {'n':>3s} {'ok':>3s} {'err':>3s} {'s/reg':>7s} "
          f"{'MI fin':>7s} {'MI base':>8s} {'dMI':>7s} {'IL fin':>7s} {'IL base':>8s} "
          f"{'jacP1':>7s} {'J<0':>7s} {'J<.5':>6s} {'incr p95':>9s}")
    print("-" * 118)
    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':5s} {r['n_cases']:>3d} {0:>3d} {'-':>3s}   [{r['reason'][:60]}]")
            continue
        print(f"{p:10s} {'yes':5s} {r['n_cases']:>3d} {r.get('n_registered', 0):>3d} "
              f"{r.get('n_error', 0):>3d} {_fmt((r.get('reg_wall_s') or {}).get('median'), '{:.0f}'):>7s} "
              f"{_fmt((r.get('mi_ct_mr') or {}).get('median'), '{:.3f}'):>7s} "
              f"{_fmt((r.get('mi_base') or {}).get('median'), '{:.3f}'):>8s} "
              f"{_fmt((r.get('d_mi') or {}).get('median'), '{:+.3f}'):>7s} "
              f"{_fmt((r.get('impact_loss') or {}).get('median'), '{:.3f}'):>7s} "
              f"{_fmt((r.get('il_base') or {}).get('median'), '{:.3f}'):>8s} "
              f"{_fmt((r.get('jac_p1') or {}).get('median'), '{:.3f}'):>7s} "
              f"{_fmt((r.get('jac_frac_neg') or {}).get('median'), '{:.4f}'):>7s} "
              f"{_fmt((r.get('jac_frac_lt05') or {}).get('median'), '{:.3f}'):>6s} "
              f"{_fmt((r.get('disp_base_p95_mm') or {}).get('median'), '{:.1f}'):>9s}")
    print("=" * 118)


def merge(args):
    out_dir = args.out_dir
    rows = []
    for p in RPD.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": ("deformable CT->T1c (each toolkit's native deformable refinement "
                      "initialized from its own rigid result); alignment scored by "
                      "achieved Mattes MI (nats, warped-CT ~ MR over the MR head, higher "
                      "= better) and the IMPACT semantic loss (TS/M730, full volumes, "
                      "lower = better), plus Jacobian-determinant statistics of the "
                      "final transform (finite differences on the T1c grid), the "
                      "deformable displacement increment vs the rigid base, and GTV "
                      "volume preservation; same contract / shared rules as the rigid "
                      "benchmark (full volumes, T1c fixed / CT moving, no mask)"),
        "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)
    md = ["| pipeline | toolkit | cases | ok | s/reg (med) | MI (med) | MI base (med) | ΔMI (med) | ImpactLoss (med) | IL base (med) | ΔIL (med) | jac p1 (med) | frac J<0 (med) | frac J<0.5 (med) | incr p95 (med, mm) | tv vol ratio (med) |",
          "|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|"]
    for r in rows:
        if not r.get("available"):
            md.append(f"| {r['pipeline']} | - | {r['n_cases']} | 0 | - | - | - | - | - | - | - | - | - | - | - | - | ({r['reason']}) |")
            continue
        md.append(
            f"| {r['pipeline']} | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | "
            f"{_fmt((r.get('reg_wall_s') or {}).get('median'), '{:.0f}')} | "
            f"{_fmt((r.get('mi_ct_mr') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('mi_base') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('d_mi') or {}).get('median'), '{:+.3f}')} | "
            f"{_fmt((r.get('impact_loss') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('il_base') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('d_il') or {}).get('median'), '{:+.3f}')} | "
            f"{_fmt((r.get('jac_p1') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('jac_frac_neg') or {}).get('median'), '{:.4f}')} | "
            f"{_fmt((r.get('jac_frac_lt05') or {}).get('median'), '{:.3f}')} | "
            f"{_fmt((r.get('disp_base_p95_mm') or {}).get('median'), '{:.1f}')} | "
            f"{_fmt((r.get('tv_vol_ratio') or {}).get('median'), '{:.3f}')} |")
    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(RPD.ALL_PIPELINES))
    ap.add_argument("--n-cases", type=int, default=10)
    ap.add_argument("--sample", choices=["first", "random"], default="first")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--keys", default=None)
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--threads", type=int, default=24)
    ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration_deformable"))
    ap.add_argument("--elastix-dir", default=None)
    ap.add_argument("--synthmorph-gpu", type=int, default=None)
    ap.add_argument("--synthmorph-extent", type=int, default=256)
    ap.add_argument("--impact-gpu", type=int, default=None)
    ap.add_argument("--impact-iters", type=int, default=None)
    ap.add_argument("--impact-samples", type=int, default=None)
    ap.add_argument("--niftyreg-backend", choices=["auto", "conda", "docker"], default="auto")
    ap.add_argument("--niftyreg-gpu", type=int, default=None)
    ap.add_argument("--merge", action="store_true")
    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")
    pipelines = [p.lower() for p in args.pipelines.split()]
    unknown = set(pipelines) - set(RPD.ALL_PIPELINES)
    if unknown:
        raise SystemExit(f"unknown pipeline(s): {unknown}")
    if args.keys:
        m6 = RP._load_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 = BR.select_cases(args)
    print(f"\n[bench-def] 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:
        os.makedirs(args.out_dir, exist_ok=True)
        fp = os.path.join(args.out_dir, r["pipeline"] + ".json")
        with open(fp, "w") as f:
            json.dump(r, f, indent=1)
        print(f"[{r['pipeline']}] wrote {fp}")
    print_table(rows)
    avail = RP.availability()
    with open(os.path.join(args.out_dir, "_availability.json"), "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("[bench-def] done. Re-run with --merge to combine.")


if __name__ == "__main__":
    main()