"""CPU parity: torch sc_terms implementation vs pinned reference martin-marek/score-centering@fdca22b
(README JAX score_centering_loss and losses.loss_fn_rl) on a fixed synthetic full-vocab fixture, float32."""
import sys, json, hashlib, re
import numpy as np, torch, jax, jax.numpy as jnp
jax.config.update("jax_enable_x64", False)
sys.path.insert(0, "/tmp/sc-ref")
import losses as ref_losses
sys.path.insert(0, "/home/ubuntu/repos/research-exchange")
from research_exchange.correctness import parity_errors
import rl_score_centering as rl

readme = open("/tmp/sc-ref/README.md").read()
code = re.search(r"```python\n(import jax\n.*?)```", readme, re.S).group(1)
ns = {}; exec(code.split("# logits: [B, T, V]")[0], ns)
score_centering_loss = ns["score_centering_loss"]; tis_weight = ns["tis_weight"]

rng = np.random.default_rng(0)
B, T, V, K = 4, 32, 4096, rl.K
logits = (rng.standard_normal((B, T, V)) * 3).astype(np.float32)
samp_logits = (logits + rng.standard_normal((B, T, V)) * 0.3).astype(np.float32)   # train/inference mismatch
samp_logp = jax.nn.log_softmax(jnp.array(samp_logits), -1)
topk_ids = np.array(jnp.argsort(-samp_logp, -1)[..., :K])
samp_head = np.take_along_axis(np.array(samp_logp), topk_ids, -1)
sampled = np.array(jax.random.categorical(jax.random.PRNGKey(1), jnp.array(samp_logits)))
samp_tok = np.take_along_axis(np.array(samp_logp), sampled[..., None], -1)[..., 0]
adv = np.array([1.0, -0.5, 0.25, -1.0], np.float32)
mask = np.ones((B, T), np.float32); mask[:, :4] = 0
fixture = dict(logits=logits, samp_logits=samp_logits, sampled=sampled, adv=adv, mask=mask)
fsha = hashlib.sha256(b"".join(np.ascontiguousarray(v).tobytes() for v in fixture.values())).hexdigest()
np.savez("ref_parity_fixture.npz", **fixture)

def ref_readme(lg, wf):
    tl = score_centering_loss(jax.nn.log_softmax(lg.astype(jnp.float32), -1), jnp.array(samp_head), jnp.array(topk_ids),
                              jnp.array(sampled), jnp.array(samp_tok), jnp.array(adv)[:, None], weight_fn=wf)
    return (tl * mask).sum() / mask.sum()

def ref_repo(lg, corr):
    # loss_fn_rl rolls targets by -1; put sampled token at t+1 so targets[t]=sampled[t].
    tokens = np.zeros((B, T), np.int32); tokens[:, 1:] = sampled[:, :-1]
    tokens = np.concatenate([tokens, sampled[:, -1:]], 1)            # length T+1
    m = np.zeros((B, T + 1), bool); m[:, 1:] = mask > 0
    pad = lambda a: np.concatenate([a, a[:, -1:]], 1)
    batch = {"tokens": jnp.array(tokens), "mask": jnp.array(m), "advantages": jnp.array(adv),
             "sampling_token_logprobs": jnp.array(pad(samp_tok)), "center_ids": jnp.array(pad(topk_ids)),
             "center_logprobs": jnp.array(pad(samp_head))}
    hidden = jnp.concatenate([lg, lg[:, -1:]], 1)
    loss, _ = ref_losses.loss_fn_rl(hidden, lambda h: h, batch, correction=corr, score_center=True, chunk=T + 1)
    return loss

def mine(wname):
    lg = torch.tensor(logits, requires_grad=True)
    lsm = torch.log_softmax(lg.float(), -1)
    ids = torch.cat([torch.tensor(sampled)[..., None], torch.tensor(topk_ids)], -1).long()
    total = 0.0
    for b in range(B):
        lp = torch.gather(lsm[b], -1, ids[b])
        hm = torch.tensor(mask[b])[:, None].expand(T, K).clone()          # as build_datums: no top-k on non-completion rows
        sh = torch.where(hm > 0, torch.tensor(samp_head[b]), torch.full((T, K), -1e9))
        d = {"samp_lp": torch.tensor(samp_tok[b]) * torch.tensor(mask[b]), "samp_head": sh,
             "head_mask": hm, "mask": torch.tensor(mask[b]), "adv": torch.full((T,), float(adv[b]))}
        tok_logp, tw, corr, _ = rl.sc_terms(lp, d, wname)
        tok = -d["adv"] * (tw.detach() * tok_logp - corr)
        total = total + (tok * d["mask"]).sum()
    loss = total / mask.sum()
    loss.backward()
    return loss.item(), lg.grad.numpy().ravel()

res = {"fixture_sha256": fsha, "fixture": f"ref_parity_fixture.npz B={B} T={T} V={V} K={K} seed0", "checks": {}}
for wname, wf, corr in [("identity", lambda r: r, ref_losses.ISCorrection()), ("ones", jnp.ones_like, None)]:
    ml, mg = mine(wname)
    rl_, rg = jax.value_and_grad(lambda x: ref_readme(x, wf))(jnp.array(logits))
    pl, pg = jax.value_and_grad(lambda x: ref_repo(x, corr))(jnp.array(logits))
    pg = np.array(pg)
    res["checks"][wname] = {
        "reference_loss_vs_readme": parity_errors([float(rl_)], [ml]),
        "reference_gradient_vs_readme": parity_errors(np.array(rg).ravel().tolist(), mg.tolist()),
        "loss_fn_rl_loss": float(pl), "readme_loss": float(rl_), "mine_loss": ml,
        "reference_gradient_vs_loss_fn_rl": parity_errors(np.array(pg).ravel().tolist(), mg.tolist())}
print(json.dumps(res, indent=1))
json.dump(res, open("ref_parity_result_v2.json", "w"), indent=1)
assert all(np.isfinite(mine("identity")[1]))