# tests/test_model.py
import pytest
import torch
from bclm.config import BCLM1ModelConfig
from bclm.model import BCLM1Model


def _small_config(**overrides) -> BCLM1ModelConfig:
    """Tiny config for fast CPU tests."""
    defaults = dict(
        vocab_size=256,
        embed_dim=64,
        n_layers=2,
        max_seq_len=128,
        dropout=0.0,
        attn_heads=2,
        attn_kv_heads=1,
        local_attn_layers=(0,),
        global_attn_layers=(1,),
        attn_window_size=32,
        conv_kernel_size=4,
        osc_n_pairs=1,
        osc_n_real=2,
        osc_clamp_min_decay=1e-5,
        bigram_table_factor=2,
    )
    defaults.update(overrides)
    return BCLM1ModelConfig(**defaults)


class TestBCLM1Model:
    def test_instantiation(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        assert model.config is cfg

    def test_num_parameters(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        n = model.num_parameters()
        assert isinstance(n, int)
        assert n > 0

    def test_num_parameters_exclude_embeddings(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        n_all = model.num_parameters(exclude_embeddings=False)
        n_no_emb = model.num_parameters(exclude_embeddings=True)
        assert n_no_emb < n_all

    def test_forward_logits_shape(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        model.eval()

        B, T = 2, 16
        idx = torch.randint(0, cfg.vocab_size, (B, T))

        with torch.no_grad():
            logits, loss = model(idx)

        assert logits is not None
        assert logits.shape == (B, T, cfg.vocab_size)
        assert loss is None

    def test_forward_with_targets(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        model.eval()

        B, T = 2, 16
        idx = torch.randint(0, cfg.vocab_size, (B, T))
        targets = torch.randint(0, cfg.vocab_size, (B, T))

        with torch.no_grad():
            logits, loss = model(idx, targets=targets)

        assert logits is not None
        assert loss is not None
        assert loss.ndim == 0
        assert loss.item() > 0

    def test_forward_seqlen_exceeds_max_raises(self):
        cfg = _small_config(max_seq_len=32)
        model = BCLM1Model(cfg)

        idx = torch.randint(0, cfg.vocab_size, (1, 64))
        with pytest.raises(ValueError, match="seqlen.*exceeds"):
            model(idx)

    def test_forward_single_token(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        model.eval()

        idx = torch.randint(0, cfg.vocab_size, (1, 1))
        with torch.no_grad():
            logits, loss = model(idx)

        assert logits.shape == (1, 1, cfg.vocab_size)

    def test_bigram_hash_deterministic(self):
        cfg = _small_config()
        idx = torch.tensor([[10, 20, 30, 40]])
        bvs = cfg.vocab_size * cfg.bigram_table_factor

        h1 = BCLM1Model._compute_bigram_hash(idx, bvs)
        h2 = BCLM1Model._compute_bigram_hash(idx, bvs)
        assert torch.equal(h1, h2)

    def test_bigram_hash_position_zero(self):
        cfg = _small_config()
        idx = torch.tensor([[5, 10, 15]])
        bvs = cfg.vocab_size * cfg.bigram_table_factor

        h = BCLM1Model._compute_bigram_hash(idx, bvs)
        assert h[0, 0].item() == bvs - 1

    def test_gradient_checkpointing_toggle(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)

        model.enable_gradient_checkpointing(True)
        for layer in model.layers:
            assert layer._use_checkpoint is True

        model.enable_gradient_checkpointing(False)
        for layer in model.layers:
            assert layer._use_checkpoint is False

    def test_gradient_checkpointing_selective(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)

        model.enable_gradient_checkpointing(True, layer_indices=[0])
        assert model.layers[0]._use_checkpoint is True
        assert model.layers[1]._use_checkpoint is False

    def test_chunked_ce(self):
        cfg = _small_config()
        model = BCLM1Model(cfg)
        model.eval()

        B, T = 2, 16
        idx = torch.randint(0, cfg.vocab_size, (B, T))
        targets = torch.randint(0, cfg.vocab_size, (B, T))

        with torch.no_grad():
            _, loss_standard = model(idx, targets=targets)

        model.set_ce_chunk_size(8)
        with torch.no_grad():
            logits_chunked, loss_chunked = model(idx, targets=targets)

        assert logits_chunked is None
        assert loss_chunked is not None
        assert torch.allclose(loss_standard, loss_chunked, atol=1e-4)

    def test_overlapping_attn_layers_raises(self):
        with pytest.raises(ValueError, match="both local and global"):
            _small_config(
                local_attn_layers=(0, 1),
                global_attn_layers=(1,),
            )
            # Config itself doesn't validate, but model construction does
            BCLM1Model(_small_config(
                local_attn_layers=(0, 1),
                global_attn_layers=(1,),
            ))