"""DPO loss adapter: batch plumbing, metric contract, ref-logprob refusal.

Run: uv run pytest tests/test_dpo_loss.py -m "not gpu" -v
(venv equivalent on this box: .venv/bin/python -m pytest tests/test_dpo_loss.py -m "not gpu" -v)

The kernel's own tests (`test_fused_dpo.py`) cover the objective. This file
covers the layer between the Trainer and the kernel: unpacking a paired batch,
the caller-side label shift, the mandatory-metric contract, and the refusals
that must not be silent.
"""

from __future__ import annotations

import math

import pytest
import torch

from align.config import ModelConfig
from align.kernels.fused_dpo import MANDATORY_METRIC_KEYS
from align.train.dpo_loss import (
    IGNORE_INDEX,
    DpoLossConfig,
    fused_dpo_available,
    make_dpo_loss_adapter,
)


def tiny_model(consts: dict):
    from align.model.mistral import NmoMistralForCausalLM

    cfg = ModelConfig(
        name="tiny-mistral-dpo",
        n_layers=consts["n_layers"],
        hidden_size=consts["hidden_size"],
        intermediate_size=consts["intermediate_size"],
        n_heads=consts["n_heads"],
        n_kv_heads=consts["n_kv_heads"],
        head_dim=consts["head_dim"],
        vocab_size=consts["vocab_size"],
        max_position_embeddings=consts["max_seq_len"],
        rope_theta=consts["rope_theta"],
        precision="fp32",
    )
    cfg.n_params = -1
    torch.manual_seed(consts["seed"])
    return NmoMistralForCausalLM(cfg, attn_implementation="eager"), cfg


def paired_batch(cfg, pairs=3, seq=9, prompt_len=3, seed=0):
    g = torch.Generator().manual_seed(seed)
    out = {}
    for side in ("chosen", "rejected"):
        ids = torch.randint(0, cfg.vocab_size, (pairs, seq), generator=g)
        labels = ids.clone()
        labels[:, :prompt_len] = IGNORE_INDEX
        out[f"input_ids_{side}"] = ids
        out[f"labels_{side}"] = labels
    out["ref_logp_chosen"] = torch.randn(pairs, generator=g) - 20.0
    out["ref_logp_rejected"] = torch.randn(pairs, generator=g) - 22.0
    return out


def test_adapter_returns_loss_output_with_mandatory_metrics(tiny_model_constants):
    """The Trainer reads `loss`, `n_tokens` and `metrics` off a LossOutput, and
    section-12's dashboard reads the mandatory metric keys off `metrics` by
    name. Both are checked here so a rename fails at the adapter rather than
    at analysis time."""
    model, cfg = tiny_model(tiny_model_constants)
    batch = paired_batch(cfg, seed=1)
    out = make_dpo_loss_adapter()(model, batch)

    assert torch.isfinite(out.loss) and out.loss.dim() == 0
    for key in MANDATORY_METRIC_KEYS:
        assert key in out.metrics, key
        assert not math.isnan(out.metrics[key]), key
    # 3 pairs x (9 - 1 shifted - 3 prompt + 1) scored tokens on each side.
    expected = 2 * 3 * (9 - 1 - 3 + 1)
    assert out.n_tokens == expected, out.n_tokens


def test_adapter_shifts_labels_exactly_once(tiny_model_constants):
    """Position t predicts token t+1, and the shift lives on THIS side of the
    boundary (matching sft_loss). A second shift inside the kernel would be
    invisible: it produces a slightly worse but entirely plausible loss. The
    token count is the observable that pins it."""
    model, cfg = tiny_model(tiny_model_constants)
    batch = paired_batch(cfg, pairs=2, seq=10, prompt_len=4, seed=2)
    out = make_dpo_loss_adapter()(model, batch)
    per_side = 10 - 1 - 4 + 1
    assert out.n_tokens == 2 * 2 * per_side


def test_adapter_gradients_reach_the_policy_but_not_the_reference(tiny_model_constants):
    """End-to-end through the real tiny model: the policy trunk and the LM head
    receive gradient, and a requires_grad reference tensor does not. This is
    the section-12 contract exercised through the adapter rather than the
    kernel, because the adapter is where a stray graph edge would be
    introduced."""
    model, cfg = tiny_model(tiny_model_constants)
    batch = paired_batch(cfg, seed=3)
    batch["ref_logp_chosen"] = batch["ref_logp_chosen"].clone().requires_grad_(True)

    out = make_dpo_loss_adapter()(model, batch)
    out.loss.backward()

    assert batch["ref_logp_chosen"].grad is None
    assert model.lm_head.weight.grad is not None
    assert torch.isfinite(model.lm_head.weight.grad).all()
    trunk = model.model.layers[0].attn.q_proj.weight
    assert trunk.grad is not None and trunk.grad.abs().sum() > 0, (
        "no gradient reached the trunk: the loss is not training the model"
    )


def test_missing_reference_logprobs_raise(tiny_model_constants):
    """Guard both directions. A batch without reference log-probs is refused;
    substituting zeros would turn DPO into likelihood maximisation on the
    chosen side, which trains and converges and is not DPO. A complete batch
    is accepted silently."""
    model, cfg = tiny_model(tiny_model_constants)
    adapter = make_dpo_loss_adapter()
    full = paired_batch(cfg, seed=4)

    for key in ("ref_logp_chosen", "ref_logp_rejected", "labels_chosen", "input_ids_rejected"):
        broken = {k: v for k, v in full.items() if k != key}
        with pytest.raises(TypeError, match=key):
            adapter(model, broken)
    assert torch.isfinite(adapter(model, full).loss)


def test_adapter_rejects_unknown_variant():
    """Guard both directions: a typo'd variant raises when the adapter is
    BUILT, not at the first step, and every real variant builds."""
    from align.kernels.fused_dpo import DPO_VARIANTS

    with pytest.raises(ValueError, match="variant"):
        make_dpo_loss_adapter(DpoLossConfig(variant="sigmiod"))
    for variant in DPO_VARIANTS:
        assert make_dpo_loss_adapter(DpoLossConfig(variant=variant)) is not None


def test_token_chunk_size_default_is_not_none():
    """INCIDENT this default prevents: pairs are the chunking unit, and the
    pre-registered config is micro_batch 1 -- ONE pair per step. With
    `token_chunk_size=None` the fused DPO loss would materialise the whole
    [1024, 32768] logit tensor while reporting a healthy loss and passing
    every parity test. The trainer-facing default must therefore set it."""
    cfg = DpoLossConfig()
    assert cfg.token_chunk_size is not None and cfg.token_chunk_size > 0
    assert cfg.beta == 0.1 and cfg.variant == "sigmoid"
    assert cfg.nll_coeff == 0.0, (
        "the NLL auxiliary is a tripwire RESPONSE; it must be off by default"
    )


def test_adapter_config_reaches_the_kernel(tiny_model_constants):
    """The config is not decorative: changing beta must change the loss, and
    the dpop variant must differ from sigmoid on a batch where its hinge is
    open. An adapter that dropped its config on the floor would run every arm
    under the default objective."""
    model, cfg = tiny_model(tiny_model_constants)
    batch = paired_batch(cfg, seed=5)
    lo = make_dpo_loss_adapter(DpoLossConfig(beta=0.1))(model, batch).loss
    hi = make_dpo_loss_adapter(DpoLossConfig(beta=0.9))(model, batch).loss
    assert abs(float(lo.detach()) - float(hi.detach())) > 1e-6

    # reference far above the policy -> the dpop hinge is open
    open_hinge = dict(batch, ref_logp_chosen=torch.zeros_like(batch["ref_logp_chosen"]))
    sig = make_dpo_loss_adapter(DpoLossConfig(variant="sigmoid"))(model, open_hinge).loss
    dpop = make_dpo_loss_adapter(DpoLossConfig(variant="dpop"))(model, open_hinge).loss
    assert float(dpop.detach()) > float(sig.detach()) + 1e-6


def test_fused_dpo_is_available_and_required(tiny_model_constants):
    """Both directions on the require_fused guard: the kernel IS importable
    here, so building with require_fused=True must succeed. The refusal branch
    is what the 7B stage scripts rely on -- a run must not quietly fall back to
    a reference path that no A10G survives."""
    assert fused_dpo_available() is True
    assert make_dpo_loss_adapter(require_fused=True).uses_fused is True


def test_dpo_loss_module_builds_no_reference_model():
    """Structural: nothing in this section constructs or loads a second model.
    The reference log-probs are data. An import that pulled a model in here
    would silently double resident weights at DPO time."""
    import pathlib

    src = pathlib.Path(__file__).resolve().parents[1] / "align" / "train" / "dpo_loss.py"
    text = src.read_text()
    for forbidden in ("from_pretrained", "load_state_dict", "AutoModel", "NmoMistralForCausalLM"):
        assert forbidden not in text, f"dpo_loss.py must not reference {forbidden!r}"