"""
Unit and integration test suite for Q-TensorFormer custom Hugging Face model.

Tests:
1. test_forward_pass_shape: Output logits shape == (batch_size, seq_len, vocab_size).
2. test_autoregressive_generation: model.generate() produces finite, valid sequences.
3. test_pid_controller_stability: 100 consecutive entropy steps suppress chattering via hysteresis.
4. test_attention_sink_retention: AdaptiveKVCache retains sink tokens during long generation.
5. test_tensor_train_linear_contraction: Einsum contraction matches reconstructed matrix and supports rank slicing.
6. test_gqa_attention: Grouped-Query Attention preserves head dimensions and causality.
"""

import math
import pytest
import torch

from configuration_qtensorformer import QTensorFormerConfig, factorize_dim
from modeling_qtensorformer import (
    TensorTrainLinear,
    AdaptiveKVCache,
    PIDResourceController,
    QTensorFormerAttention,
    QTensorFormerForCausalLM,
)


@pytest.fixture
def small_config():
    """Lightweight configuration for rapid unit testing."""
    return QTensorFormerConfig(
        vocab_size=1000,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        max_position_embeddings=256,
        tt_ranks=[1, 4, 4, 1],
        attention_sink_size=4,
        kv_cache_window_size=16,
        hysteresis_threshold=0.15,
    )


def test_factorize_dim():
    """Verify that factorize_dim decomposes numbers into exact products."""
    for dim in [128, 256, 512, 2048, 5632]:
        factors = factorize_dim(dim, num_cores=3)
        assert len(factors) == 3
        assert math.prod(factors) == dim


def test_tensor_train_linear_contraction():
    """Verify TTLinear contraction against reconstructed dense matrix."""
    torch.manual_seed(42)
    in_features, out_features = 64, 128
    tt_ranks = [1, 4, 4, 1]

    layer = TensorTrainLinear(
        in_features=in_features,
        out_features=out_features,
        tt_ranks=tt_ranks,
        bias=True,
    )

    x = torch.randn(2, 5, in_features)

    # 1. Direct TT forward contraction
    out_tt = layer(x)
    assert out_tt.shape == (2, 5, out_features)
    assert torch.isfinite(out_tt).all()

    # 2. Reconstructed dense matrix comparison
    w_dense = layer.reconstruct_weight()
    out_dense = torch.matmul(x, w_dense) + layer.bias
    assert torch.allclose(out_tt, out_dense, atol=1e-4)

    # 3. Dynamic rank-slicing test
    out_sliced = layer(x, active_rank=2)
    assert out_sliced.shape == (2, 5, out_features)
    assert torch.isfinite(out_sliced).all()


def test_pid_controller_stability():
    """Simulate 100 consecutive random entropy tokens and ensure chattering is suppressed."""
    torch.manual_seed(42)
    controller = PIDResourceController(
        kp=0.1,
        ki=0.01,
        kd=0.05,
        hysteresis_threshold=0.15,
        target_budget=1.0,
    )

    rank_history = []
    chattering_suppression_count = 0

    for step in range(100):
        # Fluctuating cost around the target budget
        cost = 1.0 + 0.3 * math.sin(step * 0.2) + (torch.rand(1).item() - 0.5) * 0.1
        info = controller(measured_cost=cost)
        rank_history.append(info["active_rank"])
        if info["chattering_suppressed"]:
            chattering_suppression_count += 1

    # Ensure rank didn't change at every single step (hysteresis worked)
    switches = sum(1 for i in range(1, len(rank_history)) if rank_history[i] != rank_history[i - 1])
    assert switches < 50, f"Too many rapid switches ({switches}/100), hysteresis failed!"
    assert chattering_suppression_count > 0, "No chattering was suppressed!"


def test_attention_sink_retention():
    """Verify that AdaptiveKVCache permanently preserves the first 4 tokens."""
    sink_size = 4
    window_size = 8
    cache = AdaptiveKVCache(
        sink_size=sink_size,
        window_size=window_size,
        quantize_cold=True,
        quant_bits=4,
    )

    bsz, num_heads, head_dim = 1, 2, 16

    # 1. Fill initial tokens (including sink)
    init_k = torch.arange(10, dtype=torch.float32).view(1, 1, 10, 1).repeat(bsz, num_heads, 1, head_dim)
    init_v = init_k.clone()
    k_out, v_out = cache.update(init_k, init_v, layer_idx=0)

    # 2. Append 20 more tokens one-by-one (exceeding window_size + sink_size)
    for t in range(10, 30):
        step_k = torch.tensor([[[[float(t)] * head_dim]]]).repeat(bsz, num_heads, 1, 1)
        step_v = step_k.clone()
        k_out, v_out = cache.update(step_k, step_v, layer_idx=0)

    # Verify sink tokens are permanently preserved at the beginning
    first_tokens = k_out[:, :, :sink_size, 0]
    expected_sink = torch.arange(sink_size, dtype=torch.float32).view(1, 1, sink_size).repeat(bsz, num_heads, 1)
    assert torch.allclose(first_tokens, expected_sink), "Sink tokens were lost during eviction!"


def test_gqa_attention(small_config):
    """Verify Grouped-Query Attention dimensions and caching."""
    layer = QTensorFormerAttention(small_config, layer_idx=0)
    x = torch.randn(2, 8, small_config.hidden_size)

    # Prefill forward
    out, _, past = layer(x, use_cache=True)
    assert out.shape == (2, 8, small_config.hidden_size)
    assert torch.isfinite(out).all()


def test_forward_pass_shape(small_config):
    """Verify model forward pass output shape matches (B, S, V)."""
    model = QTensorFormerForCausalLM(small_config)
    model.eval()

    bsz, seq_len = 2, 16
    input_ids = torch.randint(0, small_config.vocab_size, (bsz, seq_len))

    with torch.no_grad():
        outputs = model(input_ids)
        logits = outputs.logits

    assert logits.shape == (bsz, seq_len, small_config.vocab_size)
    assert torch.isfinite(logits).all(), "Logits contain NaN or Inf values!"


def test_autoregressive_generation(small_config):
    """Verify standard Hugging Face model.generate() works seamlessly."""
    model = QTensorFormerForCausalLM(small_config)
    model.eval()

    prompt_ids = torch.randint(0, small_config.vocab_size, (1, 4))
    gen_tokens = 6

    with torch.no_grad():
        generated = model.generate(
            prompt_ids,
            max_new_tokens=gen_tokens,
            do_sample=False,
            use_cache=True,
        )

    assert generated.shape == (1, 4 + gen_tokens)
    assert (generated >= 0).all() and (generated < small_config.vocab_size).all()