#!/usr/bin/env python3
"""
测试 torch.compile 兼容性的脚本
用于验证 components.py 中的修改是否与 torch.compile 兼容
"""

import os
import sys
import time

import torch
import torch.nn as nn

# 添加 src 目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.model.components import (
    AttentionPooling,
    ContextEncoder,
    CrossAttentionFusion,
    Expert,
    MoELayer,
    ResidualBlock,
    SlotMemory,
)


def test_attention_pooling():
    """测试 AttentionPooling 模块"""
    print("=" * 60)
    print("测试 AttentionPooling 模块")

    batch_size = 4
    seq_len = 10
    hidden_size = 512

    # 创建模块
    attn_pool = AttentionPooling(hidden_size)

    # 测试数据
    x = torch.randn(batch_size, seq_len, hidden_size)
    mask = torch.ones(batch_size, seq_len, dtype=torch.long)
    mask[0, 5:] = 0  # 第一个样本的部分位置mask
    token_type_ids = torch.randint(0, 3, (batch_size, seq_len))

    # 测试未编译版本
    output = attn_pool(x, mask=mask, token_type_ids=token_type_ids)
    print(f"✓ AttentionPooling 输出形状: {output.shape}")
    assert output.shape == (batch_size, hidden_size)

    # 测试编译版本
    try:
        compiled_attn_pool = torch.compile(attn_pool, mode="reduce-overhead")
        compiled_output = compiled_attn_pool(
            x, mask=mask, token_type_ids=token_type_ids
        )

        # 检查输出是否一致
        diff = torch.abs(output - compiled_output).max().item()
        print(f"✓ 编译前后输出差异: {diff:.6f}")
        assert diff < 1e-4, f"输出差异过大: {diff}"
        print("✓ AttentionPooling 通过 torch.compile 测试")
    except Exception as e:
        print(f"⚠ AttentionPooling torch.compile 测试失败: {e}")
        raise


def test_moe_layer():
    """测试 MoELayer 模块"""
    print("\n" + "=" * 60)
    print("测试 MoELayer 模块")

    batch_size = 2
    seq_len = 8
    dim = 512
    num_experts = 20
    top_k = 2

    # 创建模块
    moe = MoELayer(dim=dim, num_experts=num_experts, top_k=top_k)

    # 测试数据
    x = torch.randn(batch_size, seq_len, dim)

    # 测试未编译版本
    output = moe(x)
    print(f"✓ MoELayer 输出形状: {output.shape}")
    assert output.shape == (batch_size, seq_len, dim)

    # 检查门控权重
    gates = moe.gate(x.view(-1, dim))
    topk_vals, topk_indices = torch.topk(gates, top_k, dim=-1)
    print(f"✓ 门控权重形状: {gates.shape}, top-k 索引形状: {topk_indices.shape}")

    # 测试编译版本
    try:
        compiled_moe = torch.compile(moe, mode="reduce-overhead")

        # 预热
        for _ in range(3):
            _ = compiled_moe(x)

        compiled_output = compiled_moe(x)

        # 检查输出是否一致
        diff = torch.abs(output - compiled_output).max().item()
        print(f"✓ 编译前后输出差异: {diff:.6f}")
        assert diff < 1e-4, f"输出差异过大: {diff}"
        print("✓ MoELayer 通过 torch.compile 测试")

        # 性能测试
        n_iter = 50
        print(f"\n性能测试 ({n_iter} 次迭代):")

        # 未编译版本
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        start = time.time()
        for _ in range(n_iter):
            _ = moe(x)
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        base_time = time.time() - start

        # 编译版本
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        start = time.time()
        for _ in range(n_iter):
            _ = compiled_moe(x)
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        compiled_time = time.time() - start

        print(f"  未编译: {base_time:.4f} 秒")
        print(f"  已编译: {compiled_time:.4f} 秒")
        if compiled_time > 0:
            speedup = base_time / compiled_time
            print(f"  加速比: {speedup:.2f}x")

    except Exception as e:
        print(f"⚠ MoELayer torch.compile 测试失败: {e}")
        import traceback

        traceback.print_exc()
        raise


def test_cross_attention_fusion():
    """测试 CrossAttentionFusion 模块"""
    print("\n" + "=" * 60)
    print("测试 CrossAttentionFusion 模块")

    batch_size = 2
    num_slots = 8
    ctx_len = 16
    dim = 512

    # 创建模块
    cross_attn = CrossAttentionFusion(dim=dim, n_heads=4)

    # 测试数据
    slots_S = torch.randn(batch_size, num_slots, dim)
    context_H = torch.randn(batch_size, ctx_len, dim)
    context_mask = torch.ones(batch_size, ctx_len, dtype=torch.long)
    context_mask[0, 10:] = 0  # 第一个样本的部分位置mask

    # 测试未编译版本
    output = cross_attn(slots_S, context_H, context_mask=context_mask)
    print(f"✓ CrossAttentionFusion 输出形状: {output.shape}")
    assert output.shape == (batch_size, num_slots, dim)

    # 测试编译版本
    try:
        compiled_cross_attn = torch.compile(cross_attn, mode="reduce-overhead")
        compiled_output = compiled_cross_attn(
            slots_S, context_H, context_mask=context_mask
        )

        # 检查输出是否一致
        diff = torch.abs(output - compiled_output).max().item()
        print(f"✓ 编译前后输出差异: {diff:.6f}")
        assert diff < 1e-4, f"输出差异过大: {diff}"
        print("✓ CrossAttentionFusion 通过 torch.compile 测试")
    except Exception as e:
        print(f"⚠ CrossAttentionFusion torch.compile 测试失败: {e}")
        raise


def test_context_encoder():
    """测试 ContextEncoder 模块"""
    print("\n" + "=" * 60)
    print("测试 ContextEncoder 模块")

    batch_size = 2
    seq_len = 32
    vocab_size = 1000
    pinyin_vocab_size = 30
    dim = 512

    # 创建模块
    context_encoder = ContextEncoder(
        vocab_size=vocab_size,
        pinyin_vocab_size=pinyin_vocab_size,
        dim=dim,
        n_layers=2,  # 测试时减少层数
        n_heads=4,
        max_len=seq_len,
    )

    # 测试数据
    text_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    pinyin_ids = torch.randint(
        0, pinyin_vocab_size, (batch_size, 24)
    )  # pinyin_ids固定长度24
    mask = torch.ones(batch_size, seq_len, dtype=torch.long)
    mask[0, 20:] = 0  # 第一个样本的部分位置mask

    # 测试未编译版本
    output = context_encoder(text_ids, pinyin_ids, mask=mask)
    print(f"✓ ContextEncoder 输出形状: {output.shape}")
    assert output.shape == (batch_size, seq_len, dim)

    # 测试编译版本（ContextEncoder包含外部模型，可能不完全兼容）
    try:
        compiled_context_encoder = torch.compile(
            context_encoder, mode="reduce-overhead"
        )
        compiled_output = compiled_context_encoder(text_ids, pinyin_ids, mask=mask)

        # 检查输出是否一致
        diff = torch.abs(output - compiled_output).max().item()
        print(f"✓ 编译前后输出差异: {diff:.6f}")
        if diff < 1e-3:  # ContextEncoder 可能精度要求稍低
            print("✓ ContextEncoder 通过 torch.compile 测试")
        else:
            print(f"⚠ ContextEncoder 输出差异较大: {diff:.6f} (可能因外部模型)")
    except Exception as e:
        print(
            f"⚠ ContextEncoder torch.compile 测试失败: {e} (可能是正常现象，因为包含外部模型)"
        )


def test_full_model_compile():
    """测试完整模型编译"""
    print("\n" + "=" * 60)
    print("测试完整模型编译")

    # 导入完整模型
    from src.model.model import InputMethodEngine

    batch_size = 2
    vocab_size = 10019
    pinyin_vocab_size = 30
    dim = 512
    num_slots = 8
    seq_len = 32

    # 创建模型
    model = InputMethodEngine(
        vocab_size=vocab_size,
        pinyin_vocab_size=pinyin_vocab_size,
        dim=dim,
        num_slots=num_slots,
        n_layers=2,  # 测试时减少层数
        n_heads=4,
        num_experts=20,
        max_seq_len=seq_len,
        compile=False,  # 手动控制编译
    )

    # 测试数据
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    token_type_ids = torch.randint(0, 2, (batch_size, seq_len))
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)
    attention_mask[0, 20:] = 0
    pinyin_ids = torch.randint(0, pinyin_vocab_size, (batch_size, 24))
    history_slot_ids = torch.randint(0, vocab_size, (batch_size, num_slots))

    # 测试未编译版本
    with torch.no_grad():
        output = model(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            pinyin_ids=pinyin_ids,
            history_slot_ids=history_slot_ids,
        )

    print(f"✓ 完整模型输出形状: {output.shape}")
    assert output.shape == (batch_size, vocab_size)

    # 测试编译版本
    try:
        # 创建一个新模型用于编译测试
        model_for_compile = InputMethodEngine(
            vocab_size=vocab_size,
            pinyin_vocab_size=pinyin_vocab_size,
            dim=dim,
            num_slots=num_slots,
            n_layers=2,
            n_heads=4,
            num_experts=20,
            max_seq_len=seq_len,
            compile=False,
        )

        # 手动编译
        compiled_model = torch.compile(model_for_compile, mode="reduce-overhead")

        # 预热
        for _ in range(3):
            _ = compiled_model(
                input_ids=input_ids,
                token_type_ids=token_type_ids,
                attention_mask=attention_mask,
                pinyin_ids=pinyin_ids,
                history_slot_ids=history_slot_ids,
            )

        with torch.no_grad():
            compiled_output = compiled_model(
                input_ids=input_ids,
                token_type_ids=token_type_ids,
                attention_mask=attention_mask,
                pinyin_ids=pinyin_ids,
                history_slot_ids=history_slot_ids,
            )

        # 检查输出是否一致
        diff = torch.abs(output - compiled_output).max().item()
        print(f"✓ 编译前后输出差异: {diff:.6f}")

        # 完整模型精度要求可能较低
        if diff < 1e-3:
            print("✓ 完整模型通过 torch.compile 测试")
        else:
            print(f"⚠ 完整模型输出差异较大: {diff:.6f}")

        # 性能测试
        n_iter = 30
        print(f"\n完整模型性能测试 ({n_iter} 次迭代):")

        # 未编译版本
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        start = time.time()
        with torch.no_grad():
            for _ in range(n_iter):
                _ = model(
                    input_ids=input_ids,
                    token_type_ids=token_type_ids,
                    attention_mask=attention_mask,
                    pinyin_ids=pinyin_ids,
                    history_slot_ids=history_slot_ids,
                )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        base_time = time.time() - start

        # 编译版本
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        start = time.time()
        with torch.no_grad():
            for _ in range(n_iter):
                _ = compiled_model(
                    input_ids=input_ids,
                    token_type_ids=token_type_ids,
                    attention_mask=attention_mask,
                    pinyin_ids=pinyin_ids,
                    history_slot_ids=history_slot_ids,
                )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        compiled_time = time.time() - start

        print(f"  未编译: {base_time:.4f} 秒")
        print(f"  已编译: {compiled_time:.4f} 秒")
        if compiled_time > 0:
            speedup = base_time / compiled_time
            print(f"  加速比: {speedup:.2f}x")

    except Exception as e:
        print(f"⚠ 完整模型 torch.compile 测试失败: {e}")
        import traceback

        traceback.print_exc()


def check_compile_issues():
    """检查可能导致编译问题的代码模式"""
    print("\n" + "=" * 60)
    print("检查编译问题")

    issues = []

    # 检查 components.py 中的潜在问题
    with open("src/model/components.py", "r") as f:
        content = f.read()

    # 检查 float('-inf')
    if "float('-inf')" in content:
        issues.append("❌ 发现 float('-inf')，应替换为 -1e9")

    # 检查 .item() 调用
    if ".item()" in content and "def forward" in content:
        # 统计 forward 方法中的 .item() 调用
        lines = content.split("\n")
        in_forward = False
        item_calls = []
        for i, line in enumerate(lines):
            if "def forward" in line:
                in_forward = True
            elif in_forward and "def " in line and "def forward" not in line:
                in_forward = False
            if in_forward and ".item()" in line:
                item_calls.append((i + 1, line.strip()))

        if item_calls:
            issues.append(
                f"❌ 在 forward 方法中发现 {len(item_calls)} 个 .item() 调用:"
            )
            for line_num, line in item_calls[:3]:  # 显示前3个
                issues.append(f"   第 {line_num} 行: {line}")

    # 检查动态控制流
    dynamic_patterns = [
        "if not mask.any():",
        "if mask is not None:",
        "if token_mask.any():",
        "continue",
    ]
    for pattern in dynamic_patterns:
        if pattern in content:
            # 检查是否在 forward 方法中
            issues.append(f"⚠ 发现动态控制流: {pattern}")

    if not issues:
        print("✓ 未发现明显的编译问题")
    else:
        print("发现以下可能影响编译的问题:")
        for issue in issues:
            print(issue)


def main():
    """主测试函数"""
    print("=" * 70)
    print("torch.compile 兼容性测试")
    print("=" * 70)

    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    print(f"PyTorch 版本: {torch.__version__}")

    # 检查 torch.compile 是否可用
    if hasattr(torch, "compile"):
        print(f"✓ torch.compile 可用")
    else:
        print("❌ torch.compile 不可用，需要 PyTorch 2.0+")
        return

    # 运行测试
    try:
        # 首先检查代码问题
        check_compile_issues()

        # 模块测试
        test_attention_pooling()
        test_moe_layer()
        test_cross_attention_fusion()
        test_context_encoder()

        # 完整模型测试
        test_full_model_compile()

        print("\n" + "=" * 70)
        print("✅ 所有测试完成!")

    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()