# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: test/test_components.py
# ─────────────────────────────────────────────────────────

"""
Unit tests for model/components.py:
  RMSNorm, apply_rope, GroupedQueryAttention, SwiGLU_FFN, DropPath
"""

import math
import torch
import pytest

from axiom1.model.config import AXIOM1Config
from axiom1.model.components import (
    RMSNorm, apply_rope, GroupedQueryAttention, SwiGLU_FFN, DropPath
)


class TestRMSNorm:
    """Tests for RMSNorm."""

    def test_output_shape(self) -> None:
        norm = RMSNorm(512)
        x = torch.randn(2, 16, 512)
        out = norm(x)
        assert out.shape == x.shape, f"Shape mismatch: {out.shape} != {x.shape}"

    def test_no_nan(self) -> None:
        norm = RMSNorm(512)
        x = torch.randn(2, 16, 512)
        out = norm(x)
        assert not torch.isnan(out).any(), "RMSNorm output contains NaN"

    def test_weight_init_ones(self) -> None:
        norm = RMSNorm(512)
        assert torch.allclose(norm.weight, torch.ones(512)), \
            "RMSNorm weight should be initialized to ones"

    def test_normalization_effect(self) -> None:
        """Output RMS should be approximately 1."""
        norm = RMSNorm(64)
        x = torch.randn(4, 8, 64) * 10  # large inputs
        out = norm(x)
        rms = out.pow(2).mean(dim=-1).sqrt()
        assert rms.mean().item() == pytest.approx(1.0, abs=0.5), \
            f"RMS after normalization should be ~1.0, got {rms.mean().item()}"


class TestRoPE:
    """Tests for Rotary Position Embeddings."""

    def test_output_shape(self) -> None:
        q = torch.randn(2, 8, 16, 64)  # batch, heads, seq, head_dim
        k = torch.randn(2, 8, 16, 64)
        q_rot, k_rot = apply_rope(q, k, seq_len=16, head_dim=64)
        assert q_rot.shape == q.shape, f"q shape mismatch: {q_rot.shape} != {q.shape}"
        assert k_rot.shape == k.shape, f"k shape mismatch: {k_rot.shape} != {k.shape}"

    def test_no_nan(self) -> None:
        q = torch.randn(2, 8, 16, 64)
        k = torch.randn(2, 8, 16, 64)
        q_rot, k_rot = apply_rope(q, k, seq_len=16, head_dim=64)
        assert not torch.isnan(q_rot).any(), "RoPE q output contains NaN"
        assert not torch.isnan(k_rot).any(), "RoPE k output contains NaN"

    def test_modifies_values(self) -> None:
        """RoPE should modify the input tensors (not be identity)."""
        q = torch.randn(1, 4, 8, 64)
        k = torch.randn(1, 4, 8, 64)
        q_rot, k_rot = apply_rope(q, k, seq_len=8, head_dim=64)
        assert not torch.allclose(q, q_rot, atol=1e-3), \
            "RoPE should modify q values"


class TestGroupedQueryAttention:
    """Tests for GQA (8q / 2kv)."""

    def test_output_shape(self, cfg: AXIOM1Config) -> None:
        gqa = GroupedQueryAttention(cfg.d_model, cfg.n_q_heads, cfg.n_kv_heads)
        x = torch.randn(2, 16, cfg.d_model)
        out = gqa(x)
        assert out.shape == x.shape, f"GQA output shape {out.shape} != input {x.shape}"

    def test_no_nan(self, cfg: AXIOM1Config) -> None:
        gqa = GroupedQueryAttention(cfg.d_model, cfg.n_q_heads, cfg.n_kv_heads)
        x = torch.randn(2, 16, cfg.d_model)
        out = gqa(x)
        assert not torch.isnan(out).any(), "GQA output contains NaN"

    def test_no_bias(self, cfg: AXIOM1Config) -> None:
        """Rule: No bias in any Linear inside LoopBlock."""
        gqa = GroupedQueryAttention(cfg.d_model, cfg.n_q_heads, cfg.n_kv_heads)
        for name, module in gqa.named_modules():
            if isinstance(module, torch.nn.Linear):
                assert module.bias is None, \
                    f"GQA Linear '{name}' has bias — must be bias=False"

    def test_projection_dimensions(self, cfg: AXIOM1Config) -> None:
        gqa = GroupedQueryAttention(cfg.d_model, cfg.n_q_heads, cfg.n_kv_heads)
        assert gqa.q_proj.out_features == cfg.d_model, "q_proj out != d_model"
        assert gqa.k_proj.out_features == cfg.n_kv_heads * cfg.head_dim, \
            "k_proj out != n_kv_heads * head_dim"
        assert gqa.v_proj.out_features == cfg.n_kv_heads * cfg.head_dim, \
            "v_proj out != n_kv_heads * head_dim"


class TestSwiGLU_FFN:
    """Tests for SwiGLU Feed-Forward Network."""

    def test_output_shape(self, cfg: AXIOM1Config) -> None:
        ffn = SwiGLU_FFN(cfg.d_model, cfg.d_ff)
        x = torch.randn(2, 16, cfg.d_model)
        out = ffn(x)
        assert out.shape == x.shape, f"FFN output shape {out.shape} != input {x.shape}"

    def test_no_nan(self, cfg: AXIOM1Config) -> None:
        ffn = SwiGLU_FFN(cfg.d_model, cfg.d_ff)
        x = torch.randn(2, 16, cfg.d_model)
        out = ffn(x)
        assert not torch.isnan(out).any(), "SwiGLU output contains NaN"

    def test_no_bias(self, cfg: AXIOM1Config) -> None:
        ffn = SwiGLU_FFN(cfg.d_model, cfg.d_ff)
        for name, module in ffn.named_modules():
            if isinstance(module, torch.nn.Linear):
                assert module.bias is None, \
                    f"SwiGLU Linear '{name}' has bias — must be bias=False"

    def test_d_inner(self, cfg: AXIOM1Config) -> None:
        """d_inner should be int(d_ff * 2/3)."""
        ffn = SwiGLU_FFN(cfg.d_model, cfg.d_ff)
        expected = int(cfg.d_ff * 2 / 3)
        assert ffn.gate.out_features == expected, \
            f"d_inner should be {expected}, got {ffn.gate.out_features}"


class TestDropPath:
    """Tests for DropPath regularization."""

    def test_train_mode(self) -> None:
        dp = DropPath(0.5)
        dp.train()
        x = torch.ones(100, 4, 512)
        out = dp(x)
        assert out.shape == x.shape, f"DropPath shape mismatch"

    def test_eval_mode_identity(self) -> None:
        """DropPath should be identity during eval."""
        dp = DropPath(0.5)
        dp.eval()
        x = torch.randn(2, 8, 512)
        out = dp(x)
        assert torch.allclose(out, x), "DropPath should be identity in eval mode"

    def test_zero_drop_rate(self) -> None:
        """DropPath with p=0 should always be identity."""
        dp = DropPath(0.0)
        dp.train()
        x = torch.randn(2, 8, 512)
        out = dp(x)
        assert torch.allclose(out, x), "DropPath(0.0) should be identity"
