"""Recalibrate the LoRA readout at 15M before it is used as a campaign response.

lora-finetune-readout.md establishes rank 16 / 3 epochs / lr 3e-3 and sigma_ft
0.00233 at scale_base only, and says so explicitly: "Anything about scales other
than scale_base" is not established. ranking-transfer.md then names that
mistuning as the leading explanation for why its readout-matched correlation
(+0.537) sits below the mismatched one (+0.630) -- the LoRA readout reaches 94%
of the full fine-tune at base but only 82% at 15M. Carrying the base recipe to
15M unchanged would repeat that, on the arm this campaign actually reports.

Three passes, all on existing 15M MAE checkpoints, one GPU so the 64-arm
campaign keeps three:

  1. epoch knee  -- where val mAP stops rising while train mAP keeps climbing
  2. lr boundary -- the base recipe's 3e-3 sits close to a divergence cliff
  3. sigma_ft    -- 3 checkpoints x 3 seeds, the noise the response carries
"""
import itertools
import json
import subprocess
import sys
from pathlib import Path

REPO = "/workspace/code/eat-map-regmix"
GPU = sys.argv[1] if len(sys.argv) > 1 else "0"

CKPT = {
    "natural": "/workspace/runs/mae-validate/natural-lr0.0005/exports/step_00039833",
    "d7": "/workspace/runs/mae-validate/d7-lr0.0007/exports/step_00039833",
    "d185": "/workspace/runs/mae-validate/d185-lr0.0007/exports/step_00039833",
}


def lora(export: str, tag: str, **kw) -> dict | None:
    cmd = ["python", "-m", "eatmap.cli.lora_finetune", "--export", export, "--tag", tag]
    for k, v in kw.items():
        cmd += [f"--{k.replace('_', '-')}", str(v)]
    proc = subprocess.run(
        cmd, cwd=REPO, capture_output=True,
        env={**__import__("os").environ, "CUDA_VISIBLE_DEVICES": GPU, "PYTHONPATH": REPO},
    )
    out = Path(export) / f"lora.{tag}.json"
    if proc.returncode or not out.exists():
        print(f"    FAIL {tag}: {proc.stderr.decode()[-300:]}")
        return None
    return json.loads(out.read_text())


def val(d: dict) -> float:
    for k in ("lora/val_map", "val_map", "lora/map", "map"):
        if k in d:
            return d[k]
    return next(v for k, v in d.items() if "map" in k.lower() and isinstance(v, float))


print("[1] epoch knee (natural, lr 3e-3)")
knee = {}
for ep in (1, 2, 3, 4, 6):
    r = lora(CKPT["natural"], f"ep{ep}", epochs=ep, lr=3e-3, seed=0)
    if r:
        knee[ep] = val(r)
        train = r.get("lora/train_map", r.get("train_map", float("nan")))
        print(f"    epochs {ep}: val {knee[ep]:.5f}   train {train:.5f}")

print("\n[2] lr boundary (natural, 3 epochs)")
best_ep = max(knee, key=knee.get) if knee else 3
for lr in (1e-3, 3e-3, 1e-2):
    r = lora(CKPT["natural"], f"lr{lr:g}", epochs=best_ep, lr=lr, seed=0)
    if r:
        print(f"    lr {lr:g}: val {val(r):.5f}")

print(f"\n[3] sigma_ft grid (3 checkpoints x 3 seeds, {best_ep} epochs, lr 3e-3)")
import statistics as st
grid = {}
for name, seed in itertools.product(CKPT, (0, 1, 2)):
    r = lora(CKPT[name], f"var-s{seed}", epochs=best_ep, lr=3e-3, seed=seed)
    if r:
        grid.setdefault(name, []).append(val(r))
for name, vals in grid.items():
    print(f"    {name:<9} {' '.join(f'{v:.5f}' for v in vals)}  "
          f"mean {st.mean(vals):.5f}  sd {st.stdev(vals):.5f}")
if len(grid) >= 2:
    pooled = (sum(st.variance(v) for v in grid.values() if len(v) > 1)
              / sum(1 for v in grid.values() if len(v) > 1)) ** 0.5
    between = st.stdev([st.mean(v) for v in grid.values()])
    print(f"\n    sigma_ft (pooled within-checkpoint) = {pooled:.5f}")
    print(f"    between-checkpoint sd               = {between:.5f}")
    if between > pooled:
        print(f"    single-run reliability              = "
              f"{(between**2 - pooled**2)/between**2:.3f}")
    print(f"    (scale_base reference: sigma_ft 0.00233, reliability 0.31)")

Path("/workspace/analysis/lora_calibration_15m.json").write_text(
    json.dumps({"epoch_knee": knee, "sigma_ft_grid": grid}, indent=1))
print("\nwrote /workspace/analysis/lora_calibration_15m.json")