"""Benchmark Pipelines A (in-house Unet3D), B (nnU-Net v2), C (MONAI UNet) on one GPU:
DICE, training time, inference time, model size (params + VRAM).

Protocol
--------
- Data: identical patient-level splits (data/manifests/split_*.jsonl, from
  scripts/05_build_splits.py). A and C use the processed 1mm head-cropped
  percentile-normalized volumes (pimg/plabel); B uses the native source
  volumes (img/label, joined by key from ntuh.jsonl / m6_labeled.jsonl) —
  each pipeline's own standard data protocol.
- Training: --train-epochs on the first --max-rows train rows (single GPU,
  bf16 fp-mixing as in the study). "1 epoch" = one pass over the training
  set under each pipeline's native settings: A/C batch 3, 96^3 random
  patches, AdamW 3e-4 + warmup/cosine; B nnU-Net planned batch (2), planned
  patch size, AdamW 1e-2 PolyLR. B builds an ISOLATED dataset (--nnu-dsid,
  default 221; the study's Dataset210 is never touched); its one-time
  raw-build + plan&preprocess wall time is reported separately.
- Inference: each pipeline's standard protocol on the held-out test rows:
  A sliding window 96^3 step 48 + 4-flip TTA; B nnUNetPredictor (gaussian
  blend, nnU-Net mirroring); C MONAI sliding window (gaussian, 50% overlap)
  + 4-flip TTA. Mean s/volume over --max-eval-rows (first
  --eval-warmup cases untimed) + peak VRAM of the inference run.
- Dice: tumor probability thresholded at 0.5 vs ground truth — the primary
  metric convention of all three study pipelines (B also reports the
  argmax hard-seg dice it exports natively).
- VRAM training: max allocation over the training run (A/C in-process via
  torch; B via nvidia-smi polling around the nnUNetv2_train subprocess).

Usage (repo root, inside the `longitudinal` conda env):
  python scripts/benchmark_pipelines.py                          # train + eval A B C
  python scripts/benchmark_pipelines.py --pipelines A C
  python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
  python scripts/benchmark_pipelines.py --no-tta                 # no test-time augmentation
  python scripts/benchmark_pipelines.py --skip-train             # eval existing study checkpoints
Outputs: results/benchmark_pipelines.json (+ _per_row.jsonl, _bench.png)
"""
import argparse
import gc
import json
import os
import shutil
import subprocess
import sys
import threading
import time
from datetime import datetime

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

import numpy as np
import torch

from src.common import load_jsonl, read_nii_arr
from src.training import dice_np


def d(name):
    p = os.path.join(ROOT, name)
    os.makedirs(p, exist_ok=True)
    return p


def parse_args():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pipelines", default="A B C", help="space-separated subset of 'A B C'")
    ap.add_argument("--gpu", type=int, default=0, help="physical GPU index (all pipelines use one GPU)")
    ap.add_argument("--train-epochs", type=int, default=100, help="training budget for A and C (and B via NNU_PL_EPOCHS)")
    ap.add_argument("--max-rows", type=int, default=None,
                    help="train rows for A/C (first N of labeled split_train; default 80%% of the split)")
    ap.add_argument("--nnu-cases", type=int, default=None,
                    help="train cases for B (first N labeled native cases; default same as --max-rows)")
    ap.add_argument("--max-eval-rows", type=int, default=None,
                    help="held-out test volumes for inference + DICE (default 10%% of labeled split_test)")
    ap.add_argument("--eval-warmup", type=int, default=1, help="untimed warmup cases before inference timing")
    ap.add_argument("--workers", type=int, default=4, help="dataloader workers (A/C)")
    ap.add_argument("--batch", type=int, default=3, help="per-GPU batch (A/C)")
    ap.add_argument("--patch", type=int, default=96, help="patch/window (A/C)")
    ap.add_argument("--no-tta", action="store_true", help="disable each pipeline's standard TTA")
    ap.add_argument("--skip-train", action="store_true",
                    help="eval existing study checkpoints (runs/round0, runs_monai/round0, nnu/results 210) "
                         "instead of training")
    ap.add_argument("--nnu-dsid", type=int, default=None,
                help="nnU-Net dataset id: training mode default 221 (isolated; the study's 210 is never "
                     "touched), --skip-train mode default 210 (study model); pass explicitly to evaluate "
                     "a previously benchmarked model folder")
    ap.add_argument("--nnu-skip-prep", action="store_true", help="B/train: reuse an existing raw dataset + "
                                                                 "preprocessed tree for --nnu-dsid")
    ap.add_argument("--out", default=None, help="output json path (default results/benchmark_pipelines.json)")
    return ap.parse_args()


# ---------------- data selection ----------------

def native_manifests():
    """key -> {img, label} from the labeled source manifests (ntuh + m6_labeled)."""
    man = os.path.join(ROOT, "data", "manifests")
    out = {}
    for name in ("ntuh.jsonl", "m6_labeled.jsonl"):
        p = os.path.join(man, name)
        if not os.path.exists(p):
            continue
        for r in load_jsonl(p):
            if r.get("label"):
                out[r["key"]] = {"img": r["img"], "label": r["label"]}
    return out


def native_rows_for(keys):
    nm = native_manifests()
    rows = []
    for k in keys:
        if k not in nm:
            raise SystemExit(f"B: key {k[:60]}... not in ntuh/m6_labeled manifests; "
                             f"native img/label unavailable")
        rows.append({"key": k, **nm[k]})
    return rows


def select_rows(args):
    man = os.path.join(ROOT, "data", "manifests")
    train = [r for r in load_jsonl(os.path.join(man, "split_train.jsonl")) if r.get("plabel")]
    test = [r for r in load_jsonl(os.path.join(man, "split_test.jsonl")) if r.get("plabel")]
    if not train:
        raise SystemExit("split_train.jsonl has no labeled processed rows; run scripts/preprocess.py + 05_build_splits.py")
    if not test:
        raise SystemExit("split_test.jsonl has no labeled processed rows; run scripts/preprocess.py + 05_build_splits.py")
    if args.max_rows is None:
        args.max_rows = max(1, round(0.8 * len(train)))
    if args.max_eval_rows is None:
        args.max_eval_rows = max(1, round(0.1 * len(test)))
    if args.nnu_cases is None:
        args.nnu_cases = args.max_rows
    return train[:args.max_rows], test[:args.max_eval_rows]


# ---------------- GPU / timing helpers ----------------

def reset_peak(dev):
    torch.cuda.reset_peak_memory_stats(dev)


def peak_mb(dev):
    return round(torch.cuda.max_memory_allocated(dev) / 1e6, 1)


class SmiPoller(threading.Thread):
    """Samples nvidia-smi memory.used (MiB) for one physical GPU in the background."""

    def __init__(self, gpu):
        super().__init__(daemon=True)
        self.gpu, self._stop = gpu, False
        self.max_mib = 0.0

    def run(self):
        while not self._stop:
            try:
                out = subprocess.run(["nvidia-smi", f"--id={self.gpu}",
                                      "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
                                     capture_output=True, text=True, timeout=5).stdout.strip()
                self.max_mib = max(self.max_mib, float(out.splitlines()[0]))
            except Exception:
                pass
            time.sleep(0.4)

    def stop(self):
        self._stop = True
        self.join(timeout=2)


def free_gpu():
    gc.collect()
    torch.cuda.empty_cache()


def load_native_label(lab_path, img_itk):
    lab = read_nii_arr(lab_path).astype("uint8")
    import SimpleITK as sitk
    if lab.shape != tuple(img_itk.GetSize()[::-1]):
        lab = sitk.GetArrayFromImage(sitk.Resample(sitk.ReadImage(lab_path), img_itk,
                                                   sitk.Transform(), sitk.sitkNearestNeighbor, 0.0))
        lab = lab.astype("uint8")
    return lab


# ---------------- Pipeline A ----------------

def bench_A(args, train_rows, test_rows):
    from src.dataset import make_dataloader
    from src.losses import per_sample_loss, one_hot
    from src import training as A

    dev = torch.device(f"cuda:{args.gpu}")
    torch.cuda.set_device(dev)
    tta = not args.no_tta
    res = {"pipeline": "A", "backbone": "Unet3D(base=16, depth=4)", "tta": tta,
           "n_train": len(train_rows), "n_eval": len(test_rows)}

    ckpt_dir = d("runs_bench/A")
    if args.skip_train:
        ckpt = os.path.join(ROOT, "runs", "round0", "best.pt")
        if not os.path.exists(ckpt):
            raise FileNotFoundError(f"--skip-train: {ckpt} not found")
        model, _ = A.load_model(ckpt, dev)
        res["ckpt"] = ckpt
    else:
        model = A.build_model(base=16, device="cuda")
        patch = (args.patch, args.patch, args.patch)
        _, dl = make_dataloader(train_rows, patch, args.batch, True, num_workers=args.workers,
                                seed=0, persistent=False)
        steps = max(len(dl), 1)
        total_steps = steps * args.train_epochs
        opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
        warmup = min(300, max(10, total_steps // 10))
        sched = torch.optim.lr_scheduler.SequentialLR(
            opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup),
                  torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1),
                                                              eta_min=3e-4 * 0.05)],
            milestones=[warmup])
        model.train()
        reset_peak(dev)
        t0 = time.time()
        for ep in range(args.train_epochs):
            if hasattr(dl.sampler, "set_epoch"):
                dl.sampler.set_epoch(ep)
            loss_sum = n = 0.0
            for img, lab, wts in dl:
                img = img.to(dev, non_blocking=True)
                lab = lab.to(dev, non_blocking=True)
                wts = wts.to(dev, non_blocking=True)
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    logits = model(img)
                losses = per_sample_loss(logits.float(), one_hot(lab, 2))
                if torch.isnan(losses).any() or torch.isinf(losses).any():
                    losses = torch.zeros_like(losses)
                loss = (losses * wts).sum() / wts.sum().clamp(min=1e-6)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                opt.step()
                sched.step()
                loss_sum += float(losses.detach().mean())
                n += 1
            res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
        dt = time.time() - t0
        res["train_s"] = round(dt, 1)
        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
        res["train_vram_mb"] = peak_mb(dev)
        res["steps_per_epoch"] = steps
        ckpt = os.path.join(ckpt_dir, "best.pt")
        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
        del dl, opt, sched
    res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
    res["params"] = int(sum(p.numel() for p in model.parameters()))
    res["params_M"] = round(res["params"] / 1e6, 3)
    res["weight_mb"] = round(os.path.getsize(res["ckpt"]) / 1e6, 3)

    # inference: sliding window, 4-flip TTA (study protocol of 06/08)
    model.eval()
    step = max(args.patch // 2, 16)
    warm = test_rows[0]
    with torch.no_grad():
        A.sliding_window_probs(model, read_nii_arr(warm["pimg"]).astype("float32"), dev,
                               args.patch, step, tta)
    torch.cuda.synchronize(dev)
    reset_peak(dev)
    times, dices = [], []
    with torch.no_grad():
        for r in test_rows[args.eval_warmup:]:
            vol = read_nii_arr(r["pimg"]).astype("float32")
            torch.cuda.synchronize(dev)
            t0 = time.time()
            probs = A.sliding_window_probs(model, vol, dev, args.patch, step, tta)
            torch.cuda.synchronize(dev)
            times.append(time.time() - t0)
            # probs are on the padded grid; crop back to the volume/label shape
            osh = vol.shape
            probs = probs[:osh[0], :osh[1], :osh[2]]
            lab = read_nii_arr(r["plabel"]).astype("uint8")
            if lab.shape != osh:
                lab = np.pad(lab, [(0, max(osh[i] - lab.shape[i], 0)) for i in range(3)])[:osh]
            dices.append(dice_np(probs, lab))
    res["infer_s_per_vol"] = round(float(np.mean(times)), 2) if times else None
    res["infer_s_per_vol_std"] = round(float(np.std(times)), 2) if times else None
    res["infer_vram_mb"] = peak_mb(dev)
    res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0
    res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"],
                            "dice": round(dices[i], 4), "s": round(times[i], 2)}
                           for i in range(len(dices))]
    del model
    free_gpu()
    return res


# ---------------- Pipeline C (MONAI) ----------------

def bench_C(args, train_rows, test_rows):
    sys.path.insert(0, os.path.join(ROOT, "scripts_monai"))
    from monai_common import (build_model, WeightedDiceCELoss, make_dataloader,
                              predict_probs, prob_dice, load_model)

    dev = torch.device(f"cuda:{args.gpu}")
    torch.cuda.set_device(dev)
    tta = not args.no_tta
    res = {"pipeline": "C", "backbone": "monai UNet(16->128)", "tta": tta,
           "n_train": len(train_rows), "n_eval": len(test_rows)}

    ckpt_dir = d("runs_bench/C")
    if args.skip_train:
        ckpt = os.path.join(ROOT, "runs_monai", "round0", "best.pt")
        if not os.path.exists(ckpt):
            raise FileNotFoundError(f"--skip-train: {ckpt} not found")
        model, _ = load_model(ckpt, dev)
        res["ckpt"] = ckpt
    else:
        model = build_model(dev)
        dl = make_dataloader(train_rows, win=args.patch, batch=args.batch, workers=args.workers)
        steps = max(len(dl), 1)
        total_steps = steps * args.train_epochs
        opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
        warmup = min(300, max(10, total_steps // 10))
        sched = torch.optim.lr_scheduler.SequentialLR(
            opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup),
                  torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1),
                                                              eta_min=3e-4 * 0.05)],
            milestones=[warmup])
        loss_f = WeightedDiceCELoss()
        model.train()
        reset_peak(dev)
        t0 = time.time()
        for ep in range(args.train_epochs):
            if hasattr(dl.sampler, "set_epoch"):
                dl.sampler.set_epoch(ep)
            loss_sum = n = 0.0
            for batch in dl:
                img = batch["pimg"].to(dev, non_blocking=True)
                lab = batch["plabel"].squeeze(1).to(dev, non_blocking=True)
                wts = batch["w"].to(dev, non_blocking=True)
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    logits = model(img)
                loss = loss_f(logits, lab, wts)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                opt.step()
                sched.step()
                loss_sum += float(loss.detach())
                n += 1
            res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
        dt = time.time() - t0
        res["train_s"] = round(dt, 1)
        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
        res["train_vram_mb"] = peak_mb(dev)
        res["steps_per_epoch"] = steps
        ckpt = os.path.join(ckpt_dir, "best.pt")
        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
        del dl, opt, sched, loss_f
    res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
    res["params"] = int(sum(p.numel() for p in model.parameters()))
    res["params_M"] = round(res["params"] / 1e6, 3)
    res["weight_mb"] = round(os.path.getsize(res["ckpt"]) / 1e6, 3)

    model.eval()
    warm = test_rows[0]
    predict_probs(model, read_nii_arr(warm["pimg"]).astype("float32"), dev, args.patch, 0.5,
                  tta=tta, sw_batch=8)
    torch.cuda.synchronize(dev)
    reset_peak(dev)
    times, dices = [], []
    with torch.no_grad():
        for r in test_rows[args.eval_warmup:]:
            vol = read_nii_arr(r["pimg"]).astype("float32")
            torch.cuda.synchronize(dev)
            t0 = time.time()
            p = predict_probs(model, vol, dev, args.patch, 0.5, tta=tta, sw_batch=8)
            torch.cuda.synchronize(dev)
            times.append(time.time() - t0)
            lab = read_nii_arr(r["plabel"]).astype("uint8")
            if p.shape != lab.shape:
                lab = np.pad(lab, [(0, max(p.shape[i] - lab.shape[i], 0)) for i in range(3)])[:p.shape]
            dices.append(prob_dice(p, lab))
    res["infer_s_per_vol"] = round(float(np.mean(times)), 2) if times else None
    res["infer_s_per_vol_std"] = round(float(np.std(times)), 2) if times else None
    res["infer_vram_mb"] = peak_mb(dev)
    res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0
    res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"],
                            "dice": round(dices[i], 4), "s": round(times[i], 2)}
                           for i in range(len(dices))]
    del model
    free_gpu()
    return res


# ---------------- Pipeline B (nnU-Net) ----------------

def bench_B(args, train_rows, test_rows):
    import nnu_common as N
    import SimpleITK as sitk

    if args.skip_train:
        N.DS_ID = args.nnu_dsid or 210  # study model by default
    else:
        N.DS_ID = args.nnu_dsid or 221  # isolated dataset; the study's 210 is never touched
    N.DS_NAME = "Dataset210_NTUH_T1C_PL" if (N.DS_ID == 210 and args.skip_train) \
        else f"Dataset{N.DS_ID:03d}_T1C_BENCH"

    dev = torch.device(f"cuda:{args.gpu}")
    torch.cuda.set_device(dev)
    tta = not args.no_tta
    res = {"pipeline": "B", "backbone": "nnU-Net 3d_fullres", "tta": tta,
           "nnu_dataset": N.DS_NAME, "n_train": None, "n_eval": len(test_rows)}

    te = test_rows
    if not args.skip_train:
        tr = train_rows[:args.nnu_cases]
        nat = native_rows_for([r["key"] for r in tr])
        res["n_train"] = len(nat)

        t0 = time.time()
        N.make_raw_dataset(nat)
        res["raw_build_s"] = round(time.time() - t0, 1)
        if not args.nnu_skip_prep:
            t0 = time.time()
            N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env())
            res["plan_preprocess_s"] = round(time.time() - t0, 1)
        else:
            res["plan_preprocess_s"] = 0.0
        all_keys = N.dataset_case_keys()
        val_keys = all_keys[-max(1, len(all_keys) // 10):]
        N.write_splits(all_keys, val_keys)

        env = N.nnu_env(epoch=args.train_epochs, lr=1e-2)
        env["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
        poller = SmiPoller(args.gpu)
        poller.start()
        t0 = time.time()
        N.run(N.train_cmd(1), os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
        res["train_s"] = round(time.time() - t0, 1)
        res["train_s_per_epoch"] = round(res["train_s"] / max(args.train_epochs, 1), 1)
        res["train_vram_mb"] = round(poller.max_mib * 1.048576, 1) if poller.max_mib else None
        poller.stop()
        res["train_vram_note"] = "nvidia-smi GPU memory.used peak during nnUNetv2_train sub-process"
    else:
        if not os.path.exists(N.best_ckpt()):
            raise FileNotFoundError(
                f"--skip-train: {N.best_ckpt()} not found (study nnU-Net results missing)")

    ckpt = N.best_ckpt()
    res["ckpt"] = ckpt
    res["weight_mb"] = round(os.path.getsize(ckpt) / 1e6, 1) if os.path.exists(ckpt) else None

    # in-process inference via nnUNetPredictor
    os.environ.update(N.nnu_env())  # nnUNet_* roots + nnUNet_extTrainer (NTUHLPLTrainer discovery)
    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
    predictor = nnUNetPredictor(tile_step_size=0.5, use_gaussian=True,
                                use_mirroring=tta, perform_everything_on_device=True,
                                device=dev, verbose=False, allow_tqdm=False)
    predictor.initialize_from_trained_model_folder(N.model_folder(), use_folds=[0],
                                                   checkpoint_name="checkpoint_best.pth")
    net = predictor.network
    if hasattr(net, "module"):
        net = net.module
    res["params"] = int(sum(p.numel() for p in net.parameters()))
    res["params_M"] = round(res["params"] / 1e6, 3)

    te_keys = [r["key"] for r in te]
    te_native = native_rows_for(te_keys)
    in_dir = d("runs_bench/B/pred_in")
    out_dir = d("runs_bench/B/pred_out")
    shutil.rmtree(in_dir, ignore_errors=True)
    shutil.rmtree(out_dir, ignore_errors=True)
    os.makedirs(in_dir)
    os.makedirs(out_dir)
    for r in te_native:
        os.symlink(os.path.abspath(r["img"]), os.path.join(in_dir, r["key"] + "_0000.nii.gz"))

    def run_cases(cs):
        # list-of-lists: one channel-file list per case
        imgs = [[os.path.join(in_dir, c["key"] + "_0000.nii.gz")] for c in cs]
        predictor.predict_from_files(imgs, out_dir, save_probabilities=True, overwrite=True,
                                     num_processes_preprocessing=2,
                                     num_processes_segmentation_export=2,
                                     num_parts=1, part_id=0)

    warm, timed = te_native[:args.eval_warmup], te_native[args.eval_warmup:]
    if warm:
        run_cases(warm)
    torch.cuda.synchronize(dev)
    reset_peak(dev)
    t0 = time.time()
    if timed:
        run_cases(timed)
    res["infer_s_per_vol"] = round((time.time() - t0) / len(timed), 2) if timed else None
    res["infer_vram_mb"] = peak_mb(dev) if timed else None

    dices, dice_hard, per = [], [], []
    for r in timed:
        key = r["key"]
        npz_f = os.path.join(out_dir, key + ".npz")
        seg_f = os.path.join(out_dir, key + ".nii.gz")
        if not os.path.exists(npz_f):
            per.append({"key": key, "dice": None, "note": "no npz"})
            continue
        probs = np.load(npz_f, allow_pickle=False)["probabilities"]
        if probs.shape[0] != 2:
            per.append({"key": key, "dice": None, "note": f"bad channels {probs.shape[0]}"})
            continue
        pt = probs[1]
        img_itk = sitk.ReadImage(r["img"])
        lab = load_native_label(r["label"], img_itk)
        if pt.shape != lab.shape:
            per.append({"key": key, "dice": None,
                        "note": f"prob {pt.shape} vs label {lab.shape}"})
            continue
        dices.append(dice_np(pt, lab))
        if os.path.exists(seg_f):
            hard = read_nii_arr(seg_f)
            if hard.shape == lab.shape:
                dice_hard.append(dice_np((hard > 0).astype("float32"), lab))
        per.append({"key": key, "dice": round(dices[-1], 4)})
    res["dice_per_row"] = per
    res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0
    if dice_hard:
        res["dice_hard"] = round(float(np.mean(dice_hard)), 4)

    shutil.rmtree(in_dir, ignore_errors=True)
    del predictor
    free_gpu()
    return res


# ---------------- report ----------------

def print_table(results):
    print()
    hdr = (f"{'pipeline':46s} {'dice':>6s} {'paramsM':>8s} {'wtMB':>7s} "
           f"{'s/epoch':>8s} {'trVRAM':>8s} {'s/vol':>7s} {'inVRAM':>8s}")
    print(hdr)
    print("-" * len(hdr))
    for r in results:
        if "error" in r:
            print(f"{r['pipeline']:46s} ERROR {r['error']}")
            continue
        extra = []
        if r["pipeline"] == "B":
            extra.append(f"prep={r.get('plan_preprocess_s')}s")
        line = (f"{r['pipeline']+' '+r['backbone'][:40]:46s} "
                f"{r['dice']:6.4f} {r['params_M']:8.2f} {r.get('weight_mb', 0):7.2f} "
                f"{r.get('train_s_per_epoch', float('nan')):8.1f} "
                f"{(r.get('train_vram_mb') or float('nan')):8.0f} "
                f"{(r.get('infer_s_per_vol') or float('nan')):7.2f} "
                f"{(r.get('infer_vram_mb') or float('nan')):8.0f}")
        if extra:
            line += "  " + " ".join(extra)
        print(line)
    print()


def plot(results, out_png, epochs):
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        ok = [r for r in results if "error" not in r]
        if not ok:
            return
        names = [f"{r['pipeline']} {r['backbone']}" for r in ok]
        fig, axes = plt.subplots(2, 3, figsize=(15, 8))
        panels = [
            ("DICE (test", "dice", "{:.4f}"),
            ("Train s/epoch", "train_s_per_epoch", "  {:.0f}"),
            ("Infer s/vol (TTA as marked)", "infer_s_per_vol", "  {:.1f}"),
            ("Params (M)", "params_M", "  {:.1f}"),
            ("VRAM train (MB)", "train_vram_mb", "  {:.0f}"),
            ("VRAM infer (MB)", "infer_vram_mb", "  {:.0f}"),
        ]
        for ax, (title, key, fmt) in zip(axes.ravel(), panels):
            vals = [r.get(key) for r in ok]
            vals = [v if v is not None else 0.0 for v in vals]
            bars = ax.bar(range(len(ok)), vals, color=["#4878CF", "#6ACC65", "#D64F5C"][:len(ok)])
            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=9)
            ax.set_title(title + ")", fontsize=10)
            ax.set_xticks(range(len(ok)))
            ax.set_xticklabels([n.split(" ")[0] for n in names], fontsize=10)
        plt.suptitle(f"Pipeline benchmark: {epochs} train epochs, single GPU", fontsize=11)
        fig.tight_layout()
        fig.savefig(out_png, dpi=130)
        plt.close(fig)
        print("plot:", out_png)
    except Exception as e:  # noqa
        print("plot failed:", repr(e))


def main():
    args = parse_args()
    pipelines = [p.upper() for p in args.pipelines.split()]
    unknown = set(pipelines) - set("ABC")
    if unknown:
        raise SystemExit(f"unknown pipeline(s): {unknown}")

    out_f = args.out or os.path.join(ROOT, "results", "benchmark_pipelines.json")
    os.makedirs(os.path.dirname(out_f), exist_ok=True)
    train_rows, test_rows = select_rows(args)

    print(f"[bench] gpu={args.gpu} train_rows={len(train_rows)} eval_rows={len(test_rows)} "
          f"epochs={args.train_epochs} tta={not args.no_tta} skip_train={args.skip_train}")

    runners = {"A": bench_A, "B": bench_B, "C": bench_C}
    results = []
    for p in pipelines:
        t0 = time.time()
        print(f"\n[bench] ---- pipeline {p} ----")
        try:
            r = runners[p](args, train_rows, test_rows)
            r["wall_s"] = round(time.time() - t0, 1)
        except Exception as e:  # noqa
            r = {"pipeline": p, "error": repr(e), "wall_s": round(time.time() - t0, 1)}
            print(f"[bench] pipeline {p} FAILED: {e!r}")
        results.append(r)
        free_gpu()

    payload = {
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "settings": {k: v for k, v in vars(args).items()},
        "protocol": ("single GPU; A/C bf16 batch 3 patch 96 AdamW 3e-4; B nnU-Net planned "
                     "batch/patch fp16 AdamW 1e-2; dice = prob@0.5 vs GT; inference s/vol = "
                     "mean over timed test cases after warmup; train VRAM A/C = torch peak, "
                     "B = nvidia-smi peak; B dataset id isolated at "
                     f"{args.nnu_dsid} in train mode"),
        "results": results,
    }
    with open(out_f, "w") as f:
        json.dump(payload, f, indent=1)
    per_rows = [x for r in results for x in r.get("dice_per_row", [])]
    if per_rows:
        with open(out_f.replace(".json", "_per_row.jsonl"), "w") as f:
            for x in per_rows:
                f.write(json.dumps(x) + "\n")
    print_table(results)
    print("json:", out_f)
    plot(results, out_f.replace(".json", "_bench.png"), args.train_epochs)


if __name__ == "__main__":
    main()