"""Adversarial-review additions for the P4 CISPO objective + FP32 head.

Each test pins a failure mode the builder's suite left open (see the docstring
of each). Oracles are independent fp64 computations or exact-bit comparisons.
"""

from __future__ import annotations

import math

import pytest
import torch

import align.rl.cispo as cispo_mod
from align.rl.cispo import (
    P4CispoConformanceError,
    assert_p4_cispo_conformant,
    cispo_policy_loss,
    fp32_head_logprobs,
)
from align.rl.grpo_loss import A3ConformanceError, GRPOLossConfig, GRPOLossType, assert_a3_conformant
from align.rl.grpo_module import GRPOTrainModule
from tests.test_cispo import _Bf16LM, _Sgd, _batch, _hw, _p4_cfg

D = torch.float64


# --------------------------------------------------------------- boundaries


@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
def test_ratio_exactly_at_clip_boundary(dtype):
    """r == 1+eps_high is NOT clipped (weight == r, clipfrac 0); the next
    representable value above IS (weight pinned at the cap, clipfrac 1), and the
    weight is continuous across the boundary. Same on the low side with a floor.
    Failure mode: `>=` vs `>` flips the clip metric at the single most common
    value (ratio==1+eps after eps-sized drift) without changing the loss."""
    eps_hi, eps_lo = 0.25, 0.5  # exactly representable -> no rounding ambiguity
    hi = torch.tensor(1.0 + eps_hi, dtype=dtype)
    lo = torch.tensor(1.0 - eps_lo, dtype=dtype)
    ratios = torch.stack([torch.nextafter(hi, torch.tensor(0.0, dtype=dtype)), hi,
                          torch.nextafter(hi, torch.tensor(9.0, dtype=dtype)),
                          torch.nextafter(lo, torch.tensor(9.0, dtype=dtype)), lo,
                          torch.nextafter(lo, torch.tensor(0.0, dtype=dtype))])
    lp = torch.full_like(ratios, -2.0, requires_grad=True)
    loss, clipfrac = cispo_policy_loss(lp, ratios, torch.ones_like(ratios), eps_hi, eps_lo)
    weight = (-loss / lp.detach()).detach()  # loss = -A * w * lp
    assert clipfrac.tolist() == [0, 0, 1, 0, 0, 1]
    assert float(weight[2]) == float(hi) and float(weight[5]) == float(lo)
    assert float(weight[1]) == float(hi) and float(weight[4]) == float(lo)
    assert float((weight[0] - hi).abs()) <= 2 * torch.finfo(dtype).eps


def test_all_zero_advantages_give_exactly_zero_gradient_and_loss():
    """A=0 (every group dead / centered to zero) must be an exact no-op even with
    wild ratios. Failure mode: weight*logp leaking through when A is zero."""
    lp = torch.randn(3, 6, requires_grad=True)
    ratio = torch.tensor([[0.0, 1e-9, 0.5, 1.0, 5.0, 1e9]] * 3)
    loss, _ = cispo_policy_loss(lp, ratio, torch.zeros(3, 6), 0.272, 0.5)
    loss.sum().backward()
    assert float(loss.abs().sum()) == 0.0 and float(lp.grad.abs().sum()) == 0.0


def test_inf_ratio_is_clipped_not_propagated_but_nan_is_not_laundered():
    """exp(new-old) overflows to +inf on a runaway token. CISPO must turn that
    into the capped weight (finite loss and gradient); a NaN ratio (a real bug
    upstream) must surface as NaN, not be silently clamped into a finite weight.
    Failure modes: inf*0 -> NaN poisoning every grad; nan_to_num hiding a broken
    logprob path behind a healthy-looking loss."""
    lp = torch.tensor([-1.0, -1.0], requires_grad=True)
    loss, cf = cispo_policy_loss(lp, torch.tensor([float("inf"), 1.0]), torch.tensor([1.0, 1.0]), 0.272)
    loss.sum().backward()
    assert torch.isfinite(loss).all() and torch.isfinite(lp.grad).all() and float(cf[0]) == 1.0
    assert float(lp.grad[0]) == pytest.approx(-1.272)
    for eps_low in (None, 0.5):  # both branches of the clamp
        lp2 = torch.tensor([-1.0], requires_grad=True)
        loss2, _ = cispo_policy_loss(lp2, torch.tensor([float("nan")]), torch.tensor([1.0]), 0.272, eps_low)
        assert torch.isnan(loss2).all(), eps_low


@pytest.mark.parametrize("eps_low", [1.0, 5.0, 1e6])
def test_eps_low_ge_one_is_the_papers_no_lower_bound(eps_low):
    """Paper Sec 3.1: no lower bound is realised by setting eps_low to a LARGE
    value, so 1-eps_low <= 0 and the floor can never bind a positive ratio.
    `None` and any eps_low >= 1 must therefore give identical loss, gradient and
    clipfrac, including for ratio -> 0. Failure mode: treating a large eps_low
    as a real floor (e.g. abs()/clamp to >=eps) would re-introduce a bound."""
    ratio = torch.tensor([1e-12, 1e-3, 0.3, 1.0, 3.0, 50.0])
    adv = torch.tensor([1.0, -1.0, 1.0, -1.0, 1.0, -1.0])
    out = []
    for e in (None, eps_low):
        lp = torch.linspace(-3, -1, 6).requires_grad_(True)
        loss, cf = cispo_policy_loss(lp, ratio, adv, 0.272, e)
        loss.sum().backward()
        out.append((loss.detach(), lp.grad.clone(), cf))
    for a, b in zip(out[0], out[1]):
        assert torch.equal(a, b)


# ------------------------------------------------------- FP32 head internals


def test_head_is_checkpointed_not_stored():
    """Memory contract: backward RECOMPUTES each chunk's [chunk, V] logits
    (checkpoint) instead of saving them. Counted via the chunk function: n calls
    forward + n recompute in backward; with no grad enabled exactly n and no
    checkpoint machinery. Failure mode: dropping the checkpoint silently turns
    the head into a [B*T, V] fp32 resident tensor (10 GB at V=152k, 16k tok)."""
    h, W, labels = _hw(B=2, T=8)
    calls = []
    orig = cispo_mod._head_chunk_logprobs
    cispo_mod._head_chunk_logprobs = lambda *a, **k: (calls.append(1), orig(*a, **k))[1]
    try:
        hh = h.clone().requires_grad_(True)
        lp = fp32_head_logprobs(hh, W, labels, chunk_tokens=4)  # 16 tokens -> 4 chunks
        assert len(calls) == 4
        lp.sum().backward()
        assert len(calls) == 8
        calls.clear()
        with torch.no_grad():
            fp32_head_logprobs(h, W, labels, chunk_tokens=4)
        assert len(calls) == 4
    finally:
        cispo_mod._head_chunk_logprobs = orig


def test_strict_fp32_matmul_survives_every_precision_start_state():
    """On torch>=2.9 flipping the legacy `matmul.allow_tf32` flag while the
    global precision is 'high' (what most launchers set) leaves a mixed-API
    state in which torch.get_float32_matmul_precision() RAISES. The context must
    (a) never raise from any start state, (b) report 'highest' inside, and
    (c) restore the exact start precision."""
    start = torch.get_float32_matmul_precision()
    try:
        for s in ("highest", "high", "medium"):
            torch.set_float32_matmul_precision(s)
            with cispo_mod._strict_fp32_matmul():
                assert torch.get_float32_matmul_precision() == "highest"
                assert torch.backends.cuda.matmul.allow_tf32 is False
            assert torch.get_float32_matmul_precision() == s
            # and it restores on exceptions
            with pytest.raises(RuntimeError, match="boom"):
                with cispo_mod._strict_fp32_matmul():
                    raise RuntimeError("boom")
            assert torch.get_float32_matmul_precision() == s
    finally:
        torch.set_float32_matmul_precision(start)


# ------------------------------------------------------------ module / gate


def test_chunk_tokens_is_plumbed_to_both_module_paths(monkeypatch):
    """fp32_head_chunk_tokens is the memory knob; it must reach the head from
    BOTH the training forward and compute_logprobs_nograd. Failure mode: a
    hardcoded chunk silently ignoring the config on one path (OOM on 7B)."""
    seen = []
    orig = cispo_mod.fp32_head_logprobs

    def spy(hidden, weight, labels, temperature=1.0, chunk_tokens=2048, with_entropy=False):
        seen.append(chunk_tokens)
        return orig(hidden, weight, labels, temperature, chunk_tokens, with_entropy)

    monkeypatch.setattr(cispo_mod, "fp32_head_logprobs", spy)
    model = _Bf16LM()
    tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(fp32_head=True, fp32_head_chunk_tokens=13))
    batch = _batch(tm, model)
    tm._grpo_loss_adapter(model, batch)
    assert seen and set(seen) == {13}, seen


def test_record_entropy_off_skips_entropy_pass(monkeypatch):
    """record_entropy=False must not pay the extra no-grad head pass."""
    seen = []
    orig = cispo_mod.fp32_head_logprobs

    def spy(hidden, weight, labels, temperature=1.0, chunk_tokens=2048, with_entropy=False):
        seen.append(with_entropy)
        return orig(hidden, weight, labels, temperature, chunk_tokens, with_entropy)

    monkeypatch.setattr(cispo_mod, "fp32_head_logprobs", spy)
    model = _Bf16LM()
    tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1,
                         _p4_cfg(fp32_head=True, record_entropy=False))
    tm._grpo_loss_adapter(model, _batch(tm, model))
    assert seen and not any(seen)


def test_inert_clip_lower_knobs_are_refused():
    """`clip_lower` is never read by the cispo branch and `cispo_clip_lower` is
    never read by dapo; setting either to a non-default value used to be
    accepted and silently ignored (a config that LOOKS floored but is not)."""
    base = dict(beta=0.0, load_ref_policy=False, loss_denominator=1024.0, num_samples_per_prompt_rollout=16)
    with pytest.raises(P4CispoConformanceError, match="clip_lower"):
        assert_p4_cispo_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, clip_lower=0.5, **base))
    with pytest.raises(A3ConformanceError, match="cispo_clip_lower"):
        assert_a3_conformant(GRPOLossConfig(cispo_clip_lower=0.5, **base))
    # the sanctioned spelling still passes
    assert_p4_cispo_conformant(
        GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, cispo_clip_lower=0.5, **base)
    )


def test_gate_message_names_every_independent_violation():
    """The refusal collects, not short-circuits: a config with three problems
    must list all three (operators fix them in one pass)."""
    cfg = GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=-1.0, cispo_clip_lower=-2.0, beta=0.5,
                         load_ref_policy=False, loss_denominator=64.0, num_samples_per_prompt_rollout=16)
    with pytest.raises(P4CispoConformanceError) as ei:
        assert_p4_cispo_conformant(cfg)
    msg = str(ei.value)
    assert "clip_higher" in msg and "cispo_clip_lower" in msg and "beta" in msg