#!/usr/bin/env python3
"""
测试DC_hnet模型的脚本
用于验证时间序列分类模型能否正常运行并得到期望的输出形状
"""

import torch
import torch.nn.functional as F
import sys
import os

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from models.DC_hnet import HierEncodersSingleMainConfig, HierEncodersSingleMainClassifier

def test_dc_hnet_model():
    """测试DC_hnet时间序列分类模型"""
    print("=" * 60)
    print("测试DC_hnet时间序列分类模型")
    print("=" * 60)
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 模型参数配置
    B, L, N = 8, 512, 6        # batch_size, seq_length, num_channels
    num_classes = 10            # 分类数
    d_models = [64, 128, 256]   # 各层维度，单调递增
    
    print(f"输入形状: (B={B}, L={L}, N={N})")
    print(f"分类数: {num_classes}")
    print(f"模型维度: {d_models}")
    
    # 编码器配置（每层都是Mamba）
    encoder_cfg_per_stage = [
        dict(arch="m", height=2),  # stage 0: Mamba2, 2层
        dict(arch="m", height=3),  # stage 1: Mamba2, 3层  
    ]
    
    # 主网络配置（使用Transformer）
    main_cfg = dict(
        arch="T", height=4  # Transformer, 4层
    )
    
    # 压缩目标
    target_compression_N_per_stage = [2, 3]  # 每层压缩比例
    
    # 创建配置
    cfg = HierEncodersSingleMainConfig(
        num_channels=N,
        d_models=d_models,
        num_classes=num_classes,
        encoder_cfg_per_stage=encoder_cfg_per_stage,
        main_cfg=main_cfg,
        target_compression_N_per_stage=target_compression_N_per_stage,
        share_channel=True,                  # 通道间共享编码器
        fusion_across_channels="mean",       # 通道融合方式
        dropout=0.1,
    )
    
    print("配置创建完成")
    
    try:
        # 创建模型 - 设置正确的dtype以兼容flash attention
        dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        model = HierEncodersSingleMainClassifier(cfg, device=device, dtype=dtype)
        model = model.to(device)
        print(f"模型创建成功，参数量: {sum(p.numel() for p in model.parameters()):,}")
        
        # 创建随机输入数据 - 使用bfloat16以兼容flash attention
        dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        x = torch.randn(B, L, N, device=device, dtype=dtype)
        mask = torch.ones(B, L, dtype=torch.bool, device=device)
        
        print(f"输入数据形状: {x.shape}, 数据类型: {x.dtype}")
        
        # 前向传播测试
        print("\n开始前向传播测试...")
        model.eval()
        with torch.no_grad():
            logits, seq_debug, aux = model(x, mask=mask, return_seq=False)
        
        print(f"✅ 前向传播成功!")
        print(f"输出logits形状: {logits.shape}")  # 应该是 (B, num_classes)
        print(f"ratio_loss: {aux['ratio_loss']:.4f}")
        
        # 验证输出形状
        expected_shape = (B, num_classes)
        if logits.shape == expected_shape:
            print(f"✅ 输出形状正确: {logits.shape}")
        else:
            print(f"❌ 输出形状错误: 期望 {expected_shape}, 实际 {logits.shape}")
            return False
        
        # 测试训练模式
        print("\n开始训练模式测试...")
        model.train()
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        
        # 创建目标标签
        y = torch.randint(0, num_classes, (B,), device=device)
        
        # 前向传播
        logits, _, aux = model(x, mask=mask, return_seq=False)
        
        # 计算损失
        cls_loss = F.cross_entropy(logits, y)
        ratio_reg = 0.01 * aux["ratio_loss"]  # ratio loss正则化
        total_loss = cls_loss + ratio_reg
        
        print(f"分类损失: {cls_loss:.4f}")
        print(f"ratio损失: {ratio_reg:.4f}")  
        print(f"总损失: {total_loss:.4f}")
        
        # 反向传播
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        
        print("✅ 训练步骤成功!")
        
        # 测试序列返回功能
        print("\n测试序列调试信息返回...")
        with torch.no_grad():
            logits, seq_debug, aux = model(x, mask=mask, return_seq=True)
        
        if seq_debug is not None:
            print(f"✅ 序列调试信息获取成功，包含 {len(seq_debug)} 个通道的信息")
        else:
            print("❌ 序列调试信息获取失败")
        
        print("\n" + "=" * 60)
        print("🎉 DC_hnet模型测试全部通过!")
        print("模型可以正常进行时间序列分类任务")
        print("=" * 60)
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def test_different_configurations():
    """测试不同的模型配置"""
    print("\n测试不同配置...")
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 测试配置1: 不共享通道
    cfg1 = HierEncodersSingleMainConfig(
        num_channels=3,
        d_models=[32, 64],
        num_classes=5,
        encoder_cfg_per_stage=[dict(arch="m", height=2)],
        main_cfg=dict(arch="m", height=3),
        target_compression_N_per_stage=[2],
        share_channel=False,  # 不共享通道
        fusion_across_channels="concat",  # 连接融合
        dropout=0.1,
    )
    
    try:
        dtype1 = torch.bfloat16 if device.type == "cuda" else torch.float32
        model1 = HierEncodersSingleMainClassifier(cfg1, device=device, dtype=dtype1)
        x1 = torch.randn(4, 256, 3, device=device, dtype=dtype1)
        logits1, _, _ = model1(x1)
        print(f"✅ 配置1 (不共享通道, concat融合): 输出形状 {logits1.shape}")
    except Exception as e:
        print(f"❌ 配置1测试失败: {str(e)}")
        
    # 测试配置2: 单层模型
    cfg2 = HierEncodersSingleMainConfig(
        num_channels=2, 
        d_models=[128],  # 只有一层，没有编码器阶段
        num_classes=3,
        encoder_cfg_per_stage=[],  # 空的编码器阶段
        main_cfg=dict(arch="T", height=2),
        target_compression_N_per_stage=[],
        share_channel=True,
        fusion_across_channels="mean",
        dropout=0.1,
    )
    
    try:
        dtype2 = torch.bfloat16 if device.type == "cuda" else torch.float32
        model2 = HierEncodersSingleMainClassifier(cfg2, device=device, dtype=dtype2)
        x2 = torch.randn(2, 128, 2, device=device, dtype=dtype2)
        logits2, _, _ = model2(x2)
        print(f"✅ 配置2 (单层模型): 输出形状 {logits2.shape}")
    except Exception as e:
        print(f"❌ 配置2测试失败: {str(e)}")

if __name__ == "__main__":
    # 主测试
    success = test_dc_hnet_model()
    
    # 额外配置测试
    test_different_configurations()
    
    if success:
        print("\n🎊 所有测试完成!")
        sys.exit(0)
    else:
        print("\n💥 测试失败!")
        sys.exit(1)