"""CISPO objective (MiniMax-M1 arXiv:2506.13585 Eq. 4-5) + FP32 LM head.

Every oracle here is an independent fp64 implementation written IN the test
(explicit loops / hand-derived softmax gradient), never a call back into the
function under test. The properties that define CISPO -- a clipped token still
gets a gradient; the weight is a stop-gradient constant; the default GRPO loss
is untouched -- each have a test that fails on the corresponding mutant
(scripts in .tmp/g2/p4a/mutants.py run them).
"""

from __future__ import annotations

import math

import pytest
import torch

from align.rl.cispo import (
    P4CispoConformanceError,
    assert_p4_cispo_conformant,
    cispo_policy_loss,
    fp32_head_forward,
    fp32_head_logprobs,
)
from align.rl.grpo_loss import (
    A3ConformanceError,
    GRPOLossConfig,
    GRPOLossType,
    assert_a3_conformant,
    compute_grpo_loss,
)
from align.rl.grpo_module import GRPOTrainModule
from align.rl.types import RolloutBatch

D = torch.float64


# ------------------------------------------------------------------ oracles


def ref_cispo_grad_logits(logits, labels, adv, old_lp, mask, eps_high, eps_low, denom):
    """Hand-derived fp64 gradient of  -(1/denom) sum_t mask*A*w_t*log pi(y_t)
    wrt the logits, with w_t = clip(exp(logp_t - old_t)) a CONSTANT (Eq. 4-5).
    d logp_y / d logits = onehot(y) - softmax(logits)."""
    logits = logits.to(D)
    B, T, V = logits.shape
    grad = torch.zeros(B, T, V, dtype=D)
    loss = 0.0
    for b in range(B):
        for t in range(T):
            z = logits[b, t]
            m = z.max()
            p = torch.exp(z - m)
            p = p / p.sum()
            logp = torch.log(p[labels[b, t]])
            r = math.exp(float(logp) - float(old_lp[b, t]))
            w = min(r, 1 + eps_high)
            if eps_low is not None:
                w = max(w, 1 - eps_low)
            if not mask[b, t]:
                continue
            loss += -float(adv[b]) * w * float(logp) / denom
            onehot = torch.zeros(V, dtype=D)
            onehot[labels[b, t]] = 1.0
            grad[b, t] = -float(adv[b]) * w * (onehot - p) / denom
    return loss, grad


def _setup(B=3, T=5, V=11, seed=0):
    g = torch.Generator().manual_seed(seed)
    logits = torch.randn(B, T, V, generator=g, dtype=D)
    labels = torch.randint(0, V, (B, T), generator=g)
    adv = torch.tensor([1.0, -0.7, 0.4][:B], dtype=D)
    mask = torch.ones(B, T, dtype=torch.bool)
    mask[:, 0] = False  # a non-response position
    return logits, labels, adv, mask


def _old_from_ratios(logits, labels, ratios):
    """old = logp - log(ratio) so that exp(logp-old) == ratios exactly."""
    logp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)
    return (logp - torch.log(ratios)).detach()


def _cispo_cfg(**kw):
    base = dict(loss_fn=GRPOLossType.cispo, clip_higher=0.272, beta=0.0, load_ref_policy=False,
                loss_denominator=64.0, num_samples_per_prompt_rollout=8)
    base.update(kw)
    return GRPOLossConfig(**base)


def _run_loss(logits, labels, adv, old, mask, cfg, loss_fn=None):
    logits = logits.clone().requires_grad_(True)
    lp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)
    ratio = torch.exp(lp - old)
    a = adv[:, None].expand_as(lp)
    pg, clipfrac, _ = compute_grpo_loss(lp, ratio, a, None, cfg, torch.ones_like(lp))
    loss = (pg * mask).sum() / cfg.loss_denominator
    loss.backward()
    return loss.detach(), logits.grad, clipfrac


# ------------------------------------------------- 1. gradient == fp64 formula


@pytest.mark.parametrize("eps_low", [None, 0.5])
def test_cispo_gradient_matches_fp64_reference(eps_low):
    logits, labels, adv, mask = _setup()
    B, T, _ = logits.shape
    # ratios spanning below-floor, inside, above-cap
    ratios = torch.tensor([[0.2, 0.9, 1.0, 1.1, 3.0], [0.1, 0.6, 1.2, 1.3, 2.0], [0.95, 1.0, 5.0, 0.3, 1.25]],
                          dtype=D)
    old = _old_from_ratios(logits, labels, ratios)
    cfg = _cispo_cfg(cispo_clip_lower=eps_low)
    loss, grad, _ = _run_loss(logits, labels, adv, old, mask, cfg)
    ref_loss, ref_grad = ref_cispo_grad_logits(logits, labels, adv, old, mask, 0.272, eps_low, 64.0)
    assert abs(float(loss) - ref_loss) < 1e-12
    torch.testing.assert_close(grad, ref_grad, rtol=1e-10, atol=1e-12)


def test_cispo_matches_explicit_eq4_loop_value():
    """Eq. 4 with sum over (i,t) written as a plain python double loop."""
    logits, labels, adv, mask = _setup(seed=3)
    ratios = torch.rand(3, 5, dtype=D) * 3 + 0.05
    old = _old_from_ratios(logits, labels, ratios)
    lp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)
    pg, _ = cispo_policy_loss(lp, torch.exp(lp - old), adv[:, None].expand_as(lp), 0.272)
    total = 0.0
    for i in range(3):
        for t in range(5):
            if mask[i, t]:
                w = min(float(torch.exp(lp[i, t] - old[i, t])), 1.272)
                total += -float(adv[i]) * w * float(lp[i, t])
    assert abs(float((pg * mask).sum()) - total) < 1e-12


# ----------------------------- 2. clipped tokens still learn; PPO-clip zeroes


def test_clipped_tokens_keep_gradient_ppo_clip_zeroes_them():
    """The defining difference. One sequence, A>0, ratio >> 1+eps_high on every
    token: DAPO/PPO-clip gradient is EXACTLY zero, CISPO gradient is
    A*(1+eps_high)*(onehot-softmax) per token."""
    logits, labels, adv, mask = _setup(B=1, T=6)
    adv = torch.tensor([1.0], dtype=D)
    mask = torch.ones(1, 6, dtype=torch.bool)
    ratios = torch.full((1, 6), 4.0, dtype=D)  # far above 1.272
    old = _old_from_ratios(logits, labels, ratios)

    _, g_dapo, cf_dapo = _run_loss(logits, labels, adv, old, mask, GRPOLossConfig(loss_fn=GRPOLossType.dapo,
                                                                                  loss_denominator=6.0))
    assert torch.count_nonzero(g_dapo) == 0 and bool(cf_dapo.all())

    cfg = _cispo_cfg(loss_denominator=6.0)
    _, g_cispo, cf = _run_loss(logits, labels, adv, old, mask, cfg)
    assert bool(cf.all())  # every token was clipped ...
    p = logits.softmax(-1)
    oh = torch.nn.functional.one_hot(labels, logits.shape[-1]).to(D)
    expected = -1.0 * 1.272 * (oh - p) / 6.0  # ... and still carries full-strength gradient
    torch.testing.assert_close(g_cispo, expected, rtol=1e-10, atol=1e-13)


def test_low_side_negative_advantage_also_keeps_gradient():
    """A<0 with ratio << 1-eps_low: PPO-clip zeroes it (min-side clip). CISPO
    with no lower bound keeps weight = r (small, nonzero); with a floor it keeps
    weight = 1-eps_low. Neither is zero."""
    logits, labels, adv, mask = _setup(B=1, T=4)
    adv = torch.tensor([-1.0], dtype=D)
    mask = torch.ones(1, 4, dtype=torch.bool)
    ratios = torch.full((1, 4), 0.05, dtype=D)
    old = _old_from_ratios(logits, labels, ratios)
    _, g_dapo, _ = _run_loss(logits, labels, adv, old, mask,
                             GRPOLossConfig(loss_fn=GRPOLossType.dapo, loss_denominator=4.0))
    assert torch.count_nonzero(g_dapo) == 0
    for eps_low, w in ((None, 0.05), (0.5, 0.5)):
        _, g, _ = _run_loss(logits, labels, adv, old, mask, _cispo_cfg(cispo_clip_lower=eps_low,
                                                                       loss_denominator=4.0))
        oh = torch.nn.functional.one_hot(labels, logits.shape[-1]).to(D)
        expected = 1.0 * w * (oh - logits.softmax(-1)) / 4.0
        torch.testing.assert_close(g, expected, rtol=1e-10, atol=1e-13)


def test_weight_is_stop_gradient():
    """If `ratio` is not detached, d loss/d ratio-source is nonzero (an extra
    A*log pi*grad r term). Assert the gradient flows ONLY through log pi."""
    lp = torch.tensor([[-1.0, -2.0]], dtype=D, requires_grad=True)
    ratio_src = torch.tensor([[1.1, 0.9]], dtype=D, requires_grad=True)
    pg, _ = cispo_policy_loss(lp, ratio_src * 1.0, torch.ones(1, 2, dtype=D), 0.272)
    pg.sum().backward()
    assert ratio_src.grad is None or torch.count_nonzero(ratio_src.grad) == 0
    torch.testing.assert_close(lp.grad, -torch.tensor([[1.1, 0.9]], dtype=D))


def test_eps_cap_changes_weight_exactly_at_boundary():
    lp = torch.zeros(1, 3, dtype=D)
    ratio = torch.tensor([[1.271, 1.272, 1.2721]], dtype=D)
    a = torch.ones(1, 3, dtype=D)
    lp1 = lp - 1.0
    pg, cf = cispo_policy_loss(lp1, ratio, a, 0.272)
    torch.testing.assert_close(pg, torch.tensor([[1.271, 1.272, 1.272]], dtype=D))
    assert cf.tolist() == [[0.0, 0.0, 1.0]]


def test_clipfrac_counts_both_sides_when_floor_set():
    lp = torch.full((1, 4), -1.0, dtype=D)
    ratio = torch.tensor([[0.1, 0.6, 1.0, 2.0]], dtype=D)
    _, cf = cispo_policy_loss(lp, ratio, torch.ones(1, 4, dtype=D), 0.272, eps_low=0.5)
    assert cf.tolist() == [[1.0, 0.0, 0.0, 1.0]]
    _, cf = cispo_policy_loss(lp, ratio, torch.ones(1, 4, dtype=D), 0.272)  # paper: no floor
    assert cf.tolist() == [[0.0, 0.0, 0.0, 1.0]]


def test_invalid_eps_refused():
    x = torch.zeros(1, 1)
    for kw in ({"eps_high": 0.0}, {"eps_high": float("nan")}, {"eps_high": 0.2, "eps_low": 0.0},
               {"eps_high": 0.2, "eps_low": float("inf")}):
        with pytest.raises(ValueError):
            cispo_policy_loss(x, x + 1, x, **kw)


# --------------------------------------------- 3. default GRPO byte-identical


def _legacy_compute_grpo_loss(new_logprobs, ratio, advantages, config):
    """VERBATIM pre-change bodies of the dapo and cispo branches (HEAD 780ce3d)."""
    if config.loss_fn == GRPOLossType.dapo:
        pg_losses = -advantages * ratio
        pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - config.clip_lower, 1.0 + config.clip_higher)
        clipfrac = (pg_losses2 > pg_losses).float()
        pg_loss = torch.max(pg_losses, pg_losses2)
    else:
        clipfrac = (ratio > 1.0 + config.clip_higher).float()
        pg_loss = -advantages * torch.clamp(ratio.detach(), max=1.0 + config.clip_higher) * new_logprobs
    return pg_loss, clipfrac


@pytest.mark.parametrize("fn", [GRPOLossType.dapo, GRPOLossType.cispo])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_default_loss_byte_identical_to_pre_change(fn, dtype):
    g = torch.Generator().manual_seed(7)
    lp = (-torch.rand(4, 9, generator=g) * 5).to(dtype)
    ratio = (torch.rand(4, 9, generator=g) * 3).to(dtype)
    adv = torch.randn(4, 1, generator=g).to(dtype).expand(4, 9)
    cfg = GRPOLossConfig(loss_fn=fn)
    assert cfg.cispo_clip_lower is None and cfg.fp32_head is False  # opt-ins default off
    pg, cf, _ = compute_grpo_loss(lp, ratio, adv, None, cfg, torch.ones_like(lp))
    rpg, rcf = _legacy_compute_grpo_loss(lp, ratio, adv, cfg)
    assert torch.equal(pg, rpg) and torch.equal(cf, rcf)
    assert pg.dtype == rpg.dtype


def test_a3_gate_still_refuses_bare_cispo_and_p4_gate_accepts():
    base = dict(beta=0.0, load_ref_policy=False, loss_denominator=1024.0,
                num_samples_per_prompt_rollout=16, record_entropy=True)
    with pytest.raises(A3ConformanceError):
        assert_a3_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, **base))
    assert_p4_cispo_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, **base))
    assert_a3_conformant(GRPOLossConfig(**base))  # dapo default untouched


@pytest.mark.parametrize(
    "bad",
    [
        {"loss_denominator": "token"},
        {"beta": 0.05},
        {"load_ref_policy": True},
        {"num_samples_per_prompt_rollout": 4},
        {"use_rho_correction": False},
        {"filter_zero_std_samples": False},
        {"advantage_normalization_type": "standard"},
        {"clip_higher": 0.0},
        {"clip_higher": float("inf")},
        {"cispo_clip_lower": 0.0},
        {"cispo_clip_lower": -1.0},
        {"fp32_head_chunk_tokens": 0},
        {"loss_fn": GRPOLossType.dapo},
    ],
)
def test_p4_gate_refuses_each_violation(bad):
    cfg = dict(loss_fn=GRPOLossType.cispo, beta=0.0, load_ref_policy=False, loss_denominator=1024.0,
               num_samples_per_prompt_rollout=16, clip_higher=2.0, record_entropy=True)
    cfg.update(bad)
    with pytest.raises(P4CispoConformanceError):
        assert_p4_cispo_conformant(GRPOLossConfig(**cfg))


# ------------------------------------------------------------ 4. FP32 head


def _hw(seed=0, B=2, T=7, H=16, V=40, scale=12.0, dtype=torch.bfloat16):
    g = torch.Generator().manual_seed(seed)
    h = (torch.randn(B, T, H, generator=g) * scale).to(dtype)
    W = (torch.randn(V, H, generator=g) * 0.5).to(dtype)
    labels = torch.randint(0, V, (B, T), generator=g)
    return h, W, labels


def ref_head_logprobs(h, W, labels, temperature=1.0):
    """fp64 oracle on the bf16-VALUED operands (exact upcast, no rounding)."""
    z = (h.to(D) @ W.to(D).t()) / temperature
    out = torch.empty(labels.shape, dtype=D)
    for b in range(labels.shape[0]):
        for t in range(labels.shape[1]):
            zz = z[b, t]
            lse = zz.max() + torch.log(torch.exp(zz - zz.max()).sum())
            out[b, t] = zz[labels[b, t]] - lse
    return out


@pytest.mark.parametrize("temperature", [1.0, 0.7])
def test_fp32_head_matches_fp64_oracle(temperature):
    h, W, labels = _hw()
    lp = fp32_head_logprobs(h, W, labels, temperature, chunk_tokens=5)
    assert lp.dtype == torch.float32
    err = (lp.to(D) - ref_head_logprobs(h, W, labels, temperature)).abs().max()
    assert float(err) < 5e-5, float(err)  # |logit|~50: fp32 spacing ~4e-6


def test_fp32_head_chunk_invariance_and_gradient_oracle():
    h, W, labels = _hw(seed=2)
    outs = [fp32_head_logprobs(h, W, labels, chunk_tokens=c) for c in (1, 3, 100)]
    for o in outs[1:]:
        torch.testing.assert_close(o, outs[0], rtol=0, atol=2e-5)
    # gradient wrt hidden and weight vs an fp64 autograd oracle on the same values
    hh = h.clone().requires_grad_(True)
    Wp = W.clone().requires_grad_(True)
    w = torch.randn(labels.shape, generator=torch.Generator().manual_seed(9))
    (fp32_head_logprobs(hh, Wp, labels, chunk_tokens=4) * w).sum().backward()
    h64 = h.to(D).requires_grad_(True)
    W64 = W.to(D).requires_grad_(True)
    z = h64 @ W64.t()
    ref = (z.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) * w.to(D)).sum()
    ref.backward()
    assert hh.grad.dtype == torch.bfloat16 and Wp.grad.dtype == torch.bfloat16  # cast back to param dtype
    # bf16 grads: compare with bf16 rounding tolerance of the fp64 oracle
    torch.testing.assert_close(hh.grad.to(D), h64.grad, rtol=2 ** -7, atol=2 ** -7 * float(h64.grad.abs().max()))
    torch.testing.assert_close(Wp.grad.to(D), W64.grad, rtol=2 ** -7, atol=2 ** -7 * float(W64.grad.abs().max()))


def test_fp32_head_entropy_matches_oracle():
    h, W, labels = _hw(seed=4, scale=2.0)
    lp, ent = fp32_head_logprobs(h, W, labels, with_entropy=True, chunk_tokens=3)
    p = (h.to(D) @ W.to(D).t()).softmax(-1)
    ref = -(p * p.log()).sum(-1)
    torch.testing.assert_close(ent.to(D), ref, rtol=0, atol=1e-5)
    assert not ent.requires_grad


def test_fp32_head_disables_tf32_and_restores():
    # Global precision API only: mixing it with the legacy `allow_tf32` setter
    # leaves a state where torch.get_float32_matmul_precision() raises (torch 2.11).
    prev = torch.get_float32_matmul_precision()
    torch.set_float32_matmul_precision("high")
    try:
        h, W, labels = _hw()
        seen = []
        import align.rl.cispo as c

        import contextlib

        orig = c._strict_fp32_matmul

        @contextlib.contextmanager
        def spy():
            with orig():
                seen.append(torch.backends.cuda.matmul.allow_tf32)  # the flag AS THE MATMUL SEES IT
                yield

        c._strict_fp32_matmul = spy
        try:
            fp32_head_logprobs(h, W, labels)
        finally:
            c._strict_fp32_matmul = orig
        assert seen and not any(seen)  # off inside the head
        assert torch.backends.cuda.matmul.allow_tf32 is True  # restored
    finally:
        torch.set_float32_matmul_precision(prev)


# -- tiny bf16 model: sampler/trainer mismatch, measured -------------------


class _Bf16LM(torch.nn.Module):
    """bf16 embedding -> bf16 body -> bf16 lm_head, with the hidden_consumer seam
    (the Qwen3/Olmoe contract). Scaled so head logits reach |z|~30-60: the
    "high-magnitude activation" regime MiniMax identified."""

    def __init__(self, V=48, H=32, scale=6.0):
        super().__init__()
        torch.manual_seed(0)
        self.emb = torch.nn.Embedding(V, H).to(torch.bfloat16)
        self.body = torch.nn.Linear(H, H, bias=False).to(torch.bfloat16)
        self.lm_head = torch.nn.Linear(H, V, bias=False).to(torch.bfloat16)
        with torch.no_grad():
            self.emb.weight.mul_(scale)
            self.lm_head.weight.mul_(1.5)

    def forward(self, input_ids, attention_mask=None, hidden_consumer=None, return_hidden_states=False):
        h = self.body(self.emb(input_ids))
        if hidden_consumer is not None:
            return hidden_consumer(h, self.lm_head.weight)
        return h if return_hidden_states else self.lm_head(h)


def _split_k_logits(h, W, dtype, parts=4):
    """An INFERENCE engine's head: split-K accumulation, each partial rounded to
    `dtype`, partials summed in `dtype` -- a different but legitimate summation
    order than the trainer's single GEMM."""
    H = h.shape[-1]
    acc = None
    for k in range(parts):
        sl = slice(k * H // parts, (k + 1) * H // parts)
        part = (h[..., sl].to(dtype) @ W[:, sl].to(dtype).t()).to(dtype)
        acc = part if acc is None else (acc + part).to(dtype)
    return acc


def _sampler_lp(h, W, labels, dtype):
    z = _split_k_logits(h, W, dtype).float()  # engines softmax in fp32 over (bf16 | fp32) logits
    return z.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)


def test_fp32_head_reduces_sampler_trainer_logprob_mismatch(capsys):
    m = _Bf16LM()
    g = torch.Generator().manual_seed(5)
    ids = torch.randint(0, 48, (16, 24), generator=g)
    labels = ids[:, 1:]
    with torch.no_grad():
        h = m(ids, return_hidden_states=True)[:, :-1]
        W = m.lm_head.weight
        truth = ref_head_logprobs(h, W, labels).float()
        # bf16 pipeline: trainer = module's bf16 path, sampler = split-K bf16
        trainer_bf16 = m(ids)[:, :-1].log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1).float()
        sampler_bf16 = _sampler_lp(h, W, labels, torch.bfloat16)
        # fp32-head pipeline: trainer = fp32 head, sampler = split-K fp32
        trainer_fp32 = fp32_head_forward(m, ids, torch.ones_like(ids), chunk_tokens=64)
        sampler_fp32 = _sampler_lp(h, W, labels, torch.float32)
    mm_bf16 = (trainer_bf16 - sampler_bf16).abs()
    mm_fp32 = (trainer_fp32 - sampler_fp32).abs()
    err_bf16 = (trainer_bf16 - truth).abs()
    err_fp32 = (trainer_fp32 - truth).abs()
    print(
        f"\nMISMATCH trainer-vs-sampler |dlogp|: bf16 head mean={mm_bf16.mean():.3e} max={mm_bf16.max():.3e} | "
        f"fp32 head mean={mm_fp32.mean():.3e} max={mm_fp32.max():.3e}\n"
        f"ERROR vs fp64 truth: bf16 mean={err_bf16.mean():.3e} max={err_bf16.max():.3e} | "
        f"fp32 mean={err_fp32.mean():.3e} max={err_fp32.max():.3e}"
    )
    # the ratio r = exp(trainer - sampler) is what enters the IS weight
    r_bf16, r_fp32 = torch.exp(trainer_bf16 - sampler_bf16), torch.exp(trainer_fp32 - sampler_fp32)
    print(f"max |r-1|: bf16 {float((r_bf16 - 1).abs().max()):.3e}  fp32 {float((r_fp32 - 1).abs().max()):.3e}")
    assert mm_bf16.mean() > 1e-3  # the bf16 pipeline really mismatches in this regime
    assert mm_fp32.mean() * 100 < mm_bf16.mean()
    assert mm_fp32.max() * 100 < mm_bf16.max()
    assert err_fp32.max() < 1e-4 < err_bf16.max()


# -- module integration -----------------------------------------------------


class _Sgd:
    def __init__(self, params, lr=0.01):
        self.params = [p for p in params if p.requires_grad]
        self.inner = torch.optim.SGD(self.params, lr=lr)
        self.last_grad_norm = 0.0
        self.skipped_steps = 0

    def step(self):
        self.last_grad_norm = float(torch.linalg.vector_norm(torch.stack(
            [p.grad.float().norm() for p in self.params if p.grad is not None])))
        self.inner.step()
        return True

    def zero_grad(self):
        self.inner.zero_grad(set_to_none=True)

    def state_dict(self):
        return self.inner.state_dict()

    def load_state_dict(self, sd):
        self.inner.load_state_dict(sd)


def _p4_cfg(**kw):
    base = dict(loss_fn=GRPOLossType.cispo, beta=0.0, load_ref_policy=False, loss_denominator=256.0,
                num_samples_per_prompt_rollout=8, clip_higher=0.272, record_entropy=True)
    base.update(kw)
    return GRPOLossConfig(**base)


def _batch(tm, model, B=4, T=12, prompt=4, old_shift=0.0):
    g = torch.Generator().manual_seed(1)
    ids = torch.randint(1, 48, (B, T), generator=g)
    attn = torch.ones(B, T, dtype=torch.int64)
    resp = torch.zeros(B, T, dtype=torch.int64)
    resp[:, prompt:] = 1
    old = tm.compute_logprobs_nograd(ids, attn) + old_shift
    return RolloutBatch(input_ids=ids, attention_mask=attn, response_mask=resp,
                        advantages=torch.tensor([1.0, -1.0, 0.5, -0.5])[:B], old_logprobs=old,
                        vllm_logprobs=old.clone(), group_ids=torch.arange(B) // 2,
                        rewards={"reward/total": torch.ones(B)})


def test_module_cispo_fp32_head_end_to_end():
    model = _Bf16LM()
    tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(fp32_head=True, fp32_head_chunk_tokens=7))
    batch = _batch(tm, model)
    # old (no-grad path) and new (grad path) share one numeric basis: ratio == 1
    out = tm._grpo_loss_adapter(model, batch)
    assert out.metrics["grpo/ratio"] == pytest.approx(1.0, abs=1e-6)
    assert math.isfinite(out.metrics["entropy"]) and out.metrics["entropy"] >= 0
    out.loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad.float()).all() for p in model.parameters())
    # the no-grad logprobs equal the fp64 oracle
    ids, attn = batch.input_ids, batch.attention_mask
    with torch.no_grad():
        h = model(ids, return_hidden_states=True)[:, :-1]
        ref = ref_head_logprobs(h, model.lm_head.weight, ids[:, 1:])
    torch.testing.assert_close(batch.old_logprobs.to(D), ref, rtol=0, atol=1e-4)


def test_module_fp32_head_off_is_bit_identical_to_logits_path():
    """Default (fp32_head=False) takes the untouched branch: same loss bits as a
    manual replay of the pre-change adapter math."""
    model = _Bf16LM()
    cfg = _p4_cfg()
    assert cfg.fp32_head is False
    tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, cfg)
    batch = _batch(tm, model)
    out = tm._grpo_loss_adapter(model, batch)
    logits = model(batch.input_ids)[:, :-1]
    from align.rl.ops import log_softmax_and_gather

    nl = log_softmax_and_gather(logits, batch.input_ids[:, 1:])
    ratio = torch.exp(nl - batch.old_logprobs)
    adv = batch.advantages[:, None].expand_as(nl)
    pg, _ = _legacy_compute_grpo_loss(nl, ratio, adv, cfg)
    mask = batch.shifted_response_mask.bool()
    expect = (pg * mask).sum() / 256.0
    assert float((out.loss.detach() - expect.to(out.loss.dtype)).abs()) == 0.0


def test_module_clipped_tokens_still_update_where_dapo_does_not():
    """End to end through the module: push old_logprobs far below new so every
    ratio >> 1+eps (positive advantage rows): the DAPO module's grad on those rows
    is zero, the CISPO module's is not."""
    shift = -3.0  # old = new - 3 -> ratio e^3 ~ 20
    grads = {}
    for name, cfg in (("cispo", _p4_cfg()),
                      ("dapo", GRPOLossConfig(beta=0.0, load_ref_policy=False, loss_denominator=256.0,
                                              num_samples_per_prompt_rollout=8, record_entropy=True))):
        model = _Bf16LM()
        tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, cfg)
        batch = _batch(tm, model, old_shift=shift)
        batch.advantages = torch.tensor([1.0, 1.0, 1.0, 1.0])  # all positive: pure upper-clip regime
        batch.vllm_logprobs = batch.old_logprobs.clone()
        tm._grpo_loss_adapter(model, batch).loss.backward()
        grads[name] = sum(float(p.grad.float().abs().sum()) for p in model.parameters())
    assert grads["dapo"] == 0.0
    assert grads["cispo"] > 0.0


def test_module_p4_gate_wired_and_a3_default_kept():
    model = _Bf16LM()
    with pytest.raises(P4CispoConformanceError):
        GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(beta=0.05))
    with pytest.raises(A3ConformanceError):  # dapo path still A3-gated
        GRPOTrainModule(model, _Sgd(model.parameters()), 1,
                        GRPOLossConfig(beta=0.0, load_ref_policy=False, loss_denominator="token"))


def test_fp32_head_needs_the_seam():
    class NoSeam(torch.nn.Module):
        def forward(self, input_ids, attention_mask=None):
            return torch.zeros(*input_ids.shape, 4)

    with pytest.raises(TypeError, match="hidden_consumer"):
        fp32_head_forward(NoSeam(), torch.zeros(1, 3, dtype=torch.long), torch.ones(1, 3))