"""Regression tests for the P2 SpaceByte hierarchy.

These tests pin the patch-boundary contract separately from the full model and
prove the hierarchy can exactly collapse to the P1 plain decoder when every byte
position is promoted to a global patch.
"""
from __future__ import annotations

import torch

from ksbyte.config import BOS_ID, ByteLMConfig
from ksbyte.model import ByteDecoder, SpaceByteDecoder, build_model, build_spacebyte_boundary_mask


def _tiny_cfg(**overrides) -> ByteLMConfig:
    cfg = ByteLMConfig(
        d_model=32,
        n_layers=2,
        n_heads=4,
        n_kv_heads=2,
        mlp_ratio=2.0,
        ctx_len=16,
        dropout=0.0,
        n_local_in=0,
        n_global=2,
        n_local_out=0,
        max_patches=16,
        tie_embeddings=True,
    )
    return cfg.merge(overrides).validate()


def test_spacebyte_boundary_mask_marks_first_spacelike_byte_only():
    # ASCII bytes: 'a  b\tc', with BOS opening a fresh document.
    x = torch.tensor([[BOS_ID, ord("a"), 32, 32, ord("b"), 9, ord("c")]])
    seg_ids = torch.zeros_like(x)

    mask = build_spacebyte_boundary_mask(x, seg_ids=seg_ids)

    assert mask.tolist() == [[True, False, True, False, False, True, False]]


def test_spacebyte_gather_scatter_reuses_latest_patch_context():
    x = torch.tensor([[BOS_ID, ord("a"), 32, ord("b"), ord("c"), 32, ord("d")]])
    h = torch.arange(x.numel() * 3, dtype=torch.float32).view(1, x.numel(), 3)
    boundary = build_spacebyte_boundary_mask(x)

    gathered, patch_mask, patch_pos, patch_seg, patch_ids = SpaceByteDecoder.gather_patches(
        h,
        boundary,
        pos_ids=torch.arange(x.size(1)).unsqueeze(0),
        seg_ids=torch.zeros_like(x),
        max_patches=8,
    )
    scattered = SpaceByteDecoder.scatter_patches(gathered + 1000.0, patch_ids)

    assert patch_mask.tolist() == [[True, True, True]]
    assert patch_pos.tolist() == [[0, 2, 5]]
    assert patch_seg.tolist() == [[0, 0, 0]]
    # Positions before the first space use BOS patch; then the latest space patch.
    assert torch.equal(scattered[0, 0], gathered[0, 0] + 1000.0)
    assert torch.equal(scattered[0, 1], gathered[0, 0] + 1000.0)
    assert torch.equal(scattered[0, 2], gathered[0, 1] + 1000.0)
    assert torch.equal(scattered[0, 4], gathered[0, 1] + 1000.0)
    assert torch.equal(scattered[0, 5], gathered[0, 2] + 1000.0)
    assert torch.equal(scattered[0, 6], gathered[0, 2] + 1000.0)


def test_spacebyte_factory_builds_model():
    model = build_model(_tiny_cfg(variant="spacebyte"))
    assert isinstance(model, SpaceByteDecoder)


def test_spacebyte_reduces_to_plain_decoder_when_all_positions_are_boundaries():
    torch.manual_seed(7)
    plain_cfg = _tiny_cfg(variant="plain")
    space_cfg = _tiny_cfg(variant="spacebyte")
    plain = ByteDecoder(plain_cfg).eval()
    space = SpaceByteDecoder(space_cfg).eval()

    mapped = {}
    for name, tensor in plain.state_dict().items():
        mapped[name.replace("blocks.", "global_blocks.")] = tensor
    missing, unexpected = space.load_state_dict(mapped, strict=False)
    assert not unexpected
    assert set(missing) == set()

    x = torch.randint(0, 128, (2, 8))
    pos_ids = torch.arange(x.size(1)).unsqueeze(0).expand_as(x)
    all_boundaries = torch.ones_like(x, dtype=torch.bool)

    plain_logits, _, _ = plain(x, pos_ids=pos_ids)
    space_logits, _, _ = space(x, pos_ids=pos_ids, boundary_mask=all_boundaries)

    torch.testing.assert_close(space_logits, plain_logits, atol=1e-5, rtol=1e-5)