"""Test-time loop scaling (TRM host). Train depth is L_cycles=n_backwards_L=6;
this evaluates exact-puzzle accuracy when the loop depth is scaled up at test time
(6, 9, 12, ...). The thesis test: does loop-attn keep using extra loop depth
(accuracy rises / holds) while the carry-last baseline saturates or degrades?

Loads EMA weights. Scales BOTH L_cycles (warmup cycles) and n_backwards_L (the
grad-cycle length, also used at eval) on the inner config; H_cycles and
halt_max_steps are held fixed so only the per-cycle loop depth changes.

Usage:
  DISABLE_COMPILE=1 python3 eval_loopscale.py checkpoints/trm-lar-s0 \
      --depths 6,9,12,18,24 --max_batches 20
"""
import os, sys, argparse, glob
os.environ.setdefault("DISABLE_COMPILE", "1")
os.environ.setdefault("WANDB_MODE", "offline")

import torch
from omegaconf import OmegaConf

from pretrain_config import PretrainConfig
from create_model import create_model
from pretrain import create_dataloader, autocast_ctx
from models.losses import IGNORE_LABEL_ID


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("ckpt_dir")
    ap.add_argument("--depths", default="6,9,12,18,24")
    ap.add_argument("--step", default="")
    ap.add_argument("--max_batches", type=int, default=20)
    args = ap.parse_args()

    cfg_d = OmegaConf.to_container(OmegaConf.load(os.path.join(args.ckpt_dir, "all_config.yaml")), resolve=True)
    cfg_d["load_checkpoint"] = None
    cfg_d["resume_from"] = None
    cfg_d["metrics_out"] = None
    base = PretrainConfig(**cfg_d)

    eval_loader, eval_metadata = create_dataloader(
        base, "test", test_set_mode=True, epochs_per_iter=1,
        global_batch_size=base.global_batch_size, rank=0, world_size=1)
    model, _, _ = create_model(base, eval_metadata, rank=0, world_size=1, strict_load=False)

    strip = lambda sd: {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
    bundles = sorted(glob.glob(os.path.join(args.ckpt_dir, "step_*_train_state.pt")),
                     key=lambda p: int(p.split("step_")[1].split("_")[0]))
    bundle = (os.path.join(args.ckpt_dir, f"step_{args.step}_train_state.pt") if args.step
              else (bundles[-1] if bundles else None))
    if bundle and os.path.exists(bundle):
        # EMA-eval bundle: EMA tracks only the dense params (16 keys); puzzle_emb/H_init/
        # L_init live in d["model"] only. Load the full live state first, then overlay
        # the EMA shadows (= the eval-time weights training uses).
        d = torch.load(bundle, map_location="cpu", weights_only=False)
        model.load_state_dict(strip(d["model"]), strict=False)
        model.load_state_dict(strip(d["ema"]), strict=False)
        src = os.path.basename(bundle) + " (EMA)"
    else:
        # Raw intermediate checkpoint (checkpoint_every_n_steps) — a flat model
        # state_dict, no EMA shadow. Load it directly.
        raw = os.path.join(args.ckpt_dir, f"step_{args.step}")
        d = torch.load(raw, map_location="cpu", weights_only=False)
        model.load_state_dict(strip(d), strict=False)
        src = os.path.basename(raw) + " (raw)"
    model.eval()
    icfg = model.model.inner.config
    has_la = model.model.inner.loop_attn is not None
    print(f"{args.ckpt_dir}: loop_attn={has_la}  ckpt={src}  "
          f"train depth L_cycles={icfg.L_cycles} n_backwards_L={icfg.n_backwards_L}")

    depths = [int(x) for x in args.depths.split(",") if x]
    print(f"\n  depth | exact-puzzle acc  (max_batches={args.max_batches})")
    for D in depths:
        icfg.L_cycles = D
        icfg.n_backwards_L = D
        exact_n = exact_tot = 0
        with torch.no_grad():
            for bi, (set_name, batch, _) in enumerate(eval_loader):
                if bi >= args.max_batches:
                    break
                batch = {k: v.cuda() for k, v in batch.items()}
                with torch.device("cuda"):
                    carry = model.initial_carry(batch)
                while True:
                    with autocast_ctx(base):
                        carry, _, _, _, preds, all_finish = model(carry=carry, batch=batch, return_keys=["preds"])
                    if all_finish:
                        break
                p, lab = preds["preds"], batch["labels"]
                valid = lab != IGNORE_LABEL_ID
                cor = valid & (p == lab)
                vs, cs = valid.sum(1), cor.sum(1)
                sv = vs > 0
                exact_n += int((sv & (cs == vs)).sum().item())
                exact_tot += int(sv.sum().item())
        acc = exact_n / max(exact_tot, 1)
        print(f"  {D:5d} | {acc:.4f}   (n={exact_tot})")


if __name__ == "__main__":
    main()