#!/usr/bin/env python3
"""
CUDA和cuDNN环境测试脚本
验证PyTorch GPU支持是否正常
"""

import sys
import os


def test_cuda():
    """测试CUDA支持"""
    print("=" * 50)
    print("CUDA 环境测试")
    print("=" * 50)

    try:
        import torch
        print(f"PyTorch版本: {torch.__version__}")
        print(f"CUDA可用: {torch.cuda.is_available()}")

        if not torch.cuda.is_available():
            print("错误: CUDA不可用")
            return False

        print(f"CUDA版本: {torch.version.cuda}")

        # 检查cuDNN版本
        try:
            cudnn_ver = torch.backends.cudnn.version()
            print(f"cuDNN版本: {cudnn_ver}")
        except RuntimeError as e:
            print(f"cuDNN版本检测失败: {e}")
            print("警告: cuDNN可能无法正常使用，但GPU张量运算可能仍然正常")
            return False

        print(f"GPU数量: {torch.cuda.device_count()}")

        for i in range(torch.cuda.device_count()):
            print(f"\nGPU {i}:")
            print(f"  名称: {torch.cuda.get_device_name(i)}")
            props = torch.cuda.get_device_properties(i)
            print(f"  计算能力: {props.major}.{props.minor}")
            print(f"  总内存: {props.total_memory / 1024**3:.2f} GB")

        return True

    except ImportError:
        print("错误: PyTorch未安装")
        return False


def test_gpu_tensor_ops():
    """测试GPU张量运算"""
    print("\n" + "=" * 50)
    print("GPU 张量运算测试")
    print("=" * 50)

    try:
        import torch

        if not torch.cuda.is_available():
            print("错误: CUDA不可用，跳过GPU测试")
            return False

        # 创建张量并进行运算
        print("创建1000x1000随机张量...")
        x = torch.randn(1000, 1000).cuda()
        y = torch.randn(1000, 1000).cuda()

        print("执行矩阵乘法...")
        z = torch.matmul(x, y)

        print("执行归约操作...")
        result = z.sum().item()

        print(f"结果: {result:.2f}")
        print("GPU张量运算测试通过!")

        return True

    except Exception as e:
        print(f"错误: {e}")
        return False


def test_cudnn():
    """测试cuDNN加速"""
    print("\n" + "=" * 50)
    print("cuDNN 加速测试")
    print("=" * 50)

    try:
        import torch

        if not torch.cuda.is_available():
            print("错误: CUDA不可用")
            return False

        # 检查cuDNN是否启用
        print(f"cuDNN可用: {torch.backends.cudnn.is_available()}")
        try:
            cudnn_ver = torch.backends.cudnn.version()
            print(f"cuDNN版本: {cudnn_ver}")
        except RuntimeError as e:
            print(f"cuDNN版本检测跳过: {e}")
            print("注意: PyTorch绑定的cuDNN与系统版本可能不匹配，但不影响GPU使用")
        print(f"cuDNN基准测试启用: {torch.backends.cudnn.benchmark}")

        # cuDNN卷积测试
        print("\n执行cuDNN卷积测试...")
        x = torch.randn(1, 3, 224, 224).cuda()
        conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()

        # 预热
        _ = conv(x)

        # 计时测试
        import time
        start = time.time()
        for _ in range(10):
            _ = conv(x)
        torch.cuda.synchronize()
        elapsed = time.time() - start

        print(f"10次卷积耗时: {elapsed*1000:.2f} ms")
        print("cuDNN测试通过!")

        return True

    except Exception as e:
        print(f"错误: {e}")
        return False


def test_ultralytics_gpu():
    """测试ultralytics YOLO GPU支持"""
    print("\n" + "=" * 50)
    print("Ultralytics YOLO GPU 测试")
    print("=" * 50)

    try:
        from ultralytics import YOLO
        import torch
        import os

        # 确保使用GPU加载模型
        if not torch.cuda.is_available():
            print("错误: CUDA不可用")
            return False

        print("加载YOLOv8n模型到GPU...")
        # 关键：加载模型时指定device=0，确保模型在GPU上
        model = YOLO("yolov8n.pt", task="detect")

        # 将模型移到GPU
        model.to("cuda:0")

        print(f"模型设备: {next(model.model.parameters()).device}")

        # 使用本地测试图像进行推理
        test_image = "bus.jpg"
        if not os.path.exists(test_image):
            print("下载测试图像...")
            import urllib.request
            urllib.request.urlretrieve(
                "https://ultralytics.com/images/bus.jpg",
                test_image
            )

        print("执行GPU推理测试...")
        # 关键：device=0指定使用GPU
        results = model(test_image, device=0, verbose=False)

        # 验证结果
        num_detections = len(results[0].boxes)
        print(f"检测到 {num_detections} 个目标")

        # 检查模型是否真的在GPU上运行
        model_device = next(model.model.parameters()).device
        print(f"模型实际设备: {model_device}")
        assert model_device.type == "cuda", "模型未在GPU上运行!"

        print("YOLO GPU测试通过!")
        return True

    except ImportError:
        print("错误: ultralytics未安装")
        return False
    except Exception as e:
        print(f"错误: {e}")
        return False


def main():
    """主函数"""
    print("\n" + "#" * 50)
    print("# 红外故障检测 - CUDA/cuDNN 环境测试")
    print("#" * 50 + "\n")

    results = {
        "CUDA": test_cuda(),
        "GPU张量运算": test_gpu_tensor_ops(),
        "cuDNN": test_cudnn(),
        "YOLO GPU": test_ultralytics_gpu(),
    }

    print("\n" + "=" * 50)
    print("测试结果汇总")
    print("=" * 50)

    all_passed = True
    for name, passed in results.items():
        status = "✅ 通过" if passed else "❌ 失败"
        print(f"{name}: {status}")
        if not passed:
            all_passed = False

    print("\n" + "=" * 50)
    if all_passed:
        print("所有测试通过! 环境配置正确。")
    else:
        print("部分测试失败，请检查环境配置。")
    print("=" * 50)

    return 0 if all_passed else 1


if __name__ == "__main__":
    sys.exit(main())