"""
Test Complete AngstromNano Model
"""

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

import torch
from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM


def test_model_initialization():
    print("\n[Testing Model Initialization]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"  Total parameters: {total_params:,} ({total_params/1e9:.2f}B)")
    print(f"  Trainable parameters: {trainable_params:,}")
    print(f"  Config estimate: {config.estimate_parameters()['total_billions']:.2f}B")
    print("  [PASS]")
    
    return model


def test_forward_pass():
    print("\n[Testing Forward Pass]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    batch_size, seq_len = 2, 128
    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    
    outputs = model(input_ids, labels=labels)
    
    print(f"  Input shape: {input_ids.shape}")
    print(f"  Logits shape: {outputs['logits'].shape}")
    print(f"  Loss: {outputs['loss'].item():.4f}")
    print(f"  Aux loss: {outputs['aux_loss'].item():.6f}")
    
    assert outputs['logits'].shape == (batch_size, seq_len, config.vocab_size)
    assert outputs['loss'] is not None
    print("  [PASS]")


def test_generation():
    print("\n[Testing Text Generation]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    # Simple prompt
    prompt = torch.tensor([[1, 100, 200, 300]])  # BOS + 3 tokens
    
    generated = model.generate(
        prompt,
        max_new_tokens=20,
        temperature=1.0,
    )
    
    print(f"  Prompt length: {prompt.shape[1]}")
    print(f"  Generated length: {generated.shape[1]}")
    print(f"  Generated tokens: {generated[0].tolist()[:10]}...")
    
    assert generated.shape[1] > prompt.shape[1]
    print("  [PASS]")


def test_kv_cache():
    print("\n[Testing KV Cache]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    batch_size, seq_len = 1, 16
    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    
    # First forward pass with cache
    outputs1 = model(input_ids, use_cache=True)
    past_kv = outputs1['past_key_values']
    
    print(f"  Cached layers: {len(past_kv)}")
    print(f"  Cached K shape: {past_kv[0][0].shape}")
    print(f"  Cached V shape: {past_kv[0][1].shape}")
    
    # Second forward pass using cache
    next_token = torch.randint(0, config.vocab_size, (batch_size, 1))
    outputs2 = model(next_token, past_key_values=past_kv, use_cache=True)
    
    print(f"  New cached K shape: {outputs2['past_key_values'][0][0].shape}")
    
    assert len(past_kv) == config.num_hidden_layers
    assert outputs2['past_key_values'][0][0].shape[2] == seq_len + 1
    print("  [PASS]")


def test_gradient_flow():
    print("\n[Testing Gradient Flow]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    batch_size, seq_len = 2, 32
    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    
    outputs = model(input_ids, labels=labels)
    loss = outputs['loss']
    
    loss.backward()
    
    # Check gradients in different parts
    embed_grad = model.model.embed_tokens.weight.grad
    first_layer_grad = model.model.layers[0].self_attn.q_proj.weight.grad
    last_layer_grad = model.model.layers[-1].self_attn.q_proj.weight.grad
    lm_head_grad = model.lm_head.weight.grad
    
    print(f"  Loss: {loss.item():.4f}")
    print(f"  Embedding grad norm: {embed_grad.norm().item():.6f}")
    print(f"  First layer grad norm: {first_layer_grad.norm().item():.6f}")
    print(f"  Last layer grad norm: {last_layer_grad.norm().item():.6f}")
    print(f"  LM head grad norm: {lm_head_grad.norm().item():.6f}")
    
    assert all(g is not None for g in [embed_grad, first_layer_grad, last_layer_grad, lm_head_grad])
    print("  [PASS]")


def test_memory_usage():
    print("\n[Testing Memory Usage]")
    config = AngstromNanoConfig()
    model = AngstromNanoForCausalLM(config)
    
    # Calculate model size
    param_size = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
    total_size = param_size + buffer_size
    
    print(f"  Parameter memory: {param_size / 1e9:.2f} GB")
    print(f"  Buffer memory: {buffer_size / 1e6:.2f} MB")
    print(f"  Total model size: {total_size / 1e9:.2f} GB")
    
    # Expected: ~8GB for FP32, ~4GB for FP16
    expected_fp32 = config.estimate_parameters()['estimated_size_fp32_gb']
    print(f"  Expected FP32 size: {expected_fp32:.2f} GB")
    
    assert abs(total_size / 1e9 - expected_fp32) < 1.0  # Within 1GB tolerance
    print("  [PASS]")


def main():
    print("=" * 80)
    print("Testing Complete AngstromNano Model")
    print("=" * 80)
    
    torch.manual_seed(42)
    
    test_model_initialization()
    test_forward_pass()
    test_generation()
    test_kv_cache()
    test_gradient_flow()
    test_memory_usage()
    
    print("\n" + "=" * 80)
    print("All model tests passed!")
    print("=" * 80)


if __name__ == "__main__":
    main()