"""Correctness preflight (local, fp32): loss + dL/dlogp parity of code/losses.py vs pinned sources.
baseline: spindle@e58ebef _loss('ppo') executed from source; reference: Stable-RL@c6af1d1
compute_policy_loss_with_mask(loss_mode='dppo_binary_tv') executed from source (agg_loss with
batch_num_tokens=1, dp_size=1 => masked token SUM, matching Tinker's sum convention)."""
import ast, hashlib, json, sys, types, functools
import numpy as np, torch, tinker
sys.path.insert(0, "/home/ubuntu/rx/run_19badd842c/code")
sys.path.insert(0, "/home/ubuntu/repos/research-exchange")
from losses import make_loss, GRPO_EPS, DPPO_DELTA
from research_exchange.correctness import parity_errors

ROOT = "/home/ubuntu/rx/run_19badd842c"

def extract(path, names, ns):
    src = open(path).read(); tree = ast.parse(src)
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.Assign)) and getattr(node, "name", None) in names:
            exec(compile(ast.Module([node], []), path, "exec"), ns)
    return ns

# --- pinned baseline (spindle stock ppo)
sp_ns = extract(f"{ROOT}/spindle/src/spindle/backends/megatron_runtime/common/forward_backward.py", {"_loss"},
                {"torch": torch, "Any": object, "RL_LOSSES": {"importance_sampling", "ppo", "cispo", "dppo", "dro"}})
# --- pinned reference (Stable-RL dppo_binary_tv)
vF_ns = extract(f"{ROOT}/Stable-RL/verl/utils/torch_functional.py", {"masked_sum", "masked_mean"}, {"torch": torch})
verl_F = types.SimpleNamespace(**{k: vF_ns[k] for k in ("masked_sum", "masked_mean")})
from typing import Any, Optional
ref_ns = extract(f"{ROOT}/Stable-RL/verl/trainer/ppo/core_algos.py",
                 {"agg_loss", "compute_policy_loss_with_mask", "_compute_token_mask_for_policy_loss"},
                 {"torch": torch, "verl_F": verl_F, "Any": Any, "Optional": Optional, "ActorConfig": object})

class Cfg(dict):
    def __getattr__(self, k): return self[k]
cfg = Cfg(clip_ratio=DPPO_DELTA, clip_ratio_low=DPPO_DELTA, clip_ratio_high=DPPO_DELTA, clip_ratio_c=10.0,
          policy_loss={"loss_mode": "dppo_binary_tv"}, global_batch_info={"dp_size": 1, "batch_num_tokens": 1})

# --- fixed stale fixture: 16 seqs, ragged lengths, prompt tokens masked, drift so clips/masks fire
g = torch.Generator().manual_seed(0)
B, Tmax = 16, 96
lens = torch.randint(40, Tmax + 1, (B,), generator=g)
datums, pad = [], {k: torch.zeros(B, Tmax) for k in ("old", "cur", "adv", "mask")}
for i in range(B):
    T = int(lens[i]); P = int(torch.randint(5, 20, (1,), generator=g))
    old = torch.log(torch.rand(T, generator=g).clamp(1e-4, 1.0)) * 1.5
    cur = old + torch.randn(T, generator=g) * 0.6
    cur = torch.minimum(cur, torch.zeros(T) - 1e-6)
    mask = torch.zeros(T); mask[P:] = 1
    adv = torch.full((T,), float(torch.randn(1, generator=g))) * mask
    old = old * mask
    if i == 3: adv = adv * 0  # zero-advantage group
    datums.append(tinker.Datum(model_input=tinker.ModelInput.from_ints(list(range(T))), loss_fn_inputs={
        "target_tokens": tinker.TensorData.from_torch(torch.arange(T)), "logprobs": tinker.TensorData.from_torch(old),
        "advantages": tinker.TensorData.from_torch(adv), "mask": tinker.TensorData.from_torch(mask)}))
    for k, v in zip(("old", "cur", "adv", "mask"), (old, cur, adv, mask)): pad[k][i, :T] = v
fixture = {k: v.tolist() for k, v in pad.items()} | {"lens": lens.tolist()}
fixture_bytes = json.dumps(fixture, sort_keys=True).encode()
open(f"{ROOT}/preflight_fixture.json", "wb").write(fixture_bytes)

def ours(mode):
    lps = [pad["cur"][i, :int(lens[i])].clone().requires_grad_(True) for i in range(B)]
    loss, metrics = make_loss(mode)(datums, lps); loss.backward()
    grad = torch.zeros(B, Tmax)
    for i in range(B): grad[i, :int(lens[i])] = lps[i].grad
    return loss.detach(), grad, metrics

def ref_baseline():
    out_l, out_g = [], torch.zeros(B, Tmax)
    for i in range(B):  # spindle flattens per packed row; run per sequence and sum
        T = int(lens[i]); lp = pad["cur"][i, :T].clone().requires_grad_(True)
        loss, _ = sp_ns["_loss"](lp, {"loss_mask": pad["mask"][i, :T], "loss_name": "ppo",
            "sampling_logprobs": pad["old"][i, :T], "advantages": pad["adv"][i, :T],
            "loss_config": {"clip_low_threshold": 1 - GRPO_EPS, "clip_high_threshold": 1 + GRPO_EPS}})
        loss.backward(); out_l.append(loss.detach()); out_g[i, :T] = lp.grad
    return torch.stack(out_l).sum(), out_g

def ref_reference():
    lp = pad["cur"].clone().requires_grad_(True)  # padded [B, Tmax] layout
    loss, m = ref_ns["compute_policy_loss_with_mask"](old_log_prob=pad["old"], log_prob=lp, advantages=pad["adv"],
        response_mask=pad["mask"], loss_agg_mode="token-mean", config=cfg, rollout_log_prob=pad["old"],
        loss_mode="dppo_binary_tv", topk_kl=torch.zeros(B, Tmax), entropy=torch.zeros(B, Tmax))
    loss.backward()
    return loss.detach(), lp.grad * pad["mask"], m

res = {}
lb, gb, mb = ours("grpo"); rlb, rgb = ref_baseline()
res["baseline_loss"] = parity_errors([float(rlb)], [float(lb)])
res["baseline_gradient"] = parity_errors(rgb.flatten().tolist(), gb.flatten().tolist())
ld, gd, md = ours("dppo"); rld, rgd, rm = ref_reference()
res["reference_loss"] = parity_errors([float(rld)], [float(ld)])
res["reference_gradient"] = parity_errors(rgd.flatten().tolist(), gd.flatten().tolist())
out = {"fixture_sha256": hashlib.sha256(fixture_bytes).hexdigest(), "errors": res,
       "ours_metrics_grpo": mb, "ours_metrics_dppo": md, "ref_pg_clipfrac": rm["actor/pg_clipfrac"],
       "losses": {"grpo": float(lb), "spindle_ppo": float(rlb), "dppo": float(ld), "stable_rl": float(rld)}}
print(json.dumps(out, indent=1))
json.dump(out, open(f"{ROOT}/preflight_local_result.json", "w"), indent=1)