"""Visual QA of TV label warping: registration method x label interpolation method.

For one M6 case (default: 2BBN5XA4 / 20250411):
  - fixed = T1c (MR, picked by scripts/02_build_m6_dataset.pick_t1c),
    moving = CT (RT/ct_image.nii.gz); the TV (RT/TV-*.nii.gz) lives in the CT grid.
  - registration via ANTsPy presets (src/registration_pipelines.py), each returning
    a transform T that maps T1c -> CT, so sitk.Resample(tv, t1c, T, ...) places
    the TV in MR (T1c) coordinates:
        ants_rigid  : Rigid[0.25] (Mattes, 500x500x200, 4x2x1)
        ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1)
        ants_syn    : SyN (Affine[0.25] + SyN[0.2,3,0] stages)
  - the binary TV is then warped ONTO THE 0.5mm RESAMPLED T1c GRID by
    nearest-neighbor (nn) and by SignedMaurerDistanceMap (mm, inside<0)
    -> {linear, B-spline order-3} interp -> threshold (smd <= thr), where the
    candidate thresholds are derived per case from the CT grid spacing
    (normal-distribution theory, see smd_threshold()): the mean spacing
    halved, thr = (sx + sy + sz) / 6; the registration transform is
    grid-independent, so it applies unchanged
  - one PNG per registration: 3 rows (axial / coronal / sagittal) x 4 columns
    (leftmost = bare MR, then nn + each SMD method at its threshold); all
columns show the SAME slice (the nn warp's largest-TV cross-section per
    view), ZOOMED to the TV, rendered on a 0.5mm display grid (MR linear,
    masks NN; metrics stay on the native grid)
    bounding box + margin (the NN box, shared by all columns); only the
    1-voxel TV contour is drawn (yellow, no fill), plus the volume +
    % of the original CT-grid TV volume in the column header.

Outputs (default results/label_interp_qa/<subject>_<date>/):
  <reg>_T.h5                     registration transform (cached; re-runs skip it)
  <reg>__<warp_key>_tv05.nii.gz  warped TV on the 0.5mm T1c grid (warp_key includes thr)
  <subject>_<date>_label_qa_<reg>.png   one 3x5 figure per registration
  summary.json                   per-combination metrics (nvox, mm3, dice, shifts)
                                 + the SMD thresholding volume-preservation check

Usage (repo root):
  source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
  python scripts/qa_label_interpolation.py
  python scripts/qa_label_interpolation.py --subject 2BBN5XA4 --date 20250411 \
      --regs ants_rigid ants_affine ants_syn
"""
import argparse
import json
import os
import sys
import time

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

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

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

M6 = RP._load_m6()

REGS = ["ants_rigid", "ants_affine", "ants_syn"]
METHOD_LABELS = {"nn": "NN", "smd_linear": "SMD + linear", "smd_bspline": "SMD + B-spline(3)"}
VIEWS = ["axial", "coronal", "sagittal"]
SMD_DEFAULT = 1e5  # out-of-grid fill for the SMD volume: far OUTSIDE (inside < 0)
ZOOM_MARGIN = 24  # voxels of context around the TV bbox (0.5mm grid = 12 mm)


def smd_threshold(spacing):
    """Final SMD threshold from the normal-distribution theory.

    The SMD zero plane is inset from the true boundary by half a voxel
    projected onto the boundary normal: s(n)/2 with
    s(n) = sx|nx| + sy|ny| + sz|nz|, which ranges over the bounds
        thr_lo = min(sx, sy, sz) / 2   (normal along the smallest spacing)
        thr_hi = max(sx, sy, sz) / 2   (normal along the largest spacing)
    for the axis-aligned normals a voxelized surface actually has.  The
    final single scalar threshold (best mean-abs-volume error on the 9-case
    batch) is the mean spacing halved:
        thr = (sx + sy + sz) / 6        (e.g. (0.586+0.586+1.25)/6 = 0.404)
    Returns (thr, thr_lo, thr_hi)."""
    s = [abs(float(x)) for x in spacing]
    lo, hi = 0.5 * min(s), 0.5 * max(s)
    return round(sum(s) / 6.0, 3), lo, hi


def cols(thr):
    """Figure / warp columns: NN, then each SMD method at the threshold."""
    return [("nn", None)] + [(m, thr) for m in ("smd_linear", "smd_bspline")]
DISP_SPACING = 0.5  # display-grid spacing (mm): MR/masks are resampled to this
# for the PNG so the zoomed slices render smoothly; all metrics stay on the
# native T1c grid


def warp_key(method, thr):
    return method if thr is None else f"{method}_thr{thr:g}"


def col_label(method, thr):
    if thr is None:
        return METHOD_LABELS[method]
    return f"{METHOD_LABELS[method]} (smd<={thr:g})"


def reg_call(reg, ct, t1c):
    """ANTsPy registration, fixed = T1c / moving = CT (the shared RP rule)."""
    if reg == "ants_rigid":
        return RP.register_antspyx(ct, t1c)
    if reg == "ants_affine":
        return RP.register_ants_affine_variant(ct, t1c, "antspyx")
    if reg == "ants_syn":
        return RP.register_ants_variant_deformable(ct, t1c, "ants_syn")
    raise KeyError(f"unknown registration {reg!r} (one of {REGS})")


def find_case(subject, date):
    labeled, _unlabeled = M6.discover()
    for r in labeled:
        if r.get("subject") == subject and r.get("date") == date:
            return r
    raise SystemExit(f"case {subject}/{date} not found among m6 labeled rows")


def get_transform(reg, ct, t1c, case_dir, use_cache):
    """Register (or load the cached transform), return (T, meta)."""
    tf = os.path.join(case_dir, reg + "_T.h5")
    if use_cache and os.path.exists(tf):
        print(f"[{reg}] loading cached transform {tf}", flush=True)
        return sitk.ReadTransform(tf), {"cached": True}
    t0 = time.time()
    T, meta = reg_call(reg, ct, t1c)
    dt = time.time() - t0
    sitk.WriteTransform(T, tf)
    print(f"[{reg}] registered in {dt:.0f}s ({meta.get('toolkit')}) "
          f"inverted={meta.get('auto_inverted')} gap={meta.get('skull_ring_gap_mm')}",
          flush=True)
    return T, {**meta, "seconds_wall": round(dt, 1), "cached": False}


def make_tv_bin(tv):
    """Binary 0/1 TV image on the TV's own (CT) grid."""
    a = (sitk.GetArrayFromImage(tv) > 0.5).astype(np.uint8)
    img = sitk.GetImageFromArray(a)
    img.CopyInformation(tv)
    return img


def warp_tv(tv_bin, t1c, T, method, thr=0.0):
    """Warp the binary TV onto the t1c grid with T (T1c -> CT). Returns a
    (z, y, x) uint8 0/1 mask (the interpolated label in MR coordinates).
    For the SMD methods the threshold is mask = (smd <= thr)."""
    rs = sitk.ResampleImageFilter()
    rs.SetReferenceImage(t1c)
    rs.SetTransform(T)
    if method == "nn":
        rs.SetInterpolator(sitk.sitkNearestNeighbor)
        rs.SetDefaultPixelValue(0)
        rs.SetOutputPixelType(sitk.sitkUInt8)
        return (sitk.GetArrayFromImage(rs.Execute(tv_bin)) > 0).astype(np.uint8)
    # signed distance to the TV boundary in mm: negative inside, positive
    # outside.  Note the SimpleITK 2.5 signature is
    # (image1, insideIsPositive=False, squaredDistance=True, useImageSpacing=False,
    #  backgroundValue=0.0) -- keywords are mandatory: the defaults would give
    # a SQUARED distance map, which interpolates wrongly.  On this build the
    # boundary (outer shell) voxels of the foreground get exactly 0, so the
    # mask is recovered with <= 0, not < 0.
    smd = sitk.SignedMaurerDistanceMap(
        tv_bin, insideIsPositive=False, squaredDistance=False, useImageSpacing=True)
    smd = sitk.Cast(smd, sitk.sitkFloat32)
    if method == "smd_linear":
        rs.SetInterpolator(sitk.sitkLinear)
    else:
        rs.SetInterpolator(sitk.sitkBSpline3)  # B-spline order 3
    rs.SetDefaultPixelValue(SMD_DEFAULT)
    rs.SetOutputPixelType(sitk.sitkFloat32)
    return (sitk.GetArrayFromImage(rs.Execute(smd)) <= thr).astype(np.uint8)


def best_slices(mask, shape):
    """(idx, area) of the largest TV cross-section per orthogonal view."""
    if not mask.any():
        return {"axial": (shape[0] // 2, 0),
                "coronal": (shape[1] // 2, 0),
                "sagittal": (shape[2] // 2, 0)}
    az = mask.sum(axis=(1, 2))  # per z (axial)
    ay = mask.sum(axis=(0, 2))  # per y (coronal)
    ax = mask.sum(axis=(0, 1))  # per x (sagittal)
    return {"axial": (int(np.argmax(az)), int(az.max())),
            "coronal": (int(np.argmax(ay)), int(ay.max())),
            "sagittal": (int(np.argmax(ax)), int(ax.max()))}


def resample_grid(img, spacing=0.5, nn=False):
    """Resample ``img`` onto an isotropic ``spacing`` grid (same origin and
    direction).  Linear for scalar images, nearest for label/mask images."""
    size = [max(1, int(round(s * sp / spacing)))
            for s, sp in zip(img.GetSize(), img.GetSpacing())]
    ref = sitk.Image(size, img.GetPixelID())
    ref.SetSpacing((spacing,) * 3)
    ref.SetOrigin(img.GetOrigin())
    ref.SetDirection(img.GetDirection())
    rs = sitk.ResampleImageFilter()
    rs.SetReferenceImage(ref)
    rs.SetTransform(sitk.Transform(3, sitk.sitkIdentity))
    rs.SetInterpolator(sitk.sitkNearestNeighbor if nn else sitk.sitkLinear)
    rs.SetDefaultPixelValue(0)
    rs.SetOutputPixelType(img.GetPixelID())
    return rs.Execute(img)


def largest_cc(mask):
    """The largest connected component of the mask (for slice selection and
    the zoom box when the TV has multiple separated targets).  Returns
    (cc_mask, n_components); an empty mask gives (empty, 0)."""
    lab, n = ndimage.label(mask)
    if n == 0:
        return mask.astype(np.uint8), 0
    if n == 1:
        return (lab > 0).astype(np.uint8), 1
    sizes = ndimage.sum(mask, lab, index=range(1, n + 1))
    return (lab == (int(np.argmax(sizes)) + 1)).astype(np.uint8), n


def tv_crop_box(mask, shape, margin):
    """(lo, hi) 3-D index box around the mask (per axis of (z, y, x)) with a
    margin; centered box when the mask is empty."""
    shape = np.array(shape)
    if mask.any():
        idx = np.argwhere(mask)
        lo, hi = idx.min(0), idx.max(0)
    else:
        lo = hi = shape // 2
    lo = np.clip(lo - margin, 0, shape - 1)
    hi = np.clip(hi + margin, 0, shape - 1)
    return lo, hi


def view_crop(arr, kind, idx, lo_b, hi_b):
    """Zoomed slice display around the TV box (bench_registration.make_qa_png
    convention: coronal / sagittal flipped so superior is up)."""
    if kind == "axial":
        s = arr[idx, lo_b[1]:hi_b[1] + 1, lo_b[2]:hi_b[2] + 1]
    elif kind == "coronal":
        s = arr[lo_b[0]:hi_b[0] + 1, idx, lo_b[2]:hi_b[2] + 1][::-1, :]
    else:
        s = arr[lo_b[0]:hi_b[0] + 1, lo_b[1]:hi_b[1] + 1, idx][::-1, :]
    return np.ascontiguousarray(s)


def dice3(a, b):
    a, b = a > 0, b > 0
    if not a.any() and not b.any():
        return 1.0
    return float(2 * (a & b).sum() / max(a.sum() + b.sum(), 1))


def centroid_mm(mask, spacing):
    c = np.array(ndimage.center_of_mass(mask)) * np.array(spacing)
    return c


def check_volume(tv_bin, t1c, source_vol_mm3, colspec):
    """Does thresholding the SMD preserve the TV volume?  Resampled with the
    IDENTITY transform (CT grid -> MR grid, no registration), which isolates
    the interpolation + thresholding step from registration error.  One row
    per figure column (NN + each SMD method at its threshold)."""
    T = sitk.Transform(3, sitk.sitkIdentity)
    sp = [float(s) for s in t1c.GetSpacing()]
    vv = sp[0] * sp[1] * sp[2]
    out = {}
    for method, thr in colspec:
        mask = warp_tv(tv_bin, t1c, T, method, thr or 0.0)
        out[warp_key(method, thr)] = round(float(mask.sum() * vv), 1)
    print(f"volume preservation (identity resample CT grid -> MR grid): "
          f"source TV = {source_vol_mm3:.0f} mm3", flush=True)
    for tag, v in out.items():
        print(f"   {tag:20s}: {v:8.1f} mm3  ({100.0 * v / source_vol_mm3:5.1f}%)",
              flush=True)
    return out


def make_png(mr, combos, out_png, title, spacing, window, source_vol):
    """One registration: 3 rows (views) x (len(combos)+1) columns; the
    leftmost column is the bare MR (no overlay).  All columns show the SAME
    slices (the NN warp's largest-area slice per view); volumes go in the
    column headers (row-independent): the original CT TV volume under the
    bare-MR column, "% of it" under each warp column."""
    lo, hi = window
    ncol, nrow = len(combos) + 1, len(VIEWS)
    cw, ch, gap, m = 3.0, 3.4, 0.15, 0.62
    top = 1.5
    fig_w = ncol * cw + (ncol - 1) * gap + 2 * m
    fig_h = nrow * ch + (nrow - 1) * gap + m + top
    fig = plt.figure(figsize=(fig_w, fig_h))
    fig.patch.set_facecolor("k")
    aspects = {"axial": spacing[1] / spacing[0],
               "coronal": spacing[2] / spacing[0],
               "sagittal": spacing[2] / spacing[1]}
    cmap_edge = ListedColormap(["#ffee00"])

    def w(a):
        return np.clip((a - lo) / max(hi - lo, 1e-9), 0.0, 1.0)

    for r, kind in enumerate(VIEWS):
        for c in range(ncol):
            x0 = (m + c * (cw + gap)) / fig_w
            y0 = (m + (nrow - 1 - r) * (ch + gap)) / fig_h
            axi = fig.add_axes([x0, y0, cw / fig_w, ch / fig_h])
            # every column (incl. the bare-MR one) shows the same slice/box
            ref = combos[0]["views"][kind]
            axi.imshow(w(ref["mr"]), cmap="gray", aspect=aspects[kind])
            if c > 0:
                v = combos[c - 1]["views"][kind]
                m2 = (v["mask"] > 0)
                if m2.any():
                    cont = m2 & ~ndimage.binary_erosion(m2)
                    axi.imshow(np.ma.masked_where(~cont, cont), cmap=cmap_edge,
                               aspect=aspects[kind])
                else:
                    axi.text(0.5, 0.5, "empty warp", transform=axi.transAxes,
                             ha="center", va="center", color="r", fontsize=8)
            axi.text(0.02, 0.02, f"{kind}={ref['idx']}", transform=axi.transAxes,
                     color="w", fontsize=7,
                     bbox=dict(facecolor="k", alpha=0.55, pad=1.0))
            axi.set_xticks([]); axi.set_yticks([])
            for s_ in axi.spines.values():
                s_.set_edgecolor("0.4")
    # column headers: bare MR + original CT TV volume, then method + % volume
    cx0 = (m + cw / 2) / fig_w
    fig.text(cx0, 1 - 0.74 / fig_h, "MR (no overlay)", color="#ffffff",
             fontsize=9.5, ha="center")
    fig.text(cx0, 1 - 1.02 / fig_h, f"CT TV: {source_vol:.0f} mm3",
             color="#b6f0b6", fontsize=9, ha="center")
    for c, comb in enumerate(combos):
        cx = (m + (c + 1) * (cw + gap) + cw / 2) / fig_w
        fig.text(cx, 1 - 0.74 / fig_h, comb["ilabel"], color="#9fd0ff",
                 fontsize=9.5, ha="center")
        fig.text(cx, 1 - 1.02 / fig_h,
                 f"{comb['vol_mm3']:.0f} mm3 ({comb['pct_of_source']:.0f}%)",
                 color="#b6f0b6", fontsize=9, ha="center")
    # view labels (left)
    for r, kind in enumerate(VIEWS):
        y0 = (m + (nrow - 1 - r) * (ch + gap)) / fig_h
        fig.text(0.012, y0 + ch / 2 / fig_h, kind.capitalize(), color="w", fontsize=10,
                 ha="center", va="center", rotation=90)
    fig.suptitle(title, color="w", fontsize=11, y=1 - 0.16 / fig_h)
    fig.savefig(out_png, facecolor="k", dpi=100)
    plt.close(fig)


def ct_complete(ct):
    """Production guard (copied from scripts/bench_registration.py): skip
    corrupt CTs (upper brain missing)."""
    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_rows(args):
    """Cases to run: an explicit --subject/--date, or the first / a seeded
    random N labeled M6 cases."""
    if args.n_cases:
        labeled, _unlabeled = M6.discover()
        pool = list(labeled)
        print(f"[cases] m6 labeled candidates (CT+T1c+TV): {len(pool)}")
        if args.sample == "random":
            rng = np.random.default_rng(args.seed)
            idx = rng.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)} (sample={args.sample} seed={args.seed})")
        for r in sel:
            print(f"   - {r['key']}")
        return sel
    return [find_case(args.subject, args.date)]


def run_case(row, args, out_base):
    ap_case_dir = args.out
    case_dir = ap_case_dir or os.path.join(out_base,
                                           f"{row['subject']}_{row['date']}")
    os.makedirs(case_dir, exist_ok=True)
    ct_path = row["ct"]
    t1c_path = row["img"]
    tv_path = args.tv or row["label"]

    print(f"case {row['key']}: ct={os.path.basename(ct_path)} "
          f"t1c={os.path.basename(t1c_path)} tv={os.path.basename(tv_path)}", flush=True)
    ct = read_nii_img(ct_path)
    t1c = read_nii_img(t1c_path)
    tv = read_nii_img(tv_path)
    tv_bin = make_tv_bin(tv)
    # the TV interpolation targets the 0.5mm resampled MR grid (same origin /
    # direction; the registration transform T is grid-independent, so it
    # applies unchanged)
    t1c05 = resample_grid(t1c, DISP_SPACING)
    mr05 = sitk.GetArrayFromImage(t1c05).astype("float32")
    spacing = [DISP_SPACING] * 3  # warp-grid spacing (x, y, z)
    print(f"t1c {t1c.GetSize()} @ {[float(s) for s in t1c.GetSpacing()]}mm -> warp grid "
          f"{t1c05.GetSize()} @ {DISP_SPACING}mm, tv nvox={int((sitk.GetArrayFromImage(tv_bin) > 0).sum())} "
          f"({int((sitk.GetArrayFromImage(tv_bin) > 0).sum()) * ct.GetSpacing()[0] * ct.GetSpacing()[1] * ct.GetSpacing()[2]:.0f} mm3 on the CT grid)",
          flush=True)

    # MR display window from the head
    pos = mr05 > 0
    thr = np.percentile(mr05[pos], 2) if pos.any() else 0.0
    head = mr05 > max(thr, 1e-6)
    lo, hi = (np.percentile(mr05[head], (2, 98)) if head.any() else (0.0, 1.0))

    # volume-preservation check for the SMD thresholding (no registration);
    # the SMD threshold is derived per case from the CT grid spacing
    sp_ct = [float(s) for s in tv_bin.GetSpacing()]
    source_vol = float((sitk.GetArrayFromImage(tv_bin) > 0).sum()) * sp_ct[0] * sp_ct[1] * sp_ct[2]
    THR, THR_LO, THR_HI = smd_threshold(sp_ct)
    colspec = cols(THR)
    print(f"smd threshold from CT spacing {['%.3f' % s for s in sp_ct]}: "
          f"bounds lo={THR_LO:.3f} hi={THR_HI:.3f} -> used (sx+sy+sz)/6 = {THR:.3f}",
          flush=True)
    vol_check = check_volume(tv_bin, t1c05, source_vol, colspec)

    combos = []
    reg_meta = {}
    for reg in args.regs:
        try:
            T, meta = get_transform(reg, ct, t1c, case_dir, use_cache=not args.no_cache)
        except Exception as e:
            reg_meta[reg] = {"error": repr(e)}
            print(f"[{reg}] FAILED: {e!r}", flush=True)
            continue
        reg_meta[reg] = meta
        reg_combos = []
        # warp every column (NN + each SMD method at the case threshold) -> 0.5mm grid
        for method, thr in colspec:
            key = warp_key(method, thr)
            p = os.path.join(case_dir, f"{reg}__{key}_tv05.nii.gz")
            if not args.no_cache and os.path.exists(p):
                print(f"[{reg}/{key}] loading cached warp {p}", flush=True)
                mask = (sitk.GetArrayFromImage(read_nii_img(p)) > 0).astype(np.uint8)
                dt = 0.0
            else:
                t0 = time.time()
                mask = warp_tv(tv_bin, t1c05, T, method, thr or 0.0)
                dt = time.time() - t0
            nvox = int(mask.sum())
            mimg = sitk.GetImageFromArray(mask)
            mimg.CopyInformation(t1c05)
            if not os.path.exists(p) or args.no_cache:
                sitk.WriteImage(mimg, p, True)
            reg_combos.append({"reg": reg, "method": method, "thr": thr,
                               "key": key, "ilabel": col_label(method, thr),
                               "mask": mask, "nvox": nvox,
                               "vol_mm3": round(nvox * spacing[0] * spacing[1] * spacing[2], 1),
                               "seconds_warp": round(dt, 1)})
            print(f"[{reg}/{key}] nvox={nvox} ({reg_combos[-1]['vol_mm3']} mm3) ({dt:.1f}s)",
                  flush=True)
        # all columns show the SAME slices and the SAME zoom box, selected from the
        # LARGEST connected component of the NN mask (when the TV has several
        # separated targets, only the biggest one drives the view), on the
        # 0.5mm warp grid
        nn = next(c for c in reg_combos if c["method"] == "nn")
        nn_cc, n_ccs = largest_cc(nn["mask"])
        if n_ccs > 1:
            print(f"[{reg}] TV has {n_ccs} components; slices/zoom follow the "
                  f"largest ({int(nn_cc.sum())} of {int(nn['mask'].sum())} voxels)",
                  flush=True)
        sl = best_slices(nn_cc, mr05.shape)
        lo_b, hi_b = tv_crop_box(nn_cc, mr05.shape, ZOOM_MARGIN)
        for c in reg_combos:
            c["slices"] = [sl[k][0] for k in VIEWS]
            c["views"] = {kind: {"mr": view_crop(mr05, kind, sl[kind][0], lo_b, hi_b),
                                 "mask": view_crop(c["mask"], kind, sl[kind][0], lo_b, hi_b),
                                 "idx": sl[kind][0], "area": sl[kind][1]} for kind in VIEWS}
            c["pct_of_source"] = 100.0 * c["vol_mm3"] / source_vol
            c["pct_vs_nn"] = (100.0 * c["vol_mm3"] / nn["vol_mm3"]
                              if nn["vol_mm3"] else 0.0)
            if c["method"] == "nn":
                c["dice_vs_nn"] = 1.0
                c["centroid_shift_mm_vs_nn"] = 0.0
            else:
                c["dice_vs_nn"] = round(dice3(nn["mask"], c["mask"]), 4)
                cs = centroid_mm(c["mask"], spacing) - centroid_mm(nn["mask"], spacing)
                c["centroid_shift_mm_vs_nn"] = [round(float(x), 2) for x in cs]
        combos.extend(reg_combos)
        # warped-volume table per registration (% of the original CT-grid TV volume)
        parts = " | ".join(f"{c['key']}={c['vol_mm3']:.0f} mm3 ({c['pct_of_source']:.0f}%)"
                           for c in reg_combos)
        print(f"warped volumes {reg} (CT TV={source_vol:.0f} mm3, "
              f"slices ax/cor/sag={nn['slices']}): {parts}", flush=True)

    pngs = {}
    for reg in args.regs:
        reg_combos = [c for c in combos if c["reg"] == reg]
        if not reg_combos:
            continue
        title = (f"{row['subject']} {row['date']} -- {reg}: TV warping QA "
                 f"(same slices = NN largest-area per view)\n"
                 f"fixed=T1c ({os.path.basename(t1c_path)[:40]}...) moving=CT  |  "
                 f"yellow = TV contour  |  MR window {lo:.0f}..{hi:.0f}")
        png = os.path.join(case_dir,
                           f"{row['subject']}_{row['date']}_label_qa_{reg}.png")
        make_png(mr05, reg_combos, png, title,
                 [DISP_SPACING] * 3, (float(lo), float(hi)), source_vol)
        pngs[reg] = png
        print(f"png: {png}", flush=True)

    summary = {
        "case": row["key"], "subject": row["subject"], "date": row["date"],
        "inputs": {"ct": ct_path, "t1c": t1c_path, "tv": tv_path},
        "t1c_geometry": {"size": list(t1c.GetSize()),
                         "spacing": [float(s) for s in t1c.GetSpacing()],
                         "origin": [float(o) for o in t1c.GetOrigin()]},
        "warp_grid": {"spacing": DISP_SPACING, "size": list(t1c05.GetSize()),
                      "note": "the TV interpolation (all methods) runs on this grid"},
        "smd_threshold": {"ct_spacing": [round(s, 3) for s in sp_ct],
                          "lo": round(THR_LO, 3), "hi": round(THR_HI, 3),
                          "formula": "(sx+sy+sz)/6", "used": round(THR, 3)},
        "regs": {k: {kk: vv for kk, vv in v.items() if kk != "stages"}
                  for k, v in reg_meta.items()},
        "volume_check": {"source_tv_mm3": round(source_vol, 1),
                         "identity_resample_mm3": vol_check},
        "pngs": pngs,
        "combos": {f"{c['reg']}/{c['key']}": {
            "nvox": c["nvox"], "vol_mm3": c["vol_mm3"],
            "pct_of_source_tv": round(c["pct_of_source"], 1),
            "slices_ax_cor_sag": c["slices"],
            "dice_vs_nn": c["dice_vs_nn"],
            "centroid_shift_mm_vs_nn": c["centroid_shift_mm_vs_nn"],
            "seconds_warp": c["seconds_warp"]} for c in combos},
    }
    sf = os.path.join(case_dir, "summary.json")
    with open(sf, "w") as f:
        json.dump(summary, f, indent=1)
    print(f"summary: {sf}", flush=True)
    return summary


def print_batch_table(rows, regs):
    """Compact cross-case table of the warp volumes (% of the CT TV)."""
    short = {"nn": "nn", "smd_linear": "lin", "smd_bspline": "bspl"}

    def key_label(key):
        if key == "nn":
            return "nn"
        m, _, t = key.partition("_thr")
        return short[m] + f"<={t}"

    def reg_keys(s, r):
        return [k.split("/", 1)[1] for k in s.get("combos", {})
                if k.startswith(r + "/")]

    width = 28 + 7 + 49 * len(regs)
    print("\n" + "=" * width)
    first = next((s for s in rows if not s.get('error')), None)
    print(f"{'case':28s} {'CT TV':>6s} | " +
          " | ".join(f"{r:48s}" for r in regs)
          + ("   (% of CT TV; per reg: " +
             " ".join(key_label(k) for k in reg_keys(first, regs[0])) + ")"
             if first else ""))
    print("-" * width)
    for s in rows:
        if s.get("error"):
            print(f"{s['case']:28s} {'':>6s} | ERROR {s['error'][:70]}")
            continue
        cells = []
        for r in regs:
            cells.append(" ".join(
                f"{s['combos'].get(f'{r}/{k}', {}).get('pct_of_source_tv', '-'):>4}"
                for k in reg_keys(s, r)))
        print(f"{s['case']:28s} {s['volume_check']['source_tv_mm3']:6.0f} | "
              + " | ".join(f"{c:48s}" for c in cells))
    print("=" * width)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--subject", default="2BBN5XA4",
                    help="case (ignored when --n-cases is set)")
    ap.add_argument("--date", default="20250411",
                    help="case (ignored when --n-cases is set)")
    ap.add_argument("--n-cases", type=int, default=0,
                    help="QA the first (or --sample random) N labeled M6 cases "
                         "instead of --subject/--date")
    ap.add_argument("--sample", choices=["first", "random"], default="first",
                    help="batch selection: first N by discover order, or a "
                         "seeded random sample")
    ap.add_argument("--seed", type=int, default=0, help="seed for --sample random")
    ap.add_argument("--regs", nargs="*", default=REGS,
                    help="subset of: " + " ".join(REGS))
    ap.add_argument("--tv", default=None,
                    help="explicit TV file (single-case mode only)")
    ap.add_argument("--out", default=None,
                    help="single output dir (single-case mode only)")
    ap.add_argument("--no-cache", action="store_true",
                    help="re-register and re-warp even if cached outputs exist")
    args = ap.parse_args()
    if args.n_cases and args.tv:
        raise SystemExit("--tv is only valid for a single case")

    out_base = os.path.join(ROOT, "results", "label_interp_qa")
    rows = select_rows(args)
    results = []
    for i, row in enumerate(rows):
        print(f"\n[{i + 1}/{len(rows)}] ===== case {row['key']} =====", flush=True)
        try:
            ct = read_nii_img(row["ct"])
            if not ct_complete(ct):
                print(f"skip {row['key']}: CT head incomplete (data issue)", flush=True)
                results.append({"case": row["key"], "error": "ct_incomplete"})
                continue
            results.append(run_case(row, args, out_base))
        except Exception as e:
            print(f"case {row['key']}: ERROR {e!r}", flush=True)
            results.append({"case": row["key"], "error": repr(e)})

    if args.n_cases:
        ok = [s for s in results if not s.get("error")]
        print_batch_table(results, args.regs)
        bs = os.path.join(out_base, "batch_summary.json")
        with open(bs, "w") as f:
            json.dump({"n_cases": len(rows),
                       "selection": f"{args.sample} seed={args.seed}",
                       "regs": args.regs,
                       "cases": results}, f, indent=1)
        print(f"batch summary: {bs} ({len(ok)}/{len(rows)} cases ok)")
    print("done", flush=True)


if __name__ == "__main__":
    main()