import time
import torch

import linear_gelu_torchcode as torchcode
import linear_gelu_cudacode as cudacode


def _to_device(tensors, device):
    return [t.to(device) for t in tensors]


def _copy_params(torch_model, cuda_model):
    with torch.no_grad():
        cuda_model.weight.copy_(torch_model.linear.weight)
        cuda_model.bias.copy_(torch_model.linear.bias)


def _measure_gpu_seconds(model, args, iters=100):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    torch.cuda.synchronize()
    with torch.no_grad():
        start.record()
        for _ in range(iters):
            _ = model(*args)
        end.record()
    torch.cuda.synchronize()
    ms = start.elapsed_time(end) / iters
    return ms / 1000.0


def run_benchmark():
    if not torch.cuda.is_available():
        print("CUDA 不可用")
        return False, 0.0

    torch.manual_seed(0)
    device = torch.device("cuda")

    init_kwargs = torchcode.get_init_inputs()
    torch_model = torchcode.Model(**init_kwargs).to(device).eval()
    cuda_model = cudacode.ModelNew(**init_kwargs).to(device).eval()

    # 关闭编译，避免在当前平台上落到慢路径（CUTLASS 不可用）

    # 参数对齐
    _copy_params(torch_model, cuda_model)

    # 准备输入
    x = torchcode.get_inputs()
    x, = _to_device([x], device)

    print("-------------------- 精度对齐验证 --------------------")
    with torch.no_grad():
        # 预热
        _ = torch_model(x)
        _ = cuda_model(x)

        # 正式测试
        output_torch = torch_model(x)
        output_cuda = cuda_model(x)

    abs_diff = torch.abs(output_torch - output_cuda)
    max_diff = torch.max(abs_diff).item()
    mean_diff = torch.mean(abs_diff).item()
    if max_diff < 1e-4 and mean_diff < 1e-5:
        print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
        precision_flag = True
    else:
        print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
        precision_flag = False

    print("\n-------------------- 性能加速比测试 --------------------")
    num_iterations = 200

    # 预热
    for _ in range(10):
        _ = torch_model(x)
        _ = cuda_model(x)

    # 可选：允许 TF32（若硬件支持），提升矩阵乘性能
    try:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.set_float32_matmul_precision("medium")
    except Exception:
        pass

    # PyTorch计时（CUDA Events，秒）
    torch_time = _measure_gpu_seconds(torch_model, (x,), iters=num_iterations)
    # 优化版计时（CUDA Events，秒）
    cuda_time = _measure_gpu_seconds(cuda_model, (x,), iters=num_iterations)

    print(f"PyTorch内置Linear+GELU平均执行时间: {torch_time:.6f}秒")
    print(f"自定义CUDA Linear+GELU平均执行时间: {cuda_time:.6f}秒")
    speedup = torch_time / cuda_time if cuda_time > 0 else 0.0
    print(f"加速比 (Speedup): {speedup:.2f}x")

    return precision_flag, speedup


if __name__ == "__main__":
    precision_flag, speedup = run_benchmark()