"""Decoupled DPO backward == coupled backward, to numerical identity.

The 7B stage backwards each side's trunk separately (both checkpointed
graphs resident at once measured 21.5GB and OOM'd). The decoupling must be
a pure memory transformation: same loss, same metrics, same parameter
gradients as the coupled path the kernel tests already certify.
"""

import pytest
import torch

from align.model.mistral import NmoMistralForCausalLM
from align.train.dpo_loss import DpoLossConfig, make_dpo_loss_adapter

IGNORE_INDEX = -100


@pytest.fixture
def tiny(tiny_model_constants):
    from test_trainer import tiny_config

    return tiny_config(tiny_model_constants)


def _batch(vocab: int, b: int = 2, t: int = 12, seed: int = 7) -> dict:
    g = torch.Generator().manual_seed(seed)
    ids_c = torch.randint(0, vocab, (b, t), generator=g)
    ids_r = torch.randint(0, vocab, (b, t), generator=g)
    lab_c = ids_c.clone()
    lab_r = ids_r.clone()
    lab_c[:, : t // 2] = IGNORE_INDEX  # prompt masked, completion live
    lab_r[:, : t // 2] = IGNORE_INDEX
    return {
        "input_ids_chosen": ids_c,
        "labels_chosen": lab_c,
        "input_ids_rejected": ids_r,
        "labels_rejected": lab_r,
        "ref_logp_chosen": torch.randn(b, generator=g),
        "ref_logp_rejected": torch.randn(b, generator=g),
    }


def _grads(model) -> dict:
    return {
        n: p.grad.detach().clone()
        for n, p in model.named_parameters()
        if p.grad is not None and p.grad.abs().sum() > 0
    }


def test_decoupled_backward_matches_coupled(tiny):
    torch.manual_seed(0)
    model = NmoMistralForCausalLM(tiny, device="cpu")
    batch = _batch(tiny.vocab_size)

    coupled = make_dpo_loss_adapter(DpoLossConfig(), require_fused=True)
    out = coupled(model, batch)
    out.loss.backward()
    g_coupled = _grads(model)
    assert g_coupled, "coupled path produced no gradients"
    loss_coupled = float(out.loss.detach())
    model.zero_grad(set_to_none=True)

    dec = make_dpo_loss_adapter(
        DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=1.0
    )
    out2 = dec(model, batch)
    # The proxy honours the trainer's contract without touching model grads.
    g_before_proxy = _grads(model)
    out2.loss.backward()
    g_after_proxy = _grads(model)
    assert set(g_before_proxy) == set(g_after_proxy)
    for n in g_before_proxy:
        assert torch.equal(g_before_proxy[n], g_after_proxy[n]), (
            f"proxy backward mutated model grad {n}"
        )

    assert abs(float(out2.loss.detach()) - loss_coupled) < 1e-6
    assert set(out2.metrics) == set(out.metrics)
    assert set(g_after_proxy) == set(g_coupled)
    for n in g_coupled:
        assert torch.allclose(g_coupled[n], g_after_proxy[n], rtol=1e-5, atol=1e-7), n


def test_decoupled_grad_scale_scales_real_grads(tiny):
    torch.manual_seed(0)
    model = NmoMistralForCausalLM(tiny, device="cpu")
    batch = _batch(tiny.vocab_size)
    dec1 = make_dpo_loss_adapter(
        DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=1.0
    )
    dec1(model, batch)
    g1 = _grads(model)
    model.zero_grad(set_to_none=True)
    dec4 = make_dpo_loss_adapter(
        DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=0.25
    )
    dec4(model, batch)
    g4 = _grads(model)
    for n in g1:
        assert torch.allclose(g1[n] * 0.25, g4[n], rtol=1e-5, atol=1e-8), n