"""Holdout test evaluation of an nnU-Net round model.

Predicts the test volumes (multi-GPU sharding, --save_probabilities) and scores:
  dice       - tumor probability (softmax channel 1) thresholded at 0.5
               (same convention as scripts/08_eval.py)
  dice_hard  - argmax segmentation written by nnUNetv2_predict
Writes: <out> json + <out>.json -> per_row jsonl alongside.

Usage: python scripts_nnu/05_nnu_eval_test.py --rows data/manifests/split_test.jsonl \
    --out results/round0_test_nnu.json --gpus 3
"""
import argparse
import os
import sys
import json
import shutil
import subprocess
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.common import load_jsonl, save_jsonl, read_nii_arr, dice
from nnu_common import ROOT, d, nnu_env, predict_cmd, tumor_probs_from_npz


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rows", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--gpus", type=int, default=3)
    ap.add_argument("--chk", default="checkpoint_best.pth")
    ap.add_argument("--no-tta", action="store_true")
    ap.add_argument("--npp", type=int, default=2)
    args = ap.parse_args()

    rows = load_jsonl(args.rows)
    tag = os.path.basename(os.path.normpath(args.out)) + "_predict"
    in_dir = os.path.join(d("nnu/pred"), f"in_test_{tag}")
    out_dir = os.path.join(d("nnu/pred"), f"out_test_{tag}")
    shutil.rmtree(in_dir, ignore_errors=True)
    os.makedirs(in_dir, exist_ok=True)
    for r in rows:
        os.symlink(os.path.abspath(r["pimg"]), os.path.join(in_dir, r["key"] + "_0000.nii.gz"))

    env = nnu_env()
    procs = []
    for g in range(args.gpus):
        e = dict(env)
        e["CUDA_VISIBLE_DEVICES"] = str(g)
        logf = os.path.join(d("logs/nnu"), f"pred_test_{tag}_part{g}.log")
        with open(logf, "wb") as lf:
            procs.append(subprocess.Popen(predict_cmd(in_dir, os.path.join(out_dir, f"part{g}"),
                                                      args.gpus, g, chk=args.chk,
                                                      tta=not args.no_tta, npp=args.npp),
                                          stdout=lf, stderr=subprocess.STDOUT, env=e, cwd=ROOT))
    for g, p in enumerate(procs):
        rc = p.wait()
        if rc != 0:
            print(f"[nnu:eval] test predict part{g} failed rc={rc}; retrying once", flush=True)
            e = dict(env)
            e["CUDA_VISIBLE_DEVICES"] = str(g)
            logf = os.path.join(d("logs/nnu"), f"pred_test_{tag}_part{g}_retry.log")
            with open(logf, "wb") as lf:
                p2 = subprocess.Popen(predict_cmd(in_dir, os.path.join(out_dir, f"part{g}"),
                                                  args.gpus, g, chk=args.chk,
                                                  tta=not args.no_tta, npp=args.npp),
                                      stdout=lf, stderr=subprocess.STDOUT, env=e, cwd=ROOT)
                if p2.wait() != 0:
                    raise RuntimeError(f"test prediction part{g} failed twice; see {logf}")

    npz, segs = {}, {}
    for part in os.listdir(out_dir):
        pd = os.path.join(out_dir, part)
        if not os.path.isdir(pd):
            continue
        for fn in os.listdir(pd):
            if fn.endswith(".npz"):
                npz[fn[:-4]] = os.path.join(pd, fn)
            elif fn.endswith(".nii.gz"):
                segs[fn[:-7]] = os.path.join(pd, fn)

    per, dice_p, dice_h = [], [], []
    for r in rows:
        key = r["key"]
        lab = read_nii_arr(r.get("plabel") or r.get("label"))
        try:
            pt = tumor_probs_from_npz(npz[key])
            if pt.shape != lab.shape:
                raise ValueError(f"prob {pt.shape} vs label {lab.shape}")
            d1 = dice((pt >= 0.5).astype("uint8"), (lab > 0).astype("uint8"))
            dice_p.append(d1)
            hard = read_nii_arr(segs[key])
            if hard.shape != lab.shape:
                raise ValueError(f"hard seg {hard.shape} vs label {lab.shape}")
            d2 = dice((hard > 0).astype("uint8"), (lab > 0).astype("uint8"))
            dice_h.append(d2)
            per.append({"key": key, "dice": round(d1, 4), "dice_hard": round(d2, 4)})
        except Exception as e:  # noqa
            print(f"[nnu:eval] {key} ERR {e!r}", flush=True)

    res = {"ckpt_chron": args.chk, "n": len(per),
           "dice": float(sum(dice_p) / len(dice_p)) if dice_p else 0.0,
           "dice_hard": float(sum(dice_h) / len(dice_h)) if dice_h else 0.0,
           "n_pred": len(npz), "per_row": per}
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(res, f, indent=1)
    save_jsonl(per, args.out.replace(".json", "_per_row.jsonl"))
    print(json.dumps({k: res[k] for k in ("n", "dice", "dice_hard", "n_pred")}, indent=1), flush=True)
    print("saved", args.out, flush=True)
    shutil.rmtree(in_dir, ignore_errors=True)


if __name__ == "__main__":
    main()