"""Pseudo-label the unlabeled pool with an nnU-Net round model (multi-GPU predict).

1. Symlinks remaining pool volumes (not already in the dataset) into a predict
   folder. Volumes are NATIVE (unresampled/uncropped/unnormalized); nnU-Net
   applies the plan's preprocessing itself at predict time.
2. Runs one nnUNetv2_predict per GPU (-num_parts/-part_id), --save_probabilities.
   nnU-Net applies the plan's preprocessing at predict time and resamples the
   probabilities (and argmax seg) back to each case's native grid, so the
   selection gates run directly on that grid, in physical mm:
   pos: p_tumor >= tau_pos, median cleanup, largest-CC fraction >= min_cc_frac,
        volume within [vol_p2, vol_p98] of labeled tumor volumes (data/vols.json)
   neg: >= neg_frac of head-interior voxels have p_bg >= tau_neg
4. Per-subject longitudinal consistency filter over accepted positive timepoints
   (head-relative tumor centroid distance in mm + volume ratio; absolute
   patient-space grids are not comparable across acquisitions).
Writes <out>/rows.jsonl, <out>/<key>_label.nii.gz (native grid), <out>/accepted.jsonl,
   <out>/summary.json.

Usage:
  python scripts_nnu/04_nnu_pseudo_label.py --pool data/manifests/unlabeled_pool.jsonl \
      --out data/pseudo_nnu/round1 --gpus 3
"""
import argparse
import os
import sys
import json
import shutil
import numpy as np
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, write_arr
import SimpleITK as sitk
from nnu_common import (ROOT, d, nnu_env, predict_cmd, tumor_probs_from_npz,
                        write_label, pos_mask, neg_frac_bg, consistency_filter,
                        load_voxel_stats)


def build_input_folder(rows, in_dir):
    shutil.rmtree(in_dir, ignore_errors=True)
    os.makedirs(in_dir, exist_ok=True)
    for r in rows:
        os.symlink(os.path.abspath(r["img"]), os.path.join(in_dir, r["key"] + "_0000.nii.gz"))


def find_outputs(out_dir, ext):
    found = {}
    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(ext):
                found[fn[:-len(ext)]] = os.path.join(pd, fn)
    return found


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pool", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--already", default=None, help="jsonl of case keys already added to the dataset")
    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, help="predict subprocesses per GPU")
    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")
    ap.add_argument("--skip-predict", action="store_true", help="reuse existing prediction outputs")
    ap.add_argument("--keep-pred", action="store_true")
    args = ap.parse_args()

    outname = os.path.basename(os.path.normpath(args.out))
    out = d(args.out)
    pool_rows = load_jsonl(args.pool)
    if any("img" not in r for r in pool_rows):
        raise SystemExit(f"{args.pool} rows lack native 'img'; rerun scripts/05_build_splits.py")
    already = set()
    if args.already and os.path.exists(args.already):
        already = {r["key"] for r in load_jsonl(args.already)}
    rows = [r for r in pool_rows if r["key"] not in already]
    print(f"[nnu:pseudo:{outname}] pool={len(pool_rows)} already_in_dataset={len(already)} to_predict={len(rows)}", flush=True)

    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)

    tag = outname
    in_dir = os.path.join(d("nnu/pred"), f"in_pool_{tag}")
    out_dir = os.path.join(d("nnu/pred"), f"out_pool_{tag}")
    if rows:
        build_input_folder(rows, in_dir)
        if not args.skip_predict:
            import subprocess
            os.makedirs(out_dir, exist_ok=True)
            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_{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:pseudo:{outname}] part{g} failed rc={rc}; retrying once (skips finished cases)", flush=True)
                    e = dict(env)
                    e["CUDA_VISIBLE_DEVICES"] = str(g)
                    logf = os.path.join(d("logs/nnu"), f"pred_{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)
                        rc2 = p2.wait()
                    if rc2 != 0:
                        raise RuntimeError(f"prediction part{g} failed twice; see {logf}")

    npz = find_outputs(out_dir, ".npz") if os.path.isdir(out_dir) else {}
    rows_out, n_err = [], 0
    for r in rows:
        key = r["key"]
        entry = {"key": key, "subject": r["subject"], "date": r.get("date"),
                 "source": r.get("source"), "img": r["img"], "pimg": r.get("pimg"), "label": None,
                 "role": "rej", "vol_mm3": 0, "maxp": None}
        try:
            img_itk = sitk.ReadImage(r["img"])
            vol = read_nii_arr(r["img"]).astype("float32")
            pt = tumor_probs_from_npz(npz[key])
            if pt.shape != vol.shape:
                raise ValueError(f"prob shape {pt.shape} != native image {vol.shape} for {key}")
            entry["maxp"] = round(float(pt.max()), 4)
            voxel_vol_mm3 = float(np.prod(np.asarray(img_itk.GetSpacing(), dtype=np.float64)))
            got = pos_mask(pt, args.tau_pos, args.min_cc_frac, vol_lo, vol_hi, voxel_vol_mm3)
            if got is not None:
                mask, cc_frac, vol_mm3 = got
                lp = os.path.join(out, key + "_label.nii.gz")
                write_label(mask, lp, img_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=img_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"[nnu:pseudo:{outname}] {key} ERR {e!r}", flush=True)
        rows_out.append(entry)

    n_pos0 = sum(1 for x in rows_out if x["role"] == "pos")
    n_neg0 = sum(1 for x in rows_out if x["role"] == "neg")
    n_rej = consistency_filter(rows_out, out, args.max_rel_dist, args.vol_ratio, image_key="img")
    for x in rows_out:
        if x["role"] == "rej":
            x["label"] = None
    accepted = [x for x in rows_out if x["role"] in ("pos", "neg")]
    save_jsonl(rows_out, os.path.join(out, "rows.jsonl"))
    save_jsonl(accepted, os.path.join(out, "accepted.jsonl"))
    posv = [x["vol_mm3"] for x in rows_out if x["role"] == "pos"]
    summ = {
        "n_pool_predicted": len(rows_out),
        "n_predicted": len(npz),
        "n_pos": sum(1 for x in rows_out if x["role"] == "pos"),
        "n_neg": sum(1 for x in rows_out if x["role"] == "neg"),
        "n_pos_before_consistency": n_pos0,
        "n_rejected_consistency": n_rej,
        "n_other_rej": sum(1 for x in rows_out if x["role"] == "rej"),
        "n_error": n_err,
        "pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0,
                        "p5": float(float(np.percentile(posv, 5))) if posv else 0,
                        "p95": float(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, "chk": args.chk, "tta": not args.no_tta,
    }
    with open(os.path.join(out, "summary.json"), "w") as f:
        json.dump(summ, f, indent=1)
    print(f"[nnu:pseudo] round {outname}: {json.dumps(summ)}", flush=True)

    if not args.keep_pred and n_err == 0:
        shutil.rmtree(in_dir, ignore_errors=True)
        shutil.rmtree(out_dir, ignore_errors=True)
    elif n_err > 0:
        print(f"[nnu:pseudo:{outname}] kept prediction outputs ({in_dir}, {out_dir}) due to {n_err} errors; "
              f"rerun with --skip-predict --keep-pred after fixing", flush=True)


if __name__ == "__main__":
    main()