#!/usr/bin/env python3
"""
基础测试
测试项目的基本功能
"""

import os
import sys
import torch
import torch.nn as nn
import numpy as np

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from src.models.unet_light import UNetLight, TimestepEmbedding, ResNetBlock
from src.models.attention import MemoryEfficientAttention
from src.models.diffusion import DiffusionProcess


def test_timestep_embedding():
    """测试时间步嵌入"""
    print("测试时间步嵌入...")
    
    embedding_dim = 256
    time_embed_dim = 512
    
    embedder = TimestepEmbedding(embedding_dim, time_embed_dim)
    
    # 测试前向传播
    timesteps = torch.tensor([100, 200, 300])
    embeddings = embedder(timesteps)
    
    assert embeddings.shape == (3, time_embed_dim)
    print(f"  形状正确: {embeddings.shape}")
    
    return embedder


def test_resnet_block():
    """测试残差块"""
    print("测试残差块...")
    
    in_channels = 64
    out_channels = 128
    time_embed_dim = 256
    
    block = ResNetBlock(in_channels, out_channels, time_embed_dim)
    
    # 测试前向传播
    x = torch.randn(2, in_channels, 32, 32)
    time_emb = torch.randn(2, time_embed_dim)
    
    output = block(x, time_emb)
    
    assert output.shape == (2, out_channels, 32, 32)
    print(f"  形状正确: {output.shape}")
    
    # 测试跳跃连接
    block_same = ResNetBlock(in_channels, in_channels, time_embed_dim)
    output_same = block_same(x, time_emb)
    
    assert output_same.shape == x.shape
    print(f"  跳跃连接正确")
    
    return block


def test_attention():
    """测试注意力机制"""
    print("测试注意力机制...")
    
    dim = 256
    num_heads = 8
    
    attention = MemoryEfficientAttention(dim, num_heads)
    
    # 测试前向传播
    x = torch.randn(2, 16, dim)  # [batch, seq_len, dim]
    output = attention(x)
    
    assert output.shape == x.shape
    print(f"  形状正确: {output.shape}")
    
    return attention


def test_unet_light():
    """测试轻量UNet"""
    print("测试轻量UNet...")
    
    config = {
        'model': {
            'in_channels': 4,
            'out_channels': 4,
            'base_channels': 32,  # 测试用小模型
            'channel_mults': [1, 2, 4],
            'num_res_blocks': 1,
            'attention_resolutions': [8],
            'dropout': 0.0,
            'use_checkpoint': False,
            'num_heads': 4,
            'context_dim': 256,
            'use_linear_projection': True,
            'time_embed_dim': 128
        }
    }
    
    model = UNetLight(config)
    
    # 测试前向传播
    batch_size = 2
    x = torch.randn(batch_size, 4, 64, 64)
    timesteps = torch.randint(0, 1000, (batch_size,))
    context = torch.randn(batch_size, 77, 256)
    
    output = model(x, timesteps, context)
    
    assert output.shape == x.shape
    print(f"  形状正确: {output.shape}")
    
    # 测试梯度检查点
    model.enable_gradient_checkpointing()
    print(f"  梯度检查点已启用")
    
    return model


def test_diffusion_process():
    """测试扩散过程"""
    print("测试扩散过程...")
    
    config = {
        'diffusion': {
            'beta_schedule': 'linear',
            'beta_start': 0.0001,
            'beta_end': 0.02,
            'num_train_timesteps': 100,
            'num_inference_timesteps': 20
        }
    }
    
    diffusion = DiffusionProcess(config)
    
    # 测试前向扩散
    x_start = torch.randn(2, 3, 32, 32)
    t = torch.randint(0, 100, (2,))
    
    x_noisy = diffusion.q_sample(x_start, t)
    
    assert x_noisy.shape == x_start.shape
    print(f"  前向扩散形状正确: {x_noisy.shape}")
    
    # 测试参数提取
    extracted = diffusion.extract(diffusion.sqrt_alphas_cumprod, t, x_start.shape)
    assert extracted.shape == (2, 1, 1, 1)
    print(f"  参数提取形状正确: {extracted.shape}")
    
    return diffusion


def test_memory_efficiency():
    """测试内存效率"""
    print("测试内存效率...")
    
    # 测试模型在不同批次大小下的内存使用
    config = {
        'model': {
            'in_channels': 4,
            'out_channels': 4,
            'base_channels': 32,
            'channel_mults': [1, 2],
            'num_res_blocks': 1,
            'attention_resolutions': [],
            'dropout': 0.0,
            'use_checkpoint': False,
            'num_heads': 4,
            'context_dim': 256,
            'use_linear_projection': True,
            'time_embed_dim': 128
        }
    }
    
    model = UNetLight(config)
    model.eval()
    
    if torch.cuda.is_available():
        device = torch.device('cuda')
        model = model.to(device)
        
        print("  GPU内存测试:")
        
        for batch_size in [1, 2, 4]:
            # 清空缓存
            torch.cuda.empty_cache()
            
            # 记录初始内存
            initial_memory = torch.cuda.memory_allocated()
            
            # 前向传播
            x = torch.randn(batch_size, 4, 64, 64, device=device)
            t = torch.randint(0, 1000, (batch_size,), device=device)
            context = torch.randn(batch_size, 77, 256, device=device)
            
            with torch.no_grad():
                _ = model(x, t, context)
            
            # 记录峰值内存
            peak_memory = torch.cuda.max_memory_allocated()
            memory_used = (peak_memory - initial_memory) / 1024**3  # GB
            
            print(f"    批次大小 {batch_size}: {memory_used:.2f} GB")
    
    else:
        print("  GPU不可用，跳过内存测试")
    
    return model


def run_all_tests():
    """运行所有测试"""
    print("=" * 60)
    print("运行Lumina基础测试")
    print("=" * 60)
    
    try:
        # 测试各个组件
        test_timestep_embedding()
        test_resnet_block()
        test_attention()
        test_unet_light()
        test_diffusion_process()
        test_memory_efficiency()
        
        print("\n" + "=" * 60)
        print("所有测试通过!")
        print("=" * 60)
        
        return True
    
    except Exception as e:
        print(f"\n测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)