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

"""
Tests for inference modules:
  - Autoregressive generation
  - PerStepKVCache
  - AXIOM1_Offloaded (CPU offloading)
"""

import torch
import pytest

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.inference.kv_cache import PerStepKVCache
from axiom1.inference.generate import generate
from axiom1.inference.offload import AXIOM1_Offloaded


class TestPerStepKVCache:
    """Tests for per-loop-step KV cache."""

    def test_init_empty(self) -> None:
        cache = PerStepKVCache(num_steps=12, n_kv_heads=2, head_dim=64)
        assert cache.seq_len == 0, "New cache should have seq_len=0"

    def test_update_grows_cache(self) -> None:
        cache = PerStepKVCache(num_steps=4, n_kv_heads=2, head_dim=64)
        k = torch.randn(1, 2, 1, 64)
        v = torch.randn(1, 2, 1, 64)
        full_k, full_v = cache.update(0, k, v)
        assert full_k.shape == (1, 2, 1, 64), f"Unexpected shape: {full_k.shape}"
        assert cache.seq_len == 1

        k2 = torch.randn(1, 2, 1, 64)
        v2 = torch.randn(1, 2, 1, 64)
        full_k2, full_v2 = cache.update(0, k2, v2)
        assert full_k2.shape == (1, 2, 2, 64), f"Cache should grow: {full_k2.shape}"

    def test_separate_per_step(self) -> None:
        """Each loop step should have its own independent cache."""
        cache = PerStepKVCache(num_steps=4, n_kv_heads=2, head_dim=64)
        k0 = torch.randn(1, 2, 1, 64)
        v0 = torch.randn(1, 2, 1, 64)
        cache.update(0, k0, v0)

        k1 = torch.randn(1, 2, 1, 64)
        v1 = torch.randn(1, 2, 1, 64)
        cache.update(1, k1, v1)

        assert cache.k_cache[0] is not None, "Step 0 cache should exist"
        assert cache.k_cache[1] is not None, "Step 1 cache should exist"
        assert cache.k_cache[2] is None, "Step 2 cache should be None"

    def test_reset(self) -> None:
        cache = PerStepKVCache(num_steps=4, n_kv_heads=2, head_dim=64)
        cache.update(0, torch.randn(1, 2, 1, 64), torch.randn(1, 2, 1, 64))
        cache.reset()
        assert cache.seq_len == 0, "Cache should be empty after reset"
        assert all(c is None for c in cache.k_cache), "All caches should be None"


class TestGeneration:
    """Tests for autoregressive text generation."""

    def test_generate_produces_tokens(self) -> None:
        cfg = AXIOM1Config()
        cfg.num_loops = 2
        model = AXIOM1(cfg)
        model.eval()

        prompt = torch.tensor([[1, 100, 200]])  # BOS + 2 tokens
        with torch.no_grad():
            output = generate(model, prompt, cfg, max_new_tokens=5)

        assert output.shape[1] > prompt.shape[1], \
            f"Generation should produce more tokens: {output.shape[1]} <= {prompt.shape[1]}"

    def test_generate_no_nan(self) -> None:
        cfg = AXIOM1Config()
        cfg.num_loops = 2
        model = AXIOM1(cfg)
        model.eval()

        prompt = torch.tensor([[1, 500]])
        with torch.no_grad():
            output = generate(model, prompt, cfg, max_new_tokens=3)
        assert not torch.isnan(output.float()).any(), "Generated IDs contain NaN"

    def test_generate_respects_max_tokens(self) -> None:
        cfg = AXIOM1Config()
        cfg.num_loops = 2
        cfg.eos_token_id = -1  # Disable EOS stopping
        model = AXIOM1(cfg)
        model.eval()

        prompt = torch.tensor([[1]])
        max_new = 10
        with torch.no_grad():
            output = generate(model, prompt, cfg, max_new_tokens=max_new)
        assert output.shape[1] <= 1 + max_new, \
            f"Output length {output.shape[1]} exceeds prompt(1) + max_new({max_new})"


class TestCPUOffloading:
    """Tests for AirLLM-style CPU offloading."""

    def test_offloaded_model_on_cpu(self) -> None:
        cfg = AXIOM1Config()
        cfg.num_loops = 2
        model = AXIOM1(cfg)
        offloaded = AXIOM1_Offloaded(model, device='cpu')
        # All parameters should be on CPU
        for p in offloaded.model.parameters():
            assert p.device.type == 'cpu', "Offloaded model params should be on CPU"
