"""Orchestrates the iterative pseudo-labeling study.

Round 0: baseline supervised training on labeled T1c (patient-level split).
Round k: pseudo-label the unlabeled pool with round k-1 model (confidence +
         volume plausibility + longitudinal consistency), then fine-tune with
         labeled + pseudo data (warm restart, lower LR). After each round the
         model is evaluated on the held-out patient-level test split.

Usage: python scripts/09_run_iterative.py [--rounds 4] [--gpus 3]
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import argparse
import json
import shutil
import subprocess
from src.common import ROOT, load_jsonl, save_jsonl


def run(cmd, log, retries=2):
    for attempt in range(retries + 1):
        try:
            with open(log, "a") as f:
                f.write(f"$ (attempt {attempt + 1}) " + cmd + "\n")
                print(f"$ (attempt {attempt + 1}) " + cmd, flush=True)
                subprocess.run(cmd, shell=True, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT, check=True)
            return
        except subprocess.CalledProcessError:
            if attempt == retries:
                raise
            print(f"[orch] command failed, retrying in 60s: {cmd}", flush=True)
            import time
            time.sleep(60)


def torchrun(nproc, script, extra):
    tr = shutil.which("torchrun") or (sys.executable + " -m torch.distributed.run")
    return f"{tr} --standalone --nproc_per_node {nproc} {script} {extra}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rounds", type=int, default=4)
    ap.add_argument("--gpus", type=int, default=3)
    ap.add_argument("--base-epochs", type=int, default=40)
    ap.add_argument("--round-epochs", type=int, default=12)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--round-lr", type=float, default=8e-5)
    ap.add_argument("--pseudo-weight", type=float, default=0.3)
    ap.add_argument("--batch", type=int, default=3)
    args = ap.parse_args()

    man = os.path.join(ROOT, "data/manifests")
    train_f = os.path.join(man, "split_train.jsonl")
    val_f = os.path.join(man, "split_val.jsonl")
    test_f = os.path.join(man, "split_test.jsonl")
    pool_f = os.path.join(man, "unlabeled_pool.jsonl")
    results_dir = f"{ROOT}/results"
    os.makedirs(results_dir, exist_ok=True)
    logdir = f"{ROOT}/logs"
    os.makedirs(logdir, exist_ok=True)
    main_log = os.path.join(logdir, "iterative.log")

    # ---- round 0: baseline ----
    run(torchrun(args.gpus, "scripts/07_train.py",
                 f"--rows {train_f} --val {val_f} --epochs {args.base_epochs} --lr {args.lr} "
                 f"--batch {args.batch} --ckpt-dir runs/round0"), f"{main_log}")
    run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round0/best.pt", f"{main_log}")

    for k in range(1, args.rounds + 1):
        pdir = os.path.join(ROOT, "data/pseudo/round" + str(k))
        run(torchrun(args.gpus, "scripts/06_pseudo_label.py",
                     f"--ckpt runs/round{k-1}/best.pt --unlabeled {pool_f} --out {pdir}"), f"{main_log}")
        # build round-k training manifest: labeled + accepted pseudo rows (weighted)
        pr = load_jsonl(os.path.join(pdir, "rows.jsonl"))
        used = []
        for r in pr:
            if r["role"] == "pos":
                r["w"] = args.pseudo_weight
                used.append(r)
            elif r["role"] == "neg":
                r["w"] = args.pseudo_weight
                r["is_neg"] = True
                r["plabel"] = None
                used.append(r)
        trf = os.path.join(pdir, "train_rows.jsonl")
        save_jsonl(used, trf)
        run(torchrun(args.gpus, "scripts/07_train.py",
                     f"--rows {train_f},{trf} --val {val_f} --epochs {args.round_epochs} "
                     f"--lr {args.round_lr} --batch {args.batch} --resume runs/round{k-1}/best.pt "
                     f"--ckpt-dir runs/round{k} --val-every 1"), f"{main_log}")
        run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round{k}/best.pt", f"{main_log}")

    # ---- report ----
    table = []
    for k in list(range(args.rounds + 1)):
        f = os.path.join(results_dir, f"round{k}_test.json")
        if os.path.exists(f):
            r = json.load(open(f))
            table.append({"round": k, "test_dice": round(r["dice"], 4), "n_test": r["n"],
                          "ckpt_epoch": r.get("epoch")})
            pf = os.path.join(ROOT, f"data/pseudo/round{k}/summary.json") if k > 0 else None
            if pf and os.path.exists(pf):
                s = json.load(open(pf))
                table[-1].update({"n_pos": s["n_pos"], "n_neg": s["n_neg"],
                                  "n_rej_cons": s["n_rejected_consistency"],
                                  "pos_vol_med_mm3": s["pos_vol_mm3"]["med"]})
        sf = os.path.join(ROOT, f"runs/round{k}/final.pt")
        table[-1]["ckpt"] = sf if os.path.exists(sf) else ""
    save_jsonl(table, os.path.join(results_dir, "iterative_table.jsonl"))
    print(json.dumps(table, indent=1))
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        rs = [t["round"] for t in table]
        ds = [t["test_dice"] for t in table]
        plt.figure(figsize=(6, 4))
        plt.plot(rs, ds, "o-")
        plt.xlabel("pseudo-labeling round")
        plt.ylabel("holdout tumor Dice")
        for x, y in zip(rs, ds):
            plt.annotate(f"{y:.3f}", (x, y), textcoords="offset points", xytext=(0, 8), fontsize=8)
        plt.grid(alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(results_dir, "iterative_dice.png"), dpi=150)
    except Exception as e:  # noqa
        print("plot failed:", repr(e))


if __name__ == "__main__":
    main()