"""Pseudo-label the unlabeled pool with the MONAI round model (torchrun; rank = GPU).

torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py \
    --ckpt runs_monai/round0/best.pt --pool data/manifests/unlabeled_pool.jsonl \
    --out data/pseudo_monai/round1 [--already data/pseudo_monai/added_keys.jsonl]

Each rank (GPU) takes rows[rank::world] and computes MONAI sliding-window
(gaussian blend, 50% overlap) + 4-flip-TTA tumor probabilities, then applies
the study gates (identical to Pipelines A/B, shared from scripts_nnu/nnu_common):
  pos: p_tumor >= tau_pos, median cleanup, largest-CC fraction >= min_cc_frac,
       volume within the labeled-tumor [p2, p98] range (data/vols.json)
  neg: >= neg_frac of interior (vol > 0.02) voxels have p_bg >= tau_neg
then the per-subject frame-invariant longitudinal consistency filter over
accepted positive timepoints (head-relative centroid shift + volume ratio).

Writes <out>/rows.jsonl, <out>/<key>_label.nii.gz (tumor mask for pos, zero
mask for neg), <out>/accepted.jsonl, <out>/summary.json. Per-rank shards are
flushed to <out>/part{rank}.jsonl so interrupted runs resume.

Usage:
    python scripts_monai/03_monai_pseudo_label.py <same flags, single GPU>
    torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.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 numpy as np
import SimpleITK as sitk
from src.common import load_jsonl, save_jsonl, read_nii_arr, write_arr, d
from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
                          predict_probs, pos_mask, neg_frac_bg, consistency_filter,
                          load_voxel_stats)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--pool", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--already", default=None, help="jsonl of case keys already consumed by the dataset")
    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)
    ap.add_argument("--tau-pos", type=float, default=0.95)
    ap.add_argument("--tau-neg", type=float, default=0.98)
    ap.add_argument("--neg-frac", type=float, default=0.90)
    ap.add_argument("--vol-qp", type=float, nargs=2, default=[2, 98])
    ap.add_argument("--min-cc-frac", type=float, default=0.2)
    ap.add_argument("--max-rel-dist", type=float, default=40.0,
                    help="max mm of head-relative tumor centroid shift between compatible timepoints")
    ap.add_argument("--vol-ratio", type=float, default=10.0,
                    help="max tumor volume ratio between compatible timepoints")
    args = ap.parse_args()

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

    out = d(args.out)
    pool_rows = load_jsonl(args.pool)
    already = set()
    if args.already and os.path.exists(args.already):
        already = {r["key"] for r in load_jsonl(args.already)}
    done = set()
    for k in range(world):
        pf = os.path.join(out, f"part{k}.jsonl")
        if os.path.exists(pf):
            done |= {r["key"] for r in load_jsonl(pf)}
    rows = [r for r in pool_rows if r["key"] not in already and r["key"] not in done]
    shard = rows[rank::world]

    vstats = load_voxel_stats()
    vol_lo = vstats.get(f"p{args.vol_qp[0]:.0f}", 1.0)
    vol_hi = vstats.get(f"p{args.vol_qp[1]:.0f}", 50000.0)
    if rank == 0:
        print(f"[monai:pseudo:rank0] pool={len(pool_rows)} already={len(already)} "
              f"done={len(done)} to_predict={len(rows)} shard={len(shard)} "
              f"vol_range=[{vol_lo:.0f},{vol_hi:.0f}]mm3", flush=True)

    model, _ = load_model(args.ckpt, device)
    part = os.path.join(out, f"part{rank}.jsonl")
    buf, n_err = [], 0
    for i, r in enumerate(shard, 1):
        key = r["key"]
        entry = {"key": key, "subject": r["subject"], "date": r.get("date"),
                 "source": r.get("source"), "pimg": r["pimg"], "label": None,
                 "role": "rej", "vol_mm3": 0, "maxp": None}
        try:
            pimg_itk = sitk.ReadImage(r["pimg"])
            vol = read_nii_arr(r["pimg"]).astype("float32")
            pt = predict_probs(model, vol, device, args.win, args.overlap,
                               tta=not args.no_tta, sw_batch=args.sw_batch)
            if pt.shape != vol.shape:
                raise ValueError(f"prob shape {pt.shape} != image shape {vol.shape} for {key}")
            entry["maxp"] = round(float(pt.max()), 4)
            got = pos_mask(pt, args.tau_pos, args.min_cc_frac, vol_lo, vol_hi)
            if got is not None:
                mask, cc_frac, vol_mm3 = got
                lp = os.path.join(out, key + "_label.nii.gz")
                write_arr(mask, lp, itk_img=pimg_itk)
                entry.update({"role": "pos", "label": lp, "vol_mm3": vol_mm3, "cc_frac": round(cc_frac, 3)})
            else:
                frac = neg_frac_bg(pt, vol, args.tau_neg, args.neg_frac)
                if frac is not None:
                    lp = os.path.join(out, key + "_label.nii.gz")
                    write_arr(np.zeros(vol.shape, dtype="uint8"), lp, itk_img=pimg_itk)
                    entry.update({"role": "neg", "label": lp, "neg_conf": round(frac, 4)})
        except Exception as e:  # noqa
            entry["role"] = "error"
            n_err += 1
            print(f"[monai:pseudo:rank{rank}] {key} ERR {e!r}", flush=True)
        buf.append(entry)
        if i % 20 == 0:
            with open(part, "a") as f:
                for b in buf:
                    f.write(json.dumps(b) + "\n")
            buf = []
            print(f"[monai:pseudo:rank{rank}] {i}/{len(shard)}", flush=True)
    if buf:
        with open(part, "a") as f:
            for b in buf:
                f.write(json.dumps(b) + "\n")
    barrier()
    if rank != 0:
        destroy_dist()
        return

    merged = []
    for k in range(world):
        pf = os.path.join(out, f"part{k}.jsonl")
        if os.path.exists(pf):
            merged.extend(load_jsonl(pf))
    n_pos0 = sum(1 for x in merged if x["role"] == "pos")
    n_rej = consistency_filter(merged, out, args.max_rel_dist, args.vol_ratio)
    accepted = [x for x in merged if x["role"] in ("pos", "neg")]
    save_jsonl(merged, os.path.join(out, "rows.jsonl"))
    save_jsonl(accepted, os.path.join(out, "accepted.jsonl"))
    posv = [x["vol_mm3"] for x in merged if x["role"] == "pos"]
    summ = {
        "n_pool_predicted": len(merged),
        "n_pos": sum(1 for x in merged if x["role"] == "pos"),
        "n_neg": sum(1 for x in merged if x["role"] == "neg"),
        "n_pos_before_consistency": n_pos0,
        "n_rejected_consistency": n_rej,
        "n_other_rej": sum(1 for x in merged if x["role"] == "rej"),
        "n_error": n_err,
        "pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0,
                        "p5": float(np.percentile(posv, 5)) if posv else 0,
                        "p95": float(np.percentile(posv, 95)) if posv else 0},
        "tau_pos": args.tau_pos, "tau_neg": args.tau_neg, "neg_frac": args.neg_frac,
        "vol_range": [vol_lo, vol_hi], "max_rel_dist_mm": args.max_rel_dist,
        "vol_ratio": args.vol_ratio, "win": args.win, "overlap": args.overlap,
        "tta": not args.no_tta, "ckpt": args.ckpt,
    }
    with open(os.path.join(out, "summary.json"), "w") as f:
        json.dump(summ, f, indent=1)
    print(f"[monai:pseudo] round {os.path.basename(out)}: {json.dumps(summ)}", flush=True)
    destroy_dist()


if __name__ == "__main__":
    main()