"""run_21b1618aa7 bounded DPPO activation diagnostic (REPLACEMENT loop; prior run_19badd842c code/rx.py was not recoverable).
One frozen 32x8 GSM8K train batch, 8 minibatches of 4 complete groups, 2 epochs, <=16 AdamW updates, DPPO Binary-TV loss
(code/losses_patched.py = run_19badd842c losses.py + instrumentation-only patch). Counterfactual PPO diagnostics on same tensors."""
import asyncio, hashlib, json, os, sys, time
import numpy as np, torch, tinker
from tinker import types
sys.path.insert(0, os.path.dirname(__file__))
import losses_patched as L
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.recipes.math_rl.math_env import Gsm8kDataset, MathEnv
from tinker_cookbook.completers import TinkerTokenCompleter
from tinker_cookbook.rl.rollouts import do_group_rollout
from tinker_cookbook.rl.data_processing import compute_advantages, assemble_training_data

T0 = time.time(); CAP = 720.0
OUT = os.path.join(os.path.dirname(__file__), "..", "out"); os.makedirs(OUT, exist_ok=True)
MODEL = "Qwen/Qwen3.5-2B"; INIT = "tinker://21e32976f18f971463410adff4e364e4/weights/rx19badd842c_init"
ADAM = types.AdamParams(learning_rate=2e-5, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.0, grad_clip_norm=0.0)
DELTA, EPS, C = L.DPPO_DELTA, L.GRPO_EPS, L.DPPO_CLIP_RATIO_C
assert (DELTA, EPS, C) == (0.15, 0.2, 10.0)
log = open(os.path.join(OUT, "events.jsonl"), "a")
def ev(**k):
    k["t"] = round(time.time() - T0, 2); log.write(json.dumps(k) + "\n"); log.flush(); print(json.dumps(k)[:400], flush=True)
def left(): return CAP - (time.time() - T0)

def dhash(d):
    h = hashlib.sha256()
    h.update(np.asarray(d.model_input.to_ints(), dtype=np.int64).tobytes())
    for k in ("target_tokens", "logprobs", "advantages", "mask"):
        h.update(k.encode()); h.update(d.loss_fn_inputs[k].to_torch().numpy().tobytes())
    return h.hexdigest()
def bhash(ds): return hashlib.sha256("".join(dhash(d) for d in ds).encode()).hexdigest()

def q(x, p): return float(torch.quantile(x.double(), p)) if x.numel() else None
BINS = [0, 0.01, 0.1, 0.3, 0.5, 0.7, 0.9, 0.99, 1.0000001]

def stats(lp, olp, adv, mask):
    m = mask > 0; lp, olp, a = lp[m], olp[m], adv[m]
    p, mu = lp.exp(), olp.exp(); d = p - mu; r = (lp - olp).exp()
    masked = ((a > 0) & (d > DELTA)) | ((a < 0) & (d < -DELTA))
    clip = ((a > 0) & (r > 1 + EPS)) | ((a < 0) & (r < 1 - EPS))
    out = {"n_loss_tokens": int(m.sum()), "n_pos_adv": int((a > 0).sum()), "n_neg_adv": int((a < 0).sum()), "n_zero_adv": int((a == 0).sum()),
           "dppo_masked": int(masked.sum()), "dppo_masked_pos": int((masked & (a > 0)).sum()), "dppo_masked_neg": int((masked & (a < 0)).sum()),
           "ppo_clipped": int(clip.sum()), "ppo_clipped_pos": int((clip & (a > 0)).sum()), "ppo_clipped_neg": int((clip & (a < 0)).sum()),
           "abs_drift_gt_delta": int((d.abs() > DELTA).sum()), "drift_gt_delta": int((d > DELTA).sum()), "drift_lt_neg_delta": int((d < -DELTA).sum()),
           "ratio_cap_hits": int((r > C).sum()), "ratio_cap_hits_nonzero_adv": int(((r > C) & (a != 0)).sum()),
           "absd_q90": q(d.abs(), .9), "absd_q99": q(d.abs(), .99), "absd_max": float(d.abs().max()),
           "d_min": float(d.min()), "d_max": float(d.max()), "d_q01": q(d, .01), "d_q99": q(d, .99),
           "d_pos_adv_max": float(d[a > 0].max()) if (a > 0).any() else None, "d_neg_adv_min": float(d[a < 0].min()) if (a < 0).any() else None,
           "ratio_min": float(r.min()), "ratio_max": float(r.max()), "ratio_q99": q(r, .99), "ratio_q01": q(r, .01),
           "mu_bins": np.histogram(mu.double().numpy(), BINS)[0].tolist(), "dtype": str(lp.dtype)}
    return out, masked

async def main():
    sc = tinker.ServiceClient(base_url=os.environ["TINKER_BASE_URL"], api_key=os.environ["TINKER_API_KEY"],
        default_headers_overrides={"Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"], "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"]})
    tok = get_tokenizer(MODEL); ren = renderers.get_renderer("qwen3_5_disable_thinking", tokenizer=tok)
    ds = Gsm8kDataset(batch_size=32, group_size=8, renderer=ren, convo_prefix=MathEnv.standard_fewshot_prefix(), split="train", seed=0)
    builders = ds.get_batch(0); assert len(builders) == 32
    # restore immutable init (retry on capacity)
    t = time.time(); tc = None; delay = 30
    while tc is None:
        try:
            tc = await sc.create_training_client_from_state_async(INIT, base_model=MODEL)
        except Exception as e:
            ev(stage="restore_error", err=repr(e)[:500])
            if left() < 600: raise
            await asyncio.sleep(delay); delay = min(delay * 2, 120)
    info = await tc.get_info_async()
    ev(stage="restored", seconds=time.time() - t, init=INIT, info=json.loads(info.model_dump_json()) if hasattr(info, "model_dump_json") else str(info))
    t = time.time()
    sres = await (await tc.save_weights_for_sampler_async("rx21b_diag_sampler_init"))
    spath = sres.path
    samp = await sc.create_sampling_client_async(model_path=spath)
    ev(stage="sampler", seconds=time.time() - t, sampler_path=str(spath))
    t = time.time()
    pol = TinkerTokenCompleter(samp, max_tokens=512, temperature=1.0)
    groups = await asyncio.gather(*[do_group_rollout(b, pol) for b in builders])
    adv = compute_advantages(groups)
    data, meta = assemble_training_data(groups, adv)
    assert len(data) == 256 and [m["group_idx"] for m in meta] == [i // 8 for i in range(256)]
    rewards = [r for g in groups for r in g.get_total_rewards()]
    ev(stage="sampled", seconds=time.time() - t, reward_mean=float(np.mean(rewards)), n_datums=len(data))
    # freeze
    H = [dhash(d) for d in data]; BH = hashlib.sha256("".join(H).encode()).hexdigest()
    mbs = [list(range(32 * k, 32 * (k + 1))) for k in range(8)]
    MBH = [hashlib.sha256("".join(H[i] for i in mb).encode()).hexdigest() for mb in mbs]
    sent = [mb[0] for mb in mbs]
    frozen = {"batch_sha256": BH, "minibatch_sha256": MBH, "datum_sha256": H, "sampler_path": str(spath), "sentinel_datums": sent,
              "rewards": rewards, "group_idx": [m["group_idx"] for m in meta]}
    json.dump(frozen, open(os.path.join(OUT, "frozen-batch.json"), "w"))
    np.savez_compressed(os.path.join(OUT, "frozen-batch-tensors.npz"), **{f"{k}_{i}": (np.asarray(d.model_input.to_ints()) if k == "input" else d.loss_fn_inputs[k].to_torch().numpy())
        for i, d in enumerate(data) for k in ("input", "target_tokens", "logprobs", "advantages", "mask")})
    ev(stage="frozen", batch_sha256=BH)

    sent_data = [types.Datum(model_input=data[i].model_input, loss_fn_inputs={"target_tokens": data[i].loss_fn_inputs["target_tokens"]}) for i in sent]
    sentinel = {}
    async def snap(tag):
        out = await (await tc.forward_async(sent_data, "cross_entropy"))
        lps = [o["logprobs"].to_torch() for o in out.loss_fn_outputs]
        rows = []
        for i, lp in zip(sent, lps):
            m = data[i].loss_fn_inputs["mask"].to_torch() > 0
            rows.append(lp[m].float())
        cat = torch.cat(rows); olp = torch.cat([data[i].loss_fn_inputs["logprobs"].to_torch()[data[i].loss_fn_inputs["mask"].to_torch() > 0].float() for i in sent])
        sentinel[tag] = cat
        d0 = (cat.exp() - olp.exp()).abs()
        rec = {"stage": "sentinel", "tag": tag, "n": int(cat.numel()), "sha256": hashlib.sha256(cat.numpy().tobytes()).hexdigest(),
               "vs_sampler_absdp_q90": q(d0, .9), "vs_sampler_absdp_max": float(d0.max()), "vs_sampler_mean_abs_dlogp": float((cat - olp).abs().mean())}
        if "init" in sentinel and tag != "init":
            di = (cat.exp() - sentinel["init"].exp()).abs()
            rec.update(vs_init_absdp_q90=q(di, .9), vs_init_absdp_max=float(di.max()), vs_init_mean_abs_dlogp=float((cat - sentinel["init"]).abs().mean()))
        ev(**rec)
    await snap("init")

    calls = []
    upd = 0
    for epoch in (1, 2):
        for k, mb in enumerate(mbs):
            if left() < 25:
                ev(stage="wall_cap_stop", updates=upd); break
            batch = [data[i] for i in mb]
            assert hashlib.sha256("".join(dhash(d) for d in batch).encode()).hexdigest() == MBH[k], "frozen hash mismatch"
            rec = {}
            def loss_fn(dl, lps):
                tot = 0.0; cat = {x: [] for x in ("lp", "olp", "adv", "mask")}
                for d, lp in zip(dl, lps, strict=True):
                    inp = d.loss_fn_inputs
                    olp = inp["logprobs"].to_torch().to(lp.dtype); a = inp["advantages"].to_torch().to(lp.dtype); m = inp["mask"].to_torch().to(lp.dtype)
                    tot = tot + L.token_terms(lp, olp, a, m, "dppo").sum()
                    for x, v in zip(cat, (lp.detach(), olp, a, m)): cat[x].append(v)
                c = {x: torch.cat(v) for x, v in cat.items()}
                s, masked = stats(c["lp"], c["olp"], c["adv"], c["mask"])
                s["patched_diag"] = L.diagnostics(c["lp"], c["olp"], c["adv"], c["mask"])
                rec.update(s)
                np.savez_compressed(os.path.join(OUT, f"call-u{upd+1:02d}.npz"), current_lp=c["lp"].float().numpy(), sampler_lp=c["olp"].float().numpy(),
                    adv=c["adv"].float().numpy(), mask=c["mask"].float().numpy(), dppo_masked_on_loss_tokens=masked.numpy(),
                    target_tokens=torch.cat([d.loss_fn_inputs["target_tokens"].to_torch() for d in dl]).numpy())
                return tot, {}
            fb = await (await tc.forward_backward_custom_async(batch, loss_fn))
            # verify removed tokens actually receive zero gradient: recompute surrogate grad on saved tensors
            z = np.load(os.path.join(OUT, f"call-u{upd+1:02d}.npz"))
            lp = torch.tensor(z["current_lp"], dtype=torch.float32, requires_grad=True)
            L.token_terms(lp, torch.tensor(z["sampler_lp"]), torch.tensor(z["adv"]), torch.tensor(z["mask"]), "dppo").sum().backward()
            mk = torch.tensor(z["mask"]) > 0; g = lp.grad[mk]; a = torch.tensor(z["adv"])[mk]; msk = torch.tensor(z["dppo_masked_on_loss_tokens"])
            # PPO counterfactual: tokens whose PPO grad is zero (clipped) but nonzero adv
            lp2 = torch.tensor(z["current_lp"], dtype=torch.float32, requires_grad=True)
            L.token_terms(lp2, torch.tensor(z["sampler_lp"]), torch.tensor(z["adv"]), torch.tensor(z["mask"]), "grpo").sum().backward()
            g2 = lp2.grad[mk]
            rec.update(removed_nonzero_grad_pos=int((msk & (a > 0)).sum()), removed_nonzero_grad_neg=int((msk & (a < 0)).sum()),
                       removed_tokens_grad_max_abs=float(g[msk].abs().max()) if msk.any() else None,
                       ppo_zero_grad_nonzero_adv=int(((g2 == 0) & (a != 0)).sum()))
            o = await (await tc.optim_step_async(ADAM)); upd += 1
            rec.update(stage="update", update=upd, epoch=epoch, minibatch=k, minibatch_sha256=MBH[k], n_datums=len(batch),
                       intervening_updates_since_sampling=upd - 1, optim_metrics=getattr(o, "metrics", None))
            calls.append(rec); ev(**rec)
            if upd in (1, 2, 4, 8, 16): await snap(f"after_u{upd}")
        else:
            continue
        break
    json.dump(calls, open(os.path.join(OUT, "per-update.json"), "w"))
    t = time.time()
    try:
        st = await (await tc.save_state_async("rx21b_diag_final"))
        ev(stage="saved_state", path=st.path, seconds=time.time() - t)
    except Exception as e:
        ev(stage="save_state_error", err=repr(e)[:300])
    ev(stage="done", updates=upd, wall=time.time() - T0)

asyncio.run(main())