"""
Test Core Transformer Components
"""

import sys
sys.path.insert(0, '.')

import torch
from angstrom_nano import AngstromNanoConfig
from angstrom_nano.model import (
    AngstromNanoRMSNorm,
    AngstromNanoRotaryEmbedding,
    AngstromNanoAttention,
    AngstromNanoMLP
)


def test_rms_norm():
    print("\n[Testing RMSNorm]")
    config = AngstromNanoConfig()
    norm = AngstromNanoRMSNorm(config.hidden_size, config.rms_norm_eps)
    
    x = torch.randn(2, 512, config.hidden_size)
    out = norm(x)
    
    print(f"  Input shape: {x.shape}")
    print(f"  Output shape: {out.shape}")
    print(f"  Output mean: {out.mean():.6f}, std: {out.std():.6f}")
    assert out.shape == x.shape
    print("  [PASS]")


def test_rotary_embedding():
    print("\n[Testing Rotary Embedding]")
    config = AngstromNanoConfig()
    rope = AngstromNanoRotaryEmbedding(
        config.head_dim,
        config.max_position_embeddings,
        config.rope_theta
    )
    
    batch_size, seq_len = 2, 512
    x = torch.randn(batch_size, config.num_attention_heads, seq_len, config.head_dim)
    position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
    
    cos, sin = rope(x, position_ids)
    
    print(f"  Input shape: {x.shape}")
    print(f"  Position IDs shape: {position_ids.shape}")
    print(f"  Cos shape: {cos.shape}")
    print(f"  Sin shape: {sin.shape}")
    assert cos.shape == (batch_size, seq_len, config.head_dim)
    assert sin.shape == (batch_size, seq_len, config.head_dim)
    print("  [PASS]")


def test_attention():
    print("\n[Testing Multi-Head Attention]")
    config = AngstromNanoConfig()
    attn = AngstromNanoAttention(config, layer_idx=0)
    
    batch_size, seq_len = 2, 512
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size)
    position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
    
    # Create causal mask
    attention_mask = torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1)
    attention_mask = attention_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]
    
    output, past_kv = attn(
        hidden_states,
        attention_mask=attention_mask,
        position_ids=position_ids,
        use_cache=True
    )
    
    print(f"  Input shape: {hidden_states.shape}")
    print(f"  Output shape: {output.shape}")
    print(f"  KV cache shapes: {past_kv[0].shape}, {past_kv[1].shape}")
    assert output.shape == hidden_states.shape
    assert past_kv[0].shape[2] == seq_len  # cached sequence length
    print("  [PASS]")


def test_mlp():
    print("\n[Testing Feed-Forward Network]")
    config = AngstromNanoConfig()
    mlp = AngstromNanoMLP(config)
    
    batch_size, seq_len = 2, 512
    x = torch.randn(batch_size, seq_len, config.hidden_size)
    
    output = mlp(x)
    
    print(f"  Input shape: {x.shape}")
    print(f"  Output shape: {output.shape}")
    assert output.shape == x.shape
    print("  [PASS]")


def test_gradient_flow():
    print("\n[Testing Gradient Flow]")
    config = AngstromNanoConfig()
    attn = AngstromNanoAttention(config, layer_idx=0)
    mlp = AngstromNanoMLP(config)
    
    batch_size, seq_len = 2, 128
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size, requires_grad=True)
    position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
    
    # Forward pass
    attn_out, _ = attn(hidden_states, position_ids=position_ids)
    mlp_out = mlp(attn_out)
    
    # Compute loss and backward
    loss = mlp_out.mean()
    loss.backward()
    
    print(f"  Loss: {loss.item():.6f}")
    print(f"  Input gradient exists: {hidden_states.grad is not None}")
    print(f"  Input gradient norm: {hidden_states.grad.norm().item():.6f}")
    assert hidden_states.grad is not None
    print("  [PASS]")


def main():
    print("=" * 80)
    print("Testing Core Transformer Components")
    print("=" * 80)
    
    torch.manual_seed(42)
    
    test_rms_norm()
    test_rotary_embedding()
    test_attention()
    test_mlp()
    test_gradient_flow()
    
    print("\n" + "=" * 80)
    print("All tests passed!")
    print("=" * 80)


if __name__ == "__main__":
    main()