"""Tests for custom layers.

This module tests the custom layer implementations created for issue #35.
"""

from __future__ import annotations

import pytest
import torch

from sagellm_core.model.layers import (
    GELU,
    Attention,
    Embedding,
    LayerNorm,
    Linear,
    RMSNorm,
    SiLU,
    SwiGLU,
)


def test_linear_layer():
    """Test basic Linear layer."""
    layer = Linear(in_features=10, out_features=20, bias=True)

    x = torch.randn(2, 3, 10)
    output = layer(x)

    assert output.shape == (2, 3, 20)


def test_rmsnorm_layer():
    """Test RMSNorm layer."""
    layer = RMSNorm(hidden_size=10, eps=1e-6)

    x = torch.randn(2, 3, 10)
    output = layer(x)

    assert output.shape == x.shape


def test_layernorm_layer():
    """Test LayerNorm layer."""
    layer = LayerNorm(hidden_size=10, eps=1e-6)

    x = torch.randn(2, 3, 10)
    output = layer(x)

    assert output.shape == x.shape


def test_embedding_layer():
    """Test Embedding layer."""
    layer = Embedding(num_embeddings=100, embedding_dim=64)

    input_ids = torch.randint(0, 100, (2, 10))
    output = layer(input_ids)

    assert output.shape == (2, 10, 64)


def test_attention_layer():
    """Test Attention layer."""
    layer = Attention(
        hidden_size=128,
        num_heads=8,
        num_kv_heads=8,
        bias=False,
    )

    x = torch.randn(2, 10, 128)
    output, past_kv = layer(x, use_cache=False)

    assert output.shape == x.shape
    assert past_kv is None


def test_attention_with_cache():
    """Test Attention with KV cache."""
    layer = Attention(
        hidden_size=128,
        num_heads=8,
        num_kv_heads=8,
        bias=False,
    )

    x = torch.randn(2, 10, 128)
    output, past_kv = layer(x, use_cache=True)

    assert output.shape == x.shape
    assert past_kv is not None
    assert len(past_kv) == 2  # (k, v)


def test_attention_gqa():
    """Test Attention with grouped query attention."""
    layer = Attention(
        hidden_size=128,
        num_heads=8,
        num_kv_heads=2,  # GQA: 8/2 = 4 groups
        bias=False,
    )

    x = torch.randn(2, 10, 128)
    output, _ = layer(x, use_cache=False)

    assert output.shape == x.shape


def test_silu_activation():
    """Test SiLU activation."""
    layer = SiLU()

    x = torch.randn(2, 3, 10)
    output = layer(x)

    assert output.shape == x.shape


def test_gelu_activation():
    """Test GELU activation."""
    layer = GELU()

    x = torch.randn(2, 3, 10)
    output = layer(x)

    assert output.shape == x.shape


def test_swiglu_activation():
    """Test SwiGLU activation."""
    layer = SwiGLU(input_dim=128, hidden_dim=256)

    x = torch.randn(2, 10, 128)
    output = layer(x)

    assert output.shape == (2, 10, 256)


def test_layer_backend_injection():
    """Test backend injection for layers."""
    layer = Linear(10, 20)

    # Set backend (None for now)
    layer.set_backend(None)

    # Should work without backend
    x = torch.randn(2, 10)
    output = layer(x)

    assert output.shape == (2, 20)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
