"""Held-out test evaluation of a MONAI round model (torchrun; rank = GPU).

torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py \
    --rows data/manifests/split_test.jsonl --ckpt runs_monai/round0/best.pt \
    --out results/round0_test_monai.json [--no-tta]

Per case: MONAI sliding-window (gaussian blend, 50% overlap) + 4-flip-TTA tumor
probability map; Dice at threshold 0.5 vs the held-out label — the same
convention as Pipeline A's 08_eval and Pipeline B's primary metric.

Per-rank results are appended to <out>_part{rank}.jsonl for resume; rank 0
merges them into <out> (json) + <out>_per_row.jsonl.

Usage:
    python scripts_monai/04_monai_eval_test.py <same flags, single GPU>
    torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py ...
"""
import argparse
import os
import sys
import json
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__))))

import torch
from src.common import load_jsonl, save_jsonl, read_nii_arr
from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
                          predict_probs, prob_dice)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rows", required=True)
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--no-tta", action="store_true")
    ap.add_argument("--win", type=int, default=96)
    ap.add_argument("--overlap", type=float, default=0.5)
    ap.add_argument("--sw-batch", type=int, default=8)
    args = ap.parse_args()

    rank, world, local_rank = rank_info()
    init_dist()
    torch.cuda.set_device(local_rank)
    device = f"cuda:{local_rank}"

    rows = load_jsonl(args.rows)
    outp = os.path.join(os.path.dirname(os.path.abspath(args.out)),
                        os.path.basename(args.out) + f"_part{rank}.jsonl")
    done = set()
    if os.path.exists(outp):
        done = {r["key"] for r in load_jsonl(outp)}
    shard = [r for r in rows if r["key"] not in done][rank::world]
    if rank == 0:
        print(f"[monai:eval:rank0] n_test={len(rows)} done={len(done)} "
              f"shard(world={world})={len(shard)}", flush=True)

    model, _ = load_model(args.ckpt, device)
    with open(outp, "a") as f:
        for i, r in enumerate(shard, 1):
            key = r["key"]
            entry = {"key": key, "dice": None}
            try:
                vol = read_nii_arr(r["pimg"]).astype("float32")
                lab = read_nii_arr(r.get("plabel") or r.get("label")).astype("uint8")
                p = predict_probs(model, vol, device, args.win, args.overlap,
                                  tta=not args.no_tta, sw_batch=args.sw_batch)
                if p.shape == lab.shape:
                    entry["dice"] = round(prob_dice(p, lab), 4)
                else:
                    raise ValueError(f"prob {p.shape} vs label {lab.shape}")
            except Exception as e:  # noqa
                print(f"[monai:eval:rank{rank}] {key} ERR {e!r}", flush=True)
            f.write(json.dumps(entry) + "\n")
            f.flush()
            if i % 20 == 0:
                print(f"[monai:eval:rank{rank}] {i}/{len(shard)}", flush=True)
    barrier()
    if rank != 0:
        destroy_dist()
        return

    per = []
    for k in range(world):
        pf = os.path.join(os.path.dirname(os.path.abspath(args.out)),
                          os.path.basename(args.out) + f"_part{k}.jsonl")
        if os.path.exists(pf):
            per.extend(load_jsonl(pf))
    dice_ok = [e["dice"] for e in per if e["dice"] is not None]
    res = {"ckpt": args.ckpt, "n": len(per), "n_ok": len(dice_ok),
           "dice": float(sum(dice_ok) / len(dice_ok)) if dice_ok else 0.0,
           "tta": not args.no_tta}
    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", "n_ok", "dice")}, indent=1), flush=True)
    print("saved", args.out, flush=True)
    destroy_dist()


if __name__ == "__main__":
    main()