"""Step-2(a): PTRM test-time noise + lambda-based selection on a trained coloring model
(any backbone feature set via cfg pe). Writes a JSON per ckpt for multi-seed aggregation.

deterministic / pass@K (conflict-min, ground truth) / lambda-select (min lambda1) / random.

Run: PYTHONPATH=/home/yurenh2/rrog python3 diag/ptrm_color.py --ckpt runs/ckpt_color_rrog_trm_gin_full_...pt
"""
import argparse, json, os
import numpy as np
import torch
from diag.train_color import RecGINColor, make_split, featurize
try:
    from sklearn.metrics import roc_auc_score
except Exception:
    roc_auc_score = None
OUT = '/home/yurenh2/rrog/runs'


def rollout(model, xin, ei, sigma, n_sup, T, dev, seed):
    gen = torch.Generator(device=dev).manual_seed(seed)
    ctx = model.aggregate(xin, ei)
    y, z = ctx, torch.zeros_like(ctx)
    state = torch.cat([y, z], dim=-1)
    v = torch.randn(state.shape, generator=gen, device=dev); v = v / (v.norm() + 1e-12)
    def step(ss):
        yy, zz = ss.chunk(2, dim=-1)
        yy, zz = model.recurse(yy, zz, ctx, noise=False)
        return torch.cat([yy, zz], dim=-1)
    lam = 0.0
    for _ in range(n_sup * T):
        state_det, Jv = torch.autograd.functional.jvp(step, state, v)
        nv = Jv.norm(); lam += torch.log(nv + 1e-12).item(); v = (Jv / (nv + 1e-12)).detach()
        state = state_det.detach()
        if sigma > 0:
            state = state + sigma * torch.randn(state.shape, generator=gen, device=dev)
    lam /= max(n_sup * T, 1)
    y, _ = state.chunk(2, dim=-1)
    col = model.head(y).argmax(-1)
    conf = (col[ei[0]] == col[ei[1]]).sum().item() // 2
    return conf, lam


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True)
    ap.add_argument('--K', type=int, default=16)
    ap.add_argument('--n_graphs', type=int, default=150)
    ap.add_argument('--sigmas', type=float, nargs='+', default=[0.05, 0.1, 0.2, 0.4])
    args = ap.parse_args()
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    ck = torch.load(args.ckpt, weights_only=False); c = ck['cfg']
    deg = torch.tensor(c['deg']) if c.get('deg') else None
    model = RecGINColor(c['in_dim'], c['hidden'], c['k'], c['T'], c['n_sup'],
                        grad_mode=c['grad_mode'], conv=c.get('conv', 'gin'), deg=deg,
                        agg_layers=c.get('agg_layers', 1),
                        compute_layers=c.get('compute_layers', 2),
                        compute=(c.get('compute') if c.get('compute') == 'trm' else 'trm'),
                        attn_heads=c.get('attn_heads', 4)).to(dev)
    model.load_state_dict(ck['state']); model.eval()
    nsup, T = c['n_sup'], c['T']
    te = featurize(make_split('test', 50, 3, 0.2, 8, 500, 100000), c.get('pe', 'none'), c.get('rwse_k', 16))
    te = te[:args.n_graphs]; n = len(te)

    det = sum(rollout(model, r['xin'].to(dev), r['edge_index'].to(dev), 0.0, nsup, T, dev, 0)[0] == 0
              for r in te) / n
    out = {'conv': c.get('conv', 'gin'), 'pe': c.get('pe', 'none'), 'seed': c.get('seed'),
           'arch': c.get('arch', 'legacy'),
           'grad_mode': c['grad_mode'], 'contract': c.get('contract', False), 'det': det, 'sigmas': {}}
    print(f"[pe={out['pe']} s{out['seed']}] deterministic solve_rate = {det:.3f}  (n={n}, K={args.K})")
    print(f"{'sigma':>6} {'pass@K':>8} {'lam-sel':>8} {'random':>8} {'perRoll':>8} {'AUROC(s|-lam)':>14}")
    for sigma in args.sigmas:
        passk = lamsel = rand = 0
        L, S = [], []
        for gi, r in enumerate(te):
            xin = r['xin'].to(dev); ei = r['edge_index'].to(dev)
            res = [rollout(model, xin, ei, sigma, nsup, T, dev, 1000 * gi + j) for j in range(args.K)]
            confs = np.array([c0 for c0, _ in res]); lams = np.array([l for _, l in res])
            solved = confs == 0
            passk += int(solved.any()); lamsel += int(solved[lams.argmin()]); rand += int(solved[0])
            L += lams.tolist(); S += solved.tolist()
        L, S = np.array(L), np.array(S)
        auc = (roc_auc_score(S.astype(int), -L) if roc_auc_score and S.any() and (~S).any() else float('nan'))
        out['sigmas'][str(sigma)] = {'passk': passk / n, 'lamsel': lamsel / n, 'random': rand / n,
                                     'perRoll': float(S.mean()), 'auroc': float(auc)}
        print(f"{sigma:>6} {passk/n:>8.3f} {lamsel/n:>8.3f} {rand/n:>8.3f} {S.mean():>8.3f} {auc:>14.3f}")

    base = os.path.basename(args.ckpt).replace('ckpt_', '').replace('.pt', '')
    with open(os.path.join(OUT, f"ptrm_{base}.json"), 'w') as f:
        json.dump(out, f, indent=2)
    print("  wrote", os.path.join(OUT, f"ptrm_{base}.json"))


if __name__ == "__main__":
    main()