"""
Test Mixture of Experts Implementation
"""

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

import torch
from angstrom_nano import AngstromNanoConfig
from angstrom_nano.model import AngstromNanoSparseMoEBlock


def test_moe_forward():
    print("\n[Testing MoE Forward Pass]")
    config = AngstromNanoConfig()
    moe = AngstromNanoSparseMoEBlock(config)
    
    batch_size, seq_len = 2, 128
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size)
    
    output, router_logits = moe(hidden_states)
    
    print(f"  Input shape: {hidden_states.shape}")
    print(f"  Output shape: {output.shape}")
    print(f"  Router logits shape: {router_logits.shape}")
    assert output.shape == hidden_states.shape
    assert router_logits.shape == (batch_size * seq_len, config.num_local_experts)
    print("  [PASS]")
    
    return output, router_logits


def test_expert_utilization():
    print("\n[Testing Expert Utilization]")
    config = AngstromNanoConfig()
    moe = AngstromNanoSparseMoEBlock(config)
    
    batch_size, seq_len = 4, 256
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size)
    
    _, router_logits = moe(hidden_states)
    
    # Analyze expert selection
    routing_weights = torch.softmax(router_logits, dim=-1)
    top_experts = torch.topk(routing_weights, config.num_experts_per_tok, dim=-1).indices
    
    # Count expert usage
    expert_counts = torch.zeros(config.num_local_experts)
    for i in range(config.num_local_experts):
        expert_counts[i] = (top_experts == i).sum().item()
    
    total_selections = batch_size * seq_len * config.num_experts_per_tok
    expert_percentages = expert_counts / total_selections * 100
    
    print(f"  Total tokens: {batch_size * seq_len}")
    print(f"  Top-k: {config.num_experts_per_tok}")
    print(f"  Total expert selections: {total_selections}")
    print(f"  Expert usage distribution:")
    for i, (count, pct) in enumerate(zip(expert_counts, expert_percentages)):
        print(f"    Expert {i}: {int(count):4d} selections ({pct:5.2f}%)")
    
    # Check balance (no expert should be used more than 25% or less than 5%)
    max_usage = expert_percentages.max().item()
    min_usage = expert_percentages.min().item()
    print(f"  Max usage: {max_usage:.2f}%, Min usage: {min_usage:.2f}%")
    
    if max_usage < 25 and min_usage > 5:
        print("  [PASS] Expert utilization is balanced")
    else:
        print("  [WARNING] Expert utilization may be imbalanced")
    
    return expert_counts


def test_load_balancing_loss():
    print("\n[Testing Load Balancing Loss]")
    config = AngstromNanoConfig()
    moe = AngstromNanoSparseMoEBlock(config)
    
    batch_size, seq_len = 2, 128
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size)
    
    output, router_logits = moe(hidden_states)
    
    # Compute load balancing loss
    routing_weights = torch.softmax(router_logits, dim=-1)
    
    # Average routing probability per expert
    avg_routing_prob = routing_weights.mean(dim=0)  # [num_experts]
    
    # Fraction of tokens routed to each expert
    top_k_indices = torch.topk(routing_weights, config.num_experts_per_tok, dim=-1).indices
    expert_mask = torch.nn.functional.one_hot(top_k_indices, config.num_local_experts).float()
    expert_usage = expert_mask.sum(dim=1).mean(dim=0)  # [num_experts]
    expert_usage = expert_usage / expert_usage.sum()
    
    # Load balancing loss
    aux_loss = (avg_routing_prob * expert_usage).sum() * config.num_local_experts
    
    print(f"  Aux loss: {aux_loss.item():.6f}")
    print(f"  Weighted aux loss: {aux_loss.item() * config.router_aux_loss_coef:.6f}")
    assert aux_loss.item() > 0
    print("  [PASS]")
    
    return aux_loss


def test_gradient_flow():
    print("\n[Testing Gradient Flow Through MoE]")
    config = AngstromNanoConfig()
    moe = AngstromNanoSparseMoEBlock(config)
    
    batch_size, seq_len = 2, 64
    hidden_states = torch.randn(batch_size, seq_len, config.hidden_size, requires_grad=True)
    
    output, router_logits = moe(hidden_states)
    
    # Compute loss and backward
    loss = output.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}")
    
    # Check gradients for experts
    num_experts_with_grad = sum(
        1 for expert in moe.experts 
        if expert.gate_proj.weight.grad is not None
    )
    print(f"  Experts with gradients: {num_experts_with_grad}/{config.num_local_experts}")
    
    assert hidden_states.grad is not None
    assert num_experts_with_grad > 0
    print("  [PASS]")


def main():
    print("=" * 80)
    print("Testing Mixture of Experts Implementation")
    print("=" * 80)
    
    torch.manual_seed(42)
    
    test_moe_forward()
    test_expert_utilization()
    test_load_balancing_loss()
    test_gradient_flow()
    
    print("\n" + "=" * 80)
    print("All MoE tests passed!")
    print("=" * 80)


if __name__ == "__main__":
    main()