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

"""
Tests for AXIOM1Config — verifying default values and derived properties.
"""

from axiom1.model.config import AXIOM1Config


class TestAXIOM1Config:
    """Tests for the AXIOM1Config dataclass."""

    def test_default_d_model(self, cfg: AXIOM1Config) -> None:
        assert cfg.d_model == 512, f"Expected d_model=512, got {cfg.d_model}"

    def test_default_heads(self, cfg: AXIOM1Config) -> None:
        assert cfg.n_q_heads == 8, f"Expected n_q_heads=8, got {cfg.n_q_heads}"
        assert cfg.n_kv_heads == 2, f"Expected n_kv_heads=2, got {cfg.n_kv_heads}"

    def test_head_dim_derived(self, cfg: AXIOM1Config) -> None:
        expected = cfg.d_model // cfg.n_q_heads
        assert cfg.head_dim == expected, \
            f"head_dim should be d_model/n_q_heads={expected}, got {cfg.head_dim}"

    def test_d_ff_derived(self, cfg: AXIOM1Config) -> None:
        expected = int(cfg.d_model * 8 / 3)
        assert cfg.d_ff == expected, \
            f"d_ff should be int(d_model*8/3)={expected}, got {cfg.d_ff}"

    def test_default_num_loops(self, cfg: AXIOM1Config) -> None:
        assert cfg.num_loops == 12, f"Expected num_loops=12, got {cfg.num_loops}"

    def test_default_max_loops(self, cfg: AXIOM1Config) -> None:
        assert cfg.max_loops == 16, f"Expected max_loops=16, got {cfg.max_loops}"

    def test_default_vocab_size(self, cfg: AXIOM1Config) -> None:
        assert cfg.vocab_size == 32000, f"Expected vocab_size=32000, got {cfg.vocab_size}"

    def test_default_max_seq_len(self, cfg: AXIOM1Config) -> None:
        assert cfg.max_seq_len == 2048, f"Expected max_seq_len=2048, got {cfg.max_seq_len}"

    def test_default_drop_path(self, cfg: AXIOM1Config) -> None:
        assert cfg.drop_path == 0.10, f"Expected drop_path=0.10, got {cfg.drop_path}"

    def test_default_init_std(self, cfg: AXIOM1Config) -> None:
        assert cfg.init_std == 0.02, f"Expected init_std=0.02, got {cfg.init_std}"

    def test_default_rope_base(self, cfg: AXIOM1Config) -> None:
        assert cfg.rope_base == 10000, f"Expected rope_base=10000, got {cfg.rope_base}"

    def test_default_grad_clip(self, cfg: AXIOM1Config) -> None:
        assert cfg.grad_clip == 1.0, f"Expected grad_clip=1.0, got {cfg.grad_clip}"

    def test_special_token_ids(self, cfg: AXIOM1Config) -> None:
        assert cfg.pad_token_id == 0, f"pad_token_id should be 0"
        assert cfg.bos_token_id == 1, f"bos_token_id should be 1"
        assert cfg.eos_token_id == 2, f"eos_token_id should be 2"

    def test_gqa_ratio(self, cfg: AXIOM1Config) -> None:
        """GQA should use 4x repetition (8q / 2kv = 4)."""
        ratio = cfg.n_q_heads // cfg.n_kv_heads
        assert ratio == 4, f"GQA ratio should be 4, got {ratio}"

    def test_bfloat16_rule(self) -> None:
        """Rule 3: Never use float16. Verify this is documented in config."""
        # The precision rule is enforced in training code, not as a config field.
        # This test verifies the model's default dtype is compatible.
        import torch
        assert torch.bfloat16 is not None, "bfloat16 dtype should exist"
