"""LoRA readout on the 64 final MAE checkpoints, 3 pinned seeds, recalibrated recipe.

Recipe is the 15M recalibration, not the scale_base defaults:
  rank 16, targets qkv/proj/fc1/fc2, LayerNorm trainable, lr 3e-3,
  **4 epochs** (the knee at 15M; 3 is the knee at scale_base and costs ~0.007 here).

Seeds are pinned 0/1/2 across every arm so the common-mode component of sigma_ft
(0.00076 of 0.00167 at this scale) cancels in the arm contrast, exactly as the
base campaign did. Same CPU-thread capping as the probe stage: this is a
BLAS-bound workload and uncapped OpenMP costs ~10x.
"""
import argparse
import json
import os
import queue
import subprocess
import threading
from pathlib import Path

REPO = "/workspace/code/eat-map-regmix"
EPOCHS = 4
LR = 3e-3
SEEDS = (0, 1, 2)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("root", default="/workspace/runs/mae-64", nargs="?")
    ap.add_argument("--gpus", default="0,1,2,3")
    ap.add_argument("--per-gpu", type=int, default=2)
    ap.add_argument("--threads", type=int, default=12)
    args = ap.parse_args()

    jobs: queue.Queue = queue.Queue()
    for run in sorted(Path(args.root).iterdir()):
        exports = sorted((run / "exports").glob("step_*")) if run.is_dir() else []
        if not exports:
            continue
        final = exports[-1]
        for seed in SEEDS:
            if not (final / f"lora_s{seed}.json").exists():
                jobs.put((run.name, final, seed))
    total = jobs.qsize()
    print(f"{total} LoRA fine-tunes ({EPOCHS} epochs, lr {LR:g}, seeds {SEEDS})")
    lock = threading.Lock()
    done = [0]

    def worker(gpu):
        while True:
            try:
                name, exp, seed = jobs.get_nowait()
            except queue.Empty:
                return
            t = str(args.threads)
            proc = subprocess.run(
                ["python", "-m", "eatmap.cli.lora_finetune", "--export", str(exp),
                 "--tag", f"s{seed}", "--epochs", str(EPOCHS), "--lr", str(LR),
                 "--seed", str(seed)],
                cwd=REPO, capture_output=True,
                env={**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu), "PYTHONPATH": REPO,
                     "OMP_NUM_THREADS": t, "OPENBLAS_NUM_THREADS": t,
                     "MKL_NUM_THREADS": t, "NUMEXPR_NUM_THREADS": t},
            )
            with lock:
                done[0] += 1
                if proc.returncode:
                    print(f"  FAIL {name} s{seed}: {proc.stderr.decode()[-200:]}")
                elif done[0] % 20 == 0:
                    print(f"  {done[0]}/{total}")
            jobs.task_done()

    threads = [threading.Thread(target=worker, args=(int(g),))
               for g in args.gpus.split(",") for _ in range(args.per_gpu)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    # ---- collect, and screen for the collapse mode ranking-transfer.md found --
    rows = {}
    for run in sorted(Path(args.root).iterdir()):
        exports = sorted((run / "exports").glob("step_*")) if run.is_dir() else []
        if not exports:
            continue
        vals = []
        for seed in SEEDS:
            p = exports[-1] / f"lora_s{seed}.json"
            if p.exists():
                vals.append(json.loads(p.read_text())["lora/map"])
        if len(vals) == len(SEEDS):
            import statistics as st
            rows[run.name] = {"vals": vals, "mean": st.mean(vals), "sd": st.stdev(vals)}

    bad = [k for k, v in rows.items() if v["sd"] > 0.02]
    print(f"\n{len(rows)} arms scored")
    print(f"collapse screen (lora sd > 0.02): {bad if bad else 'none'}")
    print("  ranking-transfer.md saw ~1% incidence at lr 3e-3, non-deterministic; re-run these")
    Path(args.root, "lora_summary.json").write_text(json.dumps(rows, indent=1))
    print(f"wrote {args.root}/lora_summary.json")


if __name__ == "__main__":
    main()