import torch
import torch.nn as nn
import time
import numpy as np
from resnet_attention import resnet18_cbam, resnet34_cbam, resnet50_cbam

# 设置随机种子以确保结果可复现
torch.manual_seed(42)
np.random.seed(42)

# 设备配置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"测试设备: {device}")

# 测试参数
batch_sizes = [1, 4, 8, 16]
image_sizes = [224, 384, 512]
num_runs = 100  # 每个配置运行的次数
warmup_runs = 20  # 预热运行次数，排除启动开销

# 模型配置
model_configs = {
    "resnet18": {
        "base_model": lambda: resnet18_cbam(use_cbam=False),
        "attention_model": lambda: resnet18_cbam(use_cbam=True)
    },
    "resnet34": {
        "base_model": lambda: resnet34_cbam(use_cbam=False),
        "attention_model": lambda: resnet34_cbam(use_cbam=True)
    },
    "resnet50": {
        "base_model": lambda: resnet50_cbam(use_cbam=False),
        "attention_model": lambda: resnet50_cbam(use_cbam=True)
    }
}


# 基准测试函数
def benchmark_model(model, input_size, batch_size, num_runs, warmup_runs):
    """
    测试模型的推理性能

    参数:
    - model: 待测试的模型
    - input_size: 输入图像尺寸
    - batch_size: 批次大小
    - num_runs: 测试运行次数
    - warmup_runs: 预热运行次数

    返回:
    - 平均推理时间(毫秒)
    - 吞吐量(样本/秒)
    """
    # 设置为评估模式
    model.eval()
    model.to(device)

    # 创建随机输入
    input_tensor = torch.randn(batch_size, 3, input_size, input_size, device=device)

    # 预热
    with torch.no_grad():
        for _ in range(warmup_runs):
            _ = model(input_tensor)
            if device.type == 'cuda':
                torch.cuda.synchronize()  # 同步GPU操作

    # 测量推理时间
    start_time = time.time()
    with torch.no_grad():
        for _ in range(num_runs):
            _ = model(input_tensor)
            if device.type == 'cuda':
                torch.cuda.synchronize()  # 同步GPU操作
    end_time = time.time()

    # 计算指标
    total_time = end_time - start_time
    avg_time_per_batch = total_time / num_runs * 1000  # 毫秒
    throughput = batch_size * num_runs / total_time  # 样本/秒

    return avg_time_per_batch, throughput


# 运行测试
results = {}

for model_name, config in model_configs.items():
    results[model_name] = {}

    # 创建模型
    base_model = config["base_model"]()
    attention_model = config["attention_model"]()

    # 计算参数量
    base_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
    attention_params = sum(p.numel() for p in attention_model.parameters() if p.requires_grad)
    param_increase = (attention_params - base_params) / base_params * 100

    print(f"\n测试模型: {model_name}")
    print(f"  基础参数量: {base_params / 1e6:.2f}M")
    print(f"  带注意力参数量: {attention_params / 1e6:.2f}M")
    print(f"  参数量增加: {param_increase:.2f}%")

    for batch_size in batch_sizes:
        for image_size in image_sizes:
            key = f"batch_{batch_size}_size_{image_size}"
            results[model_name][key] = {}

            # 测试基础模型
            base_time, base_throughput = benchmark_model(
                base_model, image_size, batch_size, num_runs, warmup_runs
            )

            # 测试注意力模型
            attention_time, attention_throughput = benchmark_model(
                attention_model, image_size, batch_size, num_runs, warmup_runs
            )

            # 计算增加的百分比
            time_increase = (attention_time - base_time) / base_time * 100
            throughput_decrease = (base_throughput - attention_throughput) / base_throughput * 100

            results[model_name][key]["base_time"] = base_time
            results[model_name][key]["attention_time"] = attention_time
            results[model_name][key]["time_increase"] = time_increase
            results[model_name][key]["base_throughput"] = base_throughput
            results[model_name][key]["attention_throughput"] = attention_throughput
            results[model_name][key]["throughput_decrease"] = throughput_decrease

            print(f"  配置: 批次大小={batch_size}, 图像尺寸={image_size}x{image_size}")
            print(f"    基础模型: 平均时间={base_time:.2f}ms, 吞吐量={base_throughput:.2f}样本/秒")
            print(f"    注意力模型: 平均时间={attention_time:.2f}ms, 吞吐量={attention_throughput:.2f}样本/秒")
            print(f"    时间增加: {time_increase:.2f}%, 吞吐量下降: {throughput_decrease:.2f}%")

# 保存结果
import json

with open('benchmark_results.json', 'w') as f:
    json.dump(results, f, indent=2)

print("\n测试完成，结果已保存到 benchmark_results.json")