"""
Test script for DeepSeek-V4 open reproduction.
Validates that all architecture components work correctly with the nano config.
"""

import sys
import json
import torch
import torch.nn as nn

sys.path.insert(0, "/app/open-deepseek-v4")

from configuration_deepseek_v4 import DeepSeekV4Config
from modeling_deepseek_v4 import (
    DeepSeekV4ForCausalLM,
    DeepSeekV4Model,
    DeepSeekV4RMSNorm,
    DeepSeekV4Attention,
    DeepSeekV4MoE,
    DeepSeekV4Gate,
    DeepSeekV4Expert,
    DeepSeekV4Block,
    HyperConnection,
    HyperConnectionHead,
    KVCompressor,
    precompute_freqs_cis,
    apply_rotary_emb,
)


def count_params(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable


def test_config():
    print("=" * 60)
    print("TEST 1: Configuration loading")
    print("=" * 60)
    
    # Load from JSON
    with open("/app/open-deepseek-v4/configs/config_nano.json") as f:
        config_dict = json.load(f)
    
    config = DeepSeekV4Config(**config_dict)
    print(f"  model_type: {config.model_type}")
    print(f"  hidden_size: {config.hidden_size}")
    print(f"  num_hidden_layers: {config.num_hidden_layers}")
    print(f"  num_attention_heads: {config.num_attention_heads}")
    print(f"  n_routed_experts: {config.n_routed_experts}")
    print(f"  hc_mult: {config.hc_mult}")
    print(f"  scoring_func: {config.scoring_func}")
    print(f"  compress_ratios: {config.compress_ratios}")
    print("  ✅ Config loaded successfully\n")
    return config


def test_rmsnorm(config):
    print("=" * 60)
    print("TEST 2: RMSNorm")
    print("=" * 60)
    
    norm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
    x = torch.randn(2, 8, config.hidden_size)
    y = norm(x)
    assert y.shape == x.shape, f"Shape mismatch: {y.shape} != {x.shape}"
    print(f"  Input: {x.shape}, Output: {y.shape}")
    print("  ✅ RMSNorm works\n")


def test_rope(config):
    print("=" * 60)
    print("TEST 3: RoPE with YaRN")
    print("=" * 60)
    
    seq_len = 128
    freqs = precompute_freqs_cis(
        config.qk_rope_head_dim,
        seq_len,
        config.rope_theta,
        config.rope_scaling,
    )
    print(f"  freqs_cis shape: {freqs.shape}")
    print(f"  dtype: {freqs.dtype}")
    
    # Test apply
    x = torch.randn(2, seq_len, config.num_attention_heads, config.qk_rope_head_dim)
    y = apply_rotary_emb(x, freqs)
    assert y.shape == x.shape, f"Shape mismatch: {y.shape} != {x.shape}"
    print(f"  Input: {x.shape}, Output: {y.shape}")
    print("  ✅ RoPE works\n")


def test_kv_compressor(config):
    print("=" * 60)
    print("TEST 4: KV Compressor")
    print("=" * 60)
    
    for ratio in [4, 128]:
        comp = KVCompressor(config, compress_ratio=ratio, head_dim=config.head_dim)
        x = torch.randn(2, 256, config.hidden_size)
        freqs = precompute_freqs_cis(
            config.qk_rope_head_dim, 256, config.rope_theta, config.rope_scaling
        )
        out = comp(x, freqs)
        if out is not None:
            expected_len = 256 // ratio
            print(f"  ratio={ratio}: input {x.shape} -> compressed {out.shape} (expected seq_len={expected_len})")
            assert out.size(1) == expected_len, f"Expected {expected_len}, got {out.size(1)}"
        else:
            print(f"  ratio={ratio}: None (sequence too short)")
    
    print("  ✅ KV Compressor works\n")


def test_hyper_connections(config):
    print("=" * 60)
    print("TEST 5: Hyper-Connections (mHC)")
    print("=" * 60)
    
    hc = HyperConnection(config)
    x = torch.randn(2, 16, config.hc_mult, config.hidden_size)
    
    # Pre-connection
    y, post_w, comb = hc.pre(x)
    print(f"  Pre: {x.shape} -> {y.shape}")
    print(f"  Post weights: {post_w.shape}")
    print(f"  Combination matrix: {comb.shape}")
    assert y.shape == (2, 16, config.hidden_size)
    assert post_w.shape == (2, 16, config.hc_mult)
    assert comb.shape == (2, 16, config.hc_mult, config.hc_mult)
    
    # Simulate sublayer output
    sublayer_out = torch.randn(2, 16, config.hidden_size)
    
    # Post-connection
    z = hc.post(sublayer_out, x, post_w, comb)
    print(f"  Post: sublayer {sublayer_out.shape} + residual {x.shape} -> {z.shape}")
    assert z.shape == x.shape
    
    # Head-level HC
    hc_head = HyperConnectionHead(config)
    final = hc_head(z)
    print(f"  Head: {z.shape} -> {final.shape}")
    assert final.shape == (2, 16, config.hidden_size)
    
    print("  ✅ Hyper-Connections work\n")


def test_gate(config):
    print("=" * 60)
    print("TEST 6: MoE Gate (hash + score routing)")
    print("=" * 60)
    
    # Hash routing (layer 0)
    gate_hash = DeepSeekV4Gate(config, layer_idx=0)
    x = torch.randn(32, config.hidden_size)
    ids = torch.randint(0, config.vocab_size, (32,))
    weights, indices = gate_hash(x, ids)
    print(f"  Hash gate: weights {weights.shape}, indices {indices.shape}")
    assert weights.shape == (32, config.num_experts_per_tok)
    assert indices.shape == (32, config.num_experts_per_tok)
    
    # Score routing (layer 2)
    gate_score = DeepSeekV4Gate(config, layer_idx=2)
    weights, indices = gate_score(x)
    print(f"  Score gate: weights {weights.shape}, indices {indices.shape}")
    assert weights.shape == (32, config.num_experts_per_tok)
    
    print("  ✅ MoE Gate works\n")


def test_expert(config):
    print("=" * 60)
    print("TEST 7: MoE Expert (SwiGLU with clamping)")
    print("=" * 60)
    
    expert = DeepSeekV4Expert(config)
    x = torch.randn(8, config.hidden_size)
    y = expert(x)
    print(f"  Input: {x.shape}, Output: {y.shape}")
    assert y.shape == x.shape
    print("  ✅ Expert works\n")


def test_moe(config):
    print("=" * 60)
    print("TEST 8: Full MoE Layer")
    print("=" * 60)
    
    moe = DeepSeekV4MoE(config, layer_idx=2)
    x = torch.randn(2, 16, config.hidden_size)
    ids = torch.randint(0, config.vocab_size, (2, 16))
    y = moe(x, ids)
    print(f"  Input: {x.shape}, Output: {y.shape}")
    assert y.shape == x.shape
    print("  ✅ MoE works\n")


def test_attention(config):
    print("=" * 60)
    print("TEST 9: Attention (MLA + Compressed Sparse)")
    print("=" * 60)
    
    # Test different layer types
    for layer_idx in [0, 2, 6]:
        ratio = config.compress_ratios[layer_idx] if layer_idx < len(config.compress_ratios) else 0
        attn = DeepSeekV4Attention(config, layer_idx)
        
        seq_len = 128
        x = torch.randn(2, seq_len, config.hidden_size)
        freqs = precompute_freqs_cis(
            config.qk_rope_head_dim, seq_len, config.rope_theta, config.rope_scaling
        )
        
        out, _ = attn(x, freqs_cis=freqs)
        print(f"  Layer {layer_idx} (ratio={ratio}): {x.shape} -> {out.shape}")
        assert out.shape == x.shape
    
    print("  ✅ Attention works\n")


def test_block(config):
    print("=" * 60)
    print("TEST 10: Transformer Block")
    print("=" * 60)
    
    block = DeepSeekV4Block(config, layer_idx=2)
    x = torch.randn(2, 32, config.hc_mult, config.hidden_size)
    ids = torch.randint(0, config.vocab_size, (2, 32))
    freqs = precompute_freqs_cis(
        config.qk_rope_head_dim, 32, config.rope_theta, config.rope_scaling
    )
    
    outputs = block(x, freqs_cis=freqs, input_ids=ids)
    out = outputs[0]
    print(f"  Input: {x.shape}, Output: {out.shape}")
    assert out.shape == x.shape
    print("  ✅ Block works\n")


def test_full_model(config):
    print("=" * 60)
    print("TEST 11: Full DeepSeekV4ForCausalLM")
    print("=" * 60)
    
    model = DeepSeekV4ForCausalLM(config)
    total, trainable = count_params(model)
    print(f"  Total params: {total:,} ({total/1e6:.1f}M)")
    print(f"  Trainable params: {trainable:,} ({trainable/1e6:.1f}M)")
    
    # Forward pass
    input_ids = torch.randint(0, config.vocab_size, (1, 64))
    
    with torch.no_grad():
        outputs = model(input_ids)
    
    logits = outputs.logits
    print(f"  Input: {input_ids.shape}")
    print(f"  Logits: {logits.shape}")
    assert logits.shape == (1, 64, config.vocab_size)
    
    # Test with labels (loss computation)
    labels = torch.randint(0, config.vocab_size, (1, 64))
    labels[0, :5] = -100  # mask some positions
    
    outputs_with_loss = model(input_ids, labels=labels)
    loss = outputs_with_loss.loss
    print(f"  Loss: {loss.item():.4f}")
    assert loss.item() > 0
    
    # Test backward pass
    loss.backward()
    grad_norms = []
    for name, p in model.named_parameters():
        if p.grad is not None:
            grad_norms.append((name, p.grad.norm().item()))
    
    print(f"  Gradients computed for {len(grad_norms)} parameters")
    print(f"  Max grad norm: {max(g for _, g in grad_norms):.4f}")
    
    print("  ✅ Full model works (forward + backward)\n")


def test_generation(config):
    print("=" * 60)
    print("TEST 12: Autoregressive Generation (greedy)")
    print("=" * 60)
    
    model = DeepSeekV4ForCausalLM(config)
    model.eval()
    
    input_ids = torch.randint(0, config.vocab_size, (1, 8))
    
    with torch.no_grad():
        generated = input_ids.clone()
        for step in range(16):
            outputs = model(generated)
            next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
            generated = torch.cat([generated, next_token], dim=1)
    
    print(f"  Input length: 8")
    print(f"  Generated length: {generated.size(1)}")
    print(f"  Generated tokens: {generated[0, 8:].tolist()}")
    print("  ✅ Generation works\n")


if __name__ == "__main__":
    print("\n🔬 DeepSeek-V4 Open Reproduction — Architecture Test Suite\n")
    
    config = test_config()
    test_rmsnorm(config)
    test_rope(config)
    test_kv_compressor(config)
    test_hyper_connections(config)
    test_gate(config)
    test_expert(config)
    test_moe(config)
    test_attention(config)
    test_block(config)
    test_full_model(config)
    test_generation(config)
    
    print("=" * 60)
    print("🎉 ALL TESTS PASSED!")
    print("=" * 60)