"""Length-grouped batching + Preference collate.
Run: uv run pytest tests/test_data_loader.py -m "not gpu" -v

Padding waste is not a minor inefficiency here: the per-micro-batch cost is
dominated by a fixed host-staged all-gather, so wasted padding is wasted
wall-clock in direct proportion. And a DPO batch that reports 0 for an
undefined statistic is indistinguishable from one reporting a real 0 -- every
undefined statistic in this module is NaN.
"""

from __future__ import annotations

import math

import numpy as np
import pytest

from align.data.loader import (
    LengthGroupedSampler,
    Preference,
    collate_preference,
    preference_length_key,
)

PAD = 0


def _pair_row(
    prompt_len: int, chosen_completion: int, rejected_completion: int, pid: str = "p0"
) -> dict:
    """One preference row in the on-disk shape the pair shards produce."""
    chosen = list(range(10, 10 + prompt_len + chosen_completion))
    rejected = chosen[:prompt_len] + list(range(500, 500 + rejected_completion))
    return {
        "chosen_input_ids": np.asarray(chosen, dtype=np.int32),
        "rejected_input_ids": np.asarray(rejected, dtype=np.int32),
        "chosen_prompt_len": prompt_len,
        "rejected_prompt_len": prompt_len,
        "problem_id": pid,
        "subject": "Algebra",
    }


def test_length_grouped_batches_bound_padding_waste():
    """Batches drawn from a synthetic length distribution have intra-batch
    length spread below the configured bound (grouping actually groups).
    Both directions: the same distribution batched WITHOUT grouping blows the
    bound, so the assertion is measuring grouping and not the distribution."""
    rng = np.random.default_rng(0)
    lengths = rng.integers(64, 1024, size=2000)

    grouped = LengthGroupedSampler(lengths, max_tokens=8192, max_rows=8, seed=0, shuffle=True)
    spreads = [int(lengths[b].max() - lengths[b].min()) for b in grouped if len(b) > 1]
    assert spreads, "sampler produced no multi-row batches"
    assert max(spreads) <= 128, f"intra-batch spread too large: {max(spreads)}"

    # Firing direction: ungrouped batches of the same size are far worse.
    order = rng.permutation(len(lengths))
    naive = [order[i : i + 8] for i in range(0, len(order), 8)]
    naive_spread = max(int(lengths[b].max() - lengths[b].min()) for b in naive)
    assert naive_spread > 128

    # A token budget, not a row count: padded cost stays under the budget.
    for b in grouped:
        assert int(lengths[b].max()) * len(b) <= 8192 or len(b) == 1


def test_length_grouping_is_seeded_and_reproducible():
    """Same seed => identical batch order across two loader constructions
    (required by the section-07 seeded-resume contract). Different seed =>
    different order, so the test is not passing on a degenerate sampler."""
    lengths = np.random.default_rng(1).integers(64, 1024, size=500)

    a = [list(map(int, b)) for b in LengthGroupedSampler(lengths, 4096, seed=7)]
    b = [list(map(int, x)) for x in LengthGroupedSampler(lengths, 4096, seed=7)]
    c = [list(map(int, x)) for x in LengthGroupedSampler(lengths, 4096, seed=8)]
    assert a == b
    assert a != c

    # set_epoch changes the order deterministically, and reproducibly so.
    s1 = LengthGroupedSampler(lengths, 4096, seed=7)
    s1.set_epoch(1)
    s2 = LengthGroupedSampler(lengths, 4096, seed=7)
    s2.set_epoch(1)
    assert [list(map(int, x)) for x in s1] == [list(map(int, x)) for x in s2]
    assert [list(map(int, x)) for x in s1] != a


def test_preference_collate_shapes_and_masks():
    """Preference batch exposes chosen/rejected ids, attention masks, prompt
    lengths, and per-row metadata; completion-token masks exclude prompt
    positions.
    Incident class: a completion mask that includes prompt positions makes the
    DPO logprob a function of the (identical) prompt, which cancels in the
    difference for shared-prefix pairs and quietly shrinks the signal."""
    rows = [_pair_row(5, 3, 4, "pA"), _pair_row(7, 9, 2, "pB")]
    batch = collate_preference(rows, pad_token_id=PAD)

    assert isinstance(batch, Preference)
    assert batch.chosen_input_ids.shape == (2, 16)  # max(5+3, 7+9)
    assert batch.rejected_input_ids.shape == (2, 9)  # max(5+4, 7+2)
    assert batch.chosen_attention_mask.sum().item() == (5 + 3) + (7 + 9)
    assert batch.rejected_attention_mask.sum().item() == (5 + 4) + (7 + 2)

    # Completion mask: excludes prompt positions AND padding.
    assert batch.chosen_completion_mask[0].tolist() == [0] * 5 + [1] * 3 + [0] * 8
    assert batch.rejected_completion_mask[1].tolist() == [0] * 7 + [1] * 2

    assert batch.chosen_prompt_lens.tolist() == [5, 7]
    assert batch.rejected_prompt_lens.tolist() == [5, 7]
    assert [m["problem_id"] for m in batch.metadata] == ["pA", "pB"]

    stats = batch.stats()
    assert stats["n_pairs"] == 2
    assert stats["mean_chosen_completion_tokens"] == pytest.approx((3 + 9) / 2)
    assert stats["mean_rejected_completion_tokens"] == pytest.approx((4 + 2) / 2)
    assert stats["chosen_longer_fraction"] == pytest.approx(0.5)

    # Grouping key for preference data is max(len(chosen), len(rejected)).
    assert preference_length_key(rows[0]) == 9
    assert preference_length_key(rows[1]) == 16


def test_preference_collate_empty_stats_are_nan():
    """Any statistic undefined on an empty/degenerate batch is NaN, never 0.
    Incident: a 0 that means 'no data' is indistinguishable from a 0 that means
    'measured zero' -- the predecessor's reward channel read as an honest
    measurement of a weak policy for two full runs."""
    empty = collate_preference([], pad_token_id=PAD)
    s = empty.stats()
    assert s["n_pairs"] == 0
    for key in (
        "mean_chosen_completion_tokens",
        "mean_rejected_completion_tokens",
        "chosen_longer_fraction",
        "mean_prompt_tokens",
    ):
        assert math.isnan(s[key]), f"{key} must be NaN on an empty batch, got {s[key]}"

    # Degenerate: rows exist but no completion tokens at all.
    degenerate = collate_preference([_pair_row(4, 0, 0)], pad_token_id=PAD)
    ds = degenerate.stats()
    assert ds["n_pairs"] == 1
    assert ds["mean_prompt_tokens"] == pytest.approx(4.0)
    assert math.isnan(ds["mean_chosen_completion_tokens"])
    assert math.isnan(ds["chosen_longer_fraction"])

    # Silent direction: a normal batch reports real numbers, not NaN.
    ok = collate_preference([_pair_row(4, 2, 1)], pad_token_id=PAD).stats()
    assert not math.isnan(ok["mean_chosen_completion_tokens"])


def test_pad_token_never_enters_the_completion_mask():
    """Right-padding only: boundary receipts are absolute indices into the
    unpadded sequence, so left padding would shift every stored position.
    The test pins the padding side by construction."""
    rows = [_pair_row(2, 1, 1), _pair_row(2, 6, 6)]
    b = collate_preference(rows, pad_token_id=PAD)
    # Row 0 is the short one: real tokens first, padding at the tail.
    assert b.chosen_input_ids[0, :3].tolist() != [PAD, PAD, PAD]
    assert b.chosen_input_ids[0, 3:].tolist() == [PAD] * (b.chosen_input_ids.shape[1] - 3)
    assert b.chosen_completion_mask[0].sum().item() == 1


def test_stored_pair_rows_collate_into_a_preference_batch(tmp_workdir):
    """Shards -> Preference batch, through the documented adapter. This is the
    seam section-12 loads both DPO arms across, so it is tested rather than
    assumed: the pair shards store the CHOSEN side as the primary input_ids
    and the rejected side as a ragged extra, and a consumer that guessed that
    mapping wrong would train on chosen-vs-chosen and report a healthy loss."""
    from align.data.loader import preference_dataset
    from align.data.store import ShardWriter

    with ShardWriter(
        tmp_workdir,
        "arm-0000",
        {
            "input_ids": "int32",
            "rejected_input_ids": "int32",
            "rejected_boundary_positions": "int32",
        },
    ) as w:
        for i in range(3):
            chosen = list(range(10, 20 + i))
            rejected = chosen[:5] + [900, 901]
            w.add_row(
                chosen,
                problem_id=f"p{i}",
                subject="Algebra",
                boundary_positions=[3, 6],
                ragged={"rejected_input_ids": rejected, "rejected_boundary_positions": [3]},
                chosen_prompt_len=5,
                rejected_prompt_len=5,
                shared_prefix_len=5,
                step_index=2,
                pair_id=f"pair{i}",
                source="arm_b",
            )

    rows = preference_dataset(tmp_workdir)
    assert len(rows) == 3
    batch = collate_preference(rows, pad_token_id=PAD)
    assert batch.chosen_input_ids.shape[0] == 3
    assert batch.rejected_input_ids.shape[1] == 7
    # Prompt positions are masked out of both completion masks.
    assert batch.chosen_completion_mask[:, :5].sum().item() == 0
    assert batch.rejected_completion_mask[0].tolist() == [0] * 5 + [1, 1]
    assert [m["pair_id"] for m in batch.metadata] == ["pair0", "pair1", "pair2"]
    assert not math.isnan(batch.stats()["mean_chosen_completion_tokens"])