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

"""
Tests for training infrastructure:
  - Loss functions (CE, consistency)
  - Optimizer param groups
  - LR schedule (warmup + cosine)
  - Gradient flow and clipping
  - Training step smoke test
"""

import math
import torch
import pytest

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.training.loss import compute_ce_loss, compute_consistency_loss
from axiom1.training.optimizer import build_optimizer, get_lr, get_ponder_tau


class TestCELoss:
    """Tests for cross-entropy loss."""

    def test_loss_finite(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 32))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        assert torch.isfinite(loss), f"CE loss is not finite: {loss.item()}"

    def test_loss_positive(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 32))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        assert loss.item() > 0, f"CE loss should be positive, got {loss.item()}"

    def test_initial_loss_near_log_vocab(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        """Initial loss should be ≈ log(vocab_size) ≈ 10.37."""
        ids = torch.randint(0, small_cfg.vocab_size, (4, 32))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        expected = math.log(small_cfg.vocab_size)
        assert abs(loss.item() - expected) < 3.0, \
            f"Initial loss={loss.item():.2f}, expected ~{expected:.2f} (±3.0)"


class TestConsistencyLoss:
    """Tests for loop consistency loss."""

    def test_consistency_loss_computes(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        _, intermediates = small_model(ids, return_intermediates=True)
        loss = compute_consistency_loss(
            small_model, intermediates, ids, start_loop=1
        )
        assert torch.isfinite(loss), f"Consistency loss not finite: {loss.item()}"

    def test_consistency_loss_nonnegative(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        _, intermediates = small_model(ids, return_intermediates=True)
        loss = compute_consistency_loss(
            small_model, intermediates, ids, start_loop=1
        )
        assert loss.item() >= 0, f"Consistency loss should be >= 0, got {loss.item()}"


class TestOptimizer:
    """Tests for optimizer and param groups."""

    def test_optimizer_builds(self, model: AXIOM1, cfg: AXIOM1Config) -> None:
        optimizer = build_optimizer(model, cfg)
        assert optimizer is not None, "build_optimizer returned None"

    def test_param_groups_exist(self, model: AXIOM1, cfg: AXIOM1Config) -> None:
        optimizer = build_optimizer(model, cfg)
        assert len(optimizer.param_groups) == 2, \
            f"Expected 2 param groups (decay/no-decay), got {len(optimizer.param_groups)}"

    def test_no_decay_group_has_zero_wd(self, model: AXIOM1, cfg: AXIOM1Config) -> None:
        optimizer = build_optimizer(model, cfg)
        # Find the no-decay group
        for pg in optimizer.param_groups:
            if pg['weight_decay'] == 0.0:
                assert len(pg['params']) > 0, "No-decay group should have parameters"
                return
        pytest.fail("No param group with weight_decay=0.0 found")

    def test_beta2_is_095(self, model: AXIOM1, cfg: AXIOM1Config) -> None:
        """Rule 10: AdamW β2 must be 0.95, not 0.999."""
        optimizer = build_optimizer(model, cfg)
        for pg in optimizer.param_groups:
            betas = pg.get('betas', (0.9, 0.999))
            assert betas[1] == 0.95, \
                f"β2 must be 0.95 (Rule 10), got {betas[1]}"


class TestLRSchedule:
    """Tests for learning rate schedule."""

    def test_lr_at_zero_is_zero(self) -> None:
        """LR at step 0 should be 0 (start of warmup)."""
        lr = get_lr(0, warmup_steps=1000, total_steps=50000,
                    peak_lr=3e-4, min_lr=3e-5)
        assert lr < 1e-5, f"LR at step 0 should be ~0, got {lr}"

    def test_lr_at_warmup_end_is_peak(self) -> None:
        lr = get_lr(1000, warmup_steps=1000, total_steps=50000,
                    peak_lr=3e-4, min_lr=3e-5)
        assert abs(lr - 3e-4) < 1e-5, \
            f"LR at warmup end should be 3e-4, got {lr}"

    def test_lr_decays_after_warmup(self) -> None:
        lr_mid = get_lr(25000, warmup_steps=1000, total_steps=50000,
                        peak_lr=3e-4, min_lr=3e-5)
        lr_peak = get_lr(1000, warmup_steps=1000, total_steps=50000,
                         peak_lr=3e-4, min_lr=3e-5)
        assert lr_mid < lr_peak, \
            f"LR should decay after warmup: mid={lr_mid} >= peak={lr_peak}"

    def test_lr_at_end_near_min(self) -> None:
        lr = get_lr(49999, warmup_steps=1000, total_steps=50000,
                    peak_lr=3e-4, min_lr=3e-5)
        assert lr < 1e-4, f"LR at end should be near min_lr, got {lr}"


class TestGradientFlow:
    """Tests for gradient flow and clipping."""

    def test_gradients_flow(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        loss.backward()

        has_grad = False
        for p in small_model.parameters():
            if p.grad is not None and p.grad.abs().sum() > 0:
                has_grad = True
                break
        assert has_grad, "No gradients flowing through the model"

    def test_gradients_finite(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        loss.backward()

        for name, p in small_model.named_parameters():
            if p.grad is not None:
                assert torch.isfinite(p.grad).all(), \
                    f"Non-finite gradient in {name}"

    def test_grad_clipping(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        """Rule 7: Gradient clip at 1.0 every step."""
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        loss.backward()

        grad_norm = torch.nn.utils.clip_grad_norm_(
            small_model.parameters(), 1.0
        )
        assert torch.isfinite(grad_norm), f"Grad norm not finite: {grad_norm}"

    def test_training_step(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        """Full training step: forward + backward + optimizer step."""
        optimizer = build_optimizer(small_model, small_cfg)
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))

        logits = small_model(ids)
        loss = compute_ce_loss(logits, ids)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(small_model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

        # Verify loss decreased on second forward
        logits2 = small_model(ids)
        loss2 = compute_ce_loss(logits2, ids)
        # Not guaranteed to decrease in 1 step, but should be finite
        assert torch.isfinite(loss2), f"Loss after step not finite: {loss2.item()}"
