"""Claim-5 matched-learning-rate control.

The paper's own model-selection rule (pick the learning rate with the best
zero-shot Humanities accuracy) happened to select very different step sizes for
the two arms in our runs -- 1e-5 for all-attention (Q/K/V) but 3e-4 for
value-matrix-only (V). The two fine-tuned models therefore sit at very
different distances from the base model, which confounds any comparison of how
much in-context ability each one retained.

This control removes that confound: it trains BOTH arms at the SAME learning
rate (1e-4, the middle of the paper's grid, where both arms score well on the
zero-shot objective) with the same seeds, and re-measures 0-shot and 7-shot
accuracy on all four categories.

Emits results/claim5_matched.json.
"""
import argparse, json, os, torch
import mmlu_data as MD
import claim5_mmlu_qwen_finetune as C


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="/work/results")
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--seeds", type=int, default=2)
    ap.add_argument("--bsz", type=int, default=32)
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)

    C.log("building MMLU data ...")
    data = MD.build()
    tok = C.load_tok()
    out = {"model": C.MODEL_ID, "matched_lr": a.lr, "lora_r": C.LORA_R,
           "epochs": C.EPOCHS, "n_eval_per_cat": MD.N_EVAL_PER_CAT,
           "n_shot": MD.N_SHOT, "arms": {}}
    path = os.path.join(a.out, "claim5_matched.json")

    def dump():
        json.dump(out, open(path, "w"), indent=2)

    C.log("=== BASE MODEL (re-measured for a self-contained control) ===")
    base = C.load_base()
    out["base"], out["base_hits"] = C.eval_all(base, tok, data, batch_size=a.bsz)
    del base; torch.cuda.empty_cache()
    dump()

    for tgt in ("qkv", "v"):
        out["arms"][tgt] = {"lr": a.lr, "seeds": []}
        for s in range(a.seeds):
            C.log(f"=== MATCHED target={tgt} lr={a.lr:g} seed={s} ===")
            m, meta = C.train_lora(tok, data, tgt, a.lr, seed=s)
            acc, hits = C.eval_all(m, tok, data, batch_size=a.bsz)
            out["arms"][tgt]["seeds"].append(
                {"seed": s, "acc": acc, "hits": hits, **meta})
            del m; torch.cuda.empty_cache()
            dump()
    C.log("MATCHED CONTROL DONE")
    dump()


if __name__ == "__main__":
    main()