"""Replay pinned losses and mask/gradient probes without calling a training service."""

import argparse
import ast
import hashlib
import importlib.util
import json
import types
from pathlib import Path
from typing import Any, Optional
import torch

parser = argparse.ArgumentParser(
    description="Offline replay of run_19badd842c evidence; requires CPU torch."
)
parser.add_argument("evidence_dir", type=Path)
parser.add_argument("--patched-losses", type=Path)
args = parser.parse_args()
P = args.evidence_dir
# Verify saved evidence before executing any experiment code.
manifest = json.loads((P / "manifest.json").read_text())
for entry in manifest["files"]:
    content = (P / entry["name"]).read_bytes()
    assert len(content) == entry["size"]
    assert hashlib.sha256(content).hexdigest() == entry["sha256"]
for name, expected in {
    "losses.py": "36e744853b410170a874b80e523914ec5b0f04fa7781d96b9d04abd0e10271d3",
    "preflight_fixture.json": "aff0d4c84ff104219bf365a13dcb0de1c1cfc59b833ac72cb357c46e29a83f13",
    "core_algos.py": "4f3738abecebbe4f6e1202103fa57deeb2b9ea38b14653152b38f04fe8db7684",
    "torch_functional.py": "8ac62195a800ba784d1cddf03ee0bc7627f20236cc11e135fdba20e8d4a01f14",
    "forward_backward.py": "a71be71348c428c95c6439c4152213e046b460cbed6adb787290ffae5311b60b",
}.items():
    assert hashlib.sha256((P / name).read_bytes()).hexdigest() == expected, name
losses_path = args.patched_losses or P / "losses.py"
spec = importlib.util.spec_from_file_location("saved_losses", losses_path)
lossmod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(lossmod)


def extract(file, names, ns):
    for node in ast.parse((P / file).read_text()).body:
        if isinstance(node, ast.FunctionDef) and node.name in names:
            node.decorator_list = []
            exec(compile(ast.Module([node], []), str(P / file), "exec"), ns)
    return ns


vf = extract("torch_functional.py", {"masked_sum", "masked_mean"}, {"torch": torch})
ns = extract(
    "core_algos.py",
    {
        "agg_loss",
        "compute_policy_loss_with_mask",
        "_compute_token_mask_for_policy_loss",
    },
    {
        "torch": torch,
        "Any": Any,
        "Optional": Optional,
        "ActorConfig": object,
        "verl_F": types.SimpleNamespace(
            **{k: vf[k] for k in ("masked_sum", "masked_mean")}
        ),
    },
)
sp = extract(
    "forward_backward.py", {"_loss"}, {"torch": torch, "Any": Any, "RL_LOSSES": {"ppo"}}
)


class Cfg(dict):
    def __getattr__(self, k):
        return self[k]


def cfg(delta):
    return Cfg(
        clip_ratio=delta,
        clip_ratio_low=delta,
        clip_ratio_high=delta,
        clip_ratio_c=10,
        policy_loss={"loss_mode": "dppo_binary_tv"},
        global_batch_info={"dp_size": 1, "batch_num_tokens": 1},
    )


def ref(cur, old, adv, mask, delta=0.15):
    return ns["compute_policy_loss_with_mask"](
        old,
        cur,
        adv,
        mask,
        config=cfg(delta),
        rollout_log_prob=old,
        loss_mode="dppo_binary_tv",
        topk_kl=torch.zeros_like(cur),
        entropy=torch.zeros_like(cur),
    )


def maskref(cur, old, adv, mask, delta=0.15):
    return ns["_compute_token_mask_for_policy_loss"](
        old, old, cur, adv, mask, config=cfg(delta), loss_mode="dppo_binary_tv"
    )


def errors(a, b):
    return {
        "max_abs": float((a - b).detach().abs().max()),
        "exact": bool(torch.equal(a, b)),
    }


f = json.loads((P / "preflight_fixture.json").read_text())
t = {k: torch.tensor(f[k]) for k in ("old", "cur", "adv", "mask")}


class TD:
    def __init__(self, t):
        self.t = t

    def to_torch(self):
        return self.t


D = [
    types.SimpleNamespace(
        loss_fn_inputs={
            k: TD(t[src][i, :n])
            for k, src in [("logprobs", "old"), ("advantages", "adv"), ("mask", "mask")]
        }
    )
    for i, n in enumerate(f["lens"])
]
result = {
    "torch": torch.__version__,
    "fixture_sha256": hashlib.sha256(
        (P / "preflight_fixture.json").read_bytes()
    ).hexdigest(),
    "fixture": {},
}
for mode in ("grpo", "dppo"):
    lps = [t["cur"][i, :n].clone().requires_grad_() for i, n in enumerate(f["lens"])]
    actual, diag = lossmod.make_loss(mode)(D, lps)
    actual.backward()
    grad = torch.zeros_like(t["cur"])
    for i, n in enumerate(f["lens"]):
        grad[i, :n] = lps[i].grad
    cur = t["cur"].clone().requires_grad_()
    if mode == "dppo":
        expected, metrics = ref(cur, t["old"], t["adv"], t["mask"])
    else:
        expected = sum(
            sp["_loss"](
                cur[i, :n],
                {
                    "loss_mask": t["mask"][i, :n],
                    "loss_name": "ppo",
                    "sampling_logprobs": t["old"][i, :n],
                    "advantages": t["adv"][i, :n],
                    "loss_config": {
                        "clip_low_threshold": 0.8,
                        "clip_high_threshold": 1.2,
                    },
                },
            )[0]
            for i, n in enumerate(f["lens"])
        )
    expected.backward()
    result["fixture"][mode] = {
        "loss_actual": float(actual.detach()),
        "loss_reference": float(expected.detach()),
        "loss_error": errors(actual, expected),
        "gradient_error": errors(grad, cur.grad),
        "diagnostics": diag,
    }
    torch.testing.assert_close(grad, cur.grad, rtol=0.001, atol=1e-5)
    torch.testing.assert_close(actual, expected, rtol=0.0001, atol=1e-5)
keep = maskref(t["cur"], t["old"], t["adv"], t["mask"])
ownkeep = (
    1
    - torch.where(
        t["adv"] > 0,
        t["cur"].exp() - t["old"].exp() > 0.15,
        t["cur"].exp() - t["old"].exp() < -0.15,
    ).float()
) * t["mask"]
assert torch.equal(keep, ownkeep)
result["fixture"]["mask_exact"] = True
# low-probability tails, both signs, recovery direction, zero advantage, padding and ratio cap
cases = [
    ("positive_block", 0.5, 0.75, 1, 1, True),
    ("negative_block", 0.75, 0.5, -1, 1, True),
    ("positive_recovery", 0.75, 0.5, 1, 1, False),
    ("negative_recovery", 0.5, 0.75, -1, 1, False),
    ("positive_small", 0.5, 0.6, 1, 1, False),
    ("negative_small", 0.6, 0.5, -1, 1, False),
    ("tail_positive", 1e-8, 1e-4, 1, 1, False),
    ("tail_negative", 1e-4, 1e-8, -1, 1, False),
    ("cap_positive", 1e-8, 0.1, 1, 1, False),
    ("zero_adv", 0.75, 0.5, 0, 1, False),
    ("padding", 0.5, 0.75, 1, 0, False),
]
old = torch.tensor([c[1] for c in cases], dtype=torch.float64).log()
cur = torch.tensor([c[2] for c in cases], dtype=torch.float64).log().requires_grad_()
adv = torch.tensor([c[3] for c in cases])
mask = torch.tensor([c[4] for c in cases])
own = lossmod.token_terms(cur, old, adv, mask, "dppo").sum()
own.backward()
g = cur.grad.clone()
cur.grad = None
reference, _ = ref(cur, old, adv, mask)
reference.backward()
torch.testing.assert_close(g, cur.grad, rtol=0, atol=0)
expected_g = torch.tensor(
    [0 if c[5] or not c[4] else -c[3] * min(c[2] / c[1], 10) for c in cases],
    dtype=torch.float64,
)
torch.testing.assert_close(g, expected_g, rtol=1e-12, atol=1e-12)
result["deterministic"] = [
    {
        "case": c[0],
        "gradient": float(g[i]),
        "expected_gradient": float(expected_g[i]),
        "reference_keep": float(maskref(cur, old, adv, mask)[i]),
    }
    for i, c in enumerate(cases)
]
# Independently check the PPO clip diagnostic on the same forced cases.
cur.grad = None
ppo_loss = lossmod.token_terms(cur, old, adv, mask, "grpo").sum()
ppo_grad = torch.autograd.grad(ppo_loss, cur)[0]
stock, _ = sp["_loss"](
    cur,
    {
        "loss_mask": mask,
        "loss_name": "ppo",
        "sampling_logprobs": old,
        "advantages": adv,
        "loss_config": {"clip_low_threshold": 0.8, "clip_high_threshold": 1.2},
    },
)
stock_grad = torch.autograd.grad(stock, cur)[0]
torch.testing.assert_close(ppo_grad, stock_grad, rtol=0, atol=0)
expected_clips = sum(
    c[4] and ((c[3] > 0 and c[2] / c[1] > 1.2) or (c[3] < 0 and c[2] / c[1] < 0.8))
    for c in cases
)
clip_diagnostic = lossmod.diagnostics(cur, old, adv, mask)["grpo_clip_fraction"]
assert abs(clip_diagnostic - expected_clips / sum(c[4] for c in cases)) < 1e-12
result["deterministic_ppo_gradient_and_clip_diagnostic"] = True
# Probe just inside/outside the accepted delta for each advantage sign.
for sign in (1, -1):
    for displacement, blocked in ((0.15 - 1e-6, False), (0.15 + 1e-6, True)):
        old = torch.tensor([0.5], dtype=torch.float64).log()
        cur = (
            torch.tensor([0.5 + sign * displacement], dtype=torch.float64)
            .log()
            .requires_grad_()
        )
        adv, mask = torch.tensor([sign]), torch.ones(1)
        expected_keep = 0 if blocked else 1
        assert float(maskref(cur, old, adv, mask)[0]) == expected_keep
        gradient = torch.autograd.grad(
            lossmod.token_terms(cur, old, adv, mask, "dppo").sum(), cur
        )[0]
        assert (float(gradient) == 0) == blocked
result["default_delta_both_sides_both_signs"] = True
# Strict boundary: binary-exact delta avoids conflating log/exp rounding with > vs >=.
for sign in (1, -1):
    old = torch.tensor([0.375 if sign > 0 else 0.5], dtype=torch.float64).log()
    cur = (
        torch.tensor([0.5 if sign > 0 else 0.375], dtype=torch.float64)
        .log()
        .requires_grad_()
    )
    adv = torch.tensor([sign])
    mask = torch.ones(1)
    assert float(maskref(cur, old, adv, mask, 0.125)[0]) == 1
    g = torch.autograd.grad(
        lossmod.token_terms(cur, old, adv, mask, "dppo", delta=0.125).sum(), cur
    )[0]
    assert float(g) != 0
result["strict_boundary_both_signs"] = True
# Instrumentation counts zero-advantage negative shifts as masks, though they remove no gradient.
result["zero_adv_diagnostic"] = lossmod.diagnostics(
    torch.tensor([0.5]).log(), torch.tensor([0.75]).log(), torch.zeros(1), torch.ones(1)
)
rows = [json.loads(s) for s in (P / "pilot-dppo-steps.jsonl").read_text().splitlines()]
result["pilot"] = {
    "tokens": sum(m["n_loss_tokens"] for r in rows for m in r["mb"]),
    "per_update": [],
    "raw_token_tails_available": False,
}
for i in (0, 1):
    ms = [r["mb"][i] for r in rows]
    n = sum(m["n_loss_tokens"] for m in ms)
    result["pilot"]["per_update"].append(
        {
            "update": i + 1,
            "tokens": n,
            "masked_tokens": sum(
                m["n_loss_tokens"] * m["dppo_masked_token_fraction"] for m in ms
            ),
            "ppo_clipped_tokens": round(
                sum(m["n_loss_tokens"] * m["grpo_clip_fraction"] for m in ms)
            ),
            "q90_range": [
                min(m["policy_drift_abs_p_q90"] for m in ms),
                max(m["policy_drift_abs_p_q90"] for m in ms),
            ],
            "token_weighted_abs_drift_mean": sum(
                m["n_loss_tokens"] * m["policy_drift_abs_p_mean"] for m in ms
            )
            / n,
        }
    )
(
    P / ("audit-result-patched.json" if args.patched_losses else "audit-result.json")
).write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))

if args.patched_losses:
    assert result["zero_adv_diagnostic"]["dppo_masked_token_fraction"] == 0
    assert result["zero_adv_diagnostic"]["dppo_removed_gradient_tokens"] == 0