"""Does deep supervision make the recurrence EXTRAPOLATE past its training depth?

The claim behind supervising every recurrence is that it forces each application
of the shared block to be a valid one-step refinement. If that holds, the map is
iterable beyond where it was fit, and running more recurrences at inference than
at training should keep helping - or at least not collapse.

This is free: no retraining, only re-evaluating trained checkpoints at different
recurrence counts. A model whose accuracy peaks exactly at its training depth and
degrades past it has learned one entangled R-step function rather than a
repeatable refinement step.
"""
import argparse, glob, json, os, sys
from collections import defaultdict

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import RecurrentDenoiser
import diffusion as dfn
from sudoku import CELLS, SudokuTokenizer

p = argparse.ArgumentParser()
p.add_argument("--hard", default="data/sudoku_hard")
p.add_argument("--runs", default="runs")
p.add_argument("--pattern", default="rec1-*")
p.add_argument("--passes", type=int, default=16)
p.add_argument("--eval_bs", type=int, default=250)
p.add_argument("--out", default="figures/sudoku_rsweep.json")
a = p.parse_args()

dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = SudokuTokenizer()
P = np.load(os.path.join(a.hard, "test_puz.npy"))
S = np.load(os.path.join(a.hard, "test_sol.npy"))
Pt, St = torch.from_numpy(P), torch.from_numpy(S - 1)
blank = Pt == 0

res = defaultdict(dict)
for d in sorted(glob.glob(os.path.join(a.runs, a.pattern))):
    rf = os.path.join(d, "result.json")
    if not (os.path.exists(rf) and os.path.exists(os.path.join(d, "model.pt"))):
        continue
    cfg = json.load(open(rf))["args"]
    if not cfg.get("recurrent"):
        continue
    m = RecurrentDenoiser(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
                          max_len=CELLS + 8, recurrences=cfg["R"]).to(dev)
    m.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location=dev))
    m.eval()
    name = os.path.basename(d)
    trained_R = cfg["R"]
    print(f"\n--- {name}  (trained with R={trained_R})", flush=True)
    print(f"    {'R at inference':>16}{'board':>10}", flush=True)
    for R in (1, 2, 4, 8, 16, 32, 64, 128):
        m.recurrences = R
        preds = []
        with torch.no_grad():
            for i in range(0, len(Pt), a.eval_bs):
                gold = St[i:i + a.eval_bs].to(dev)
                bl = blank[i:i + a.eval_bs].to(dev)
                x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
                y, _ = dfn.fixed_k_decode(m, x, tok, dev, a.passes, fillable=bl)
                preds.append(y.cpu().numpy() + 1)
        acc = float((np.concatenate(preds) == S).all(axis=1).mean())
        res[name][R] = acc
        flag = "  <- trained here" if R == trained_R else ""
        print(f"    {R:>16}{acc*100:>9.2f}%{flag}", flush=True)
    del m
    torch.cuda.empty_cache()

os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
json.dump(res, open(a.out, "w"), indent=2)
print(f"\nwrote {a.out}")