import torch

def check_gpu_availability():
    """检查GPU可用性并打印详细信息"""
    print("="*50)
    print("GPU/CUDA 可用性检查")
    print("="*50)
    
    # 1. 检查PyTorch是否支持CUDA
    if not torch.cuda.is_available():
        print("❌ 错误: PyTorch未编译CUDA支持或未检测到NVIDIA驱动")
        return False
    
    # 2. 检查GPU数量
    gpu_count = torch.cuda.device_count()
    print(f"✅ 找到 {gpu_count} 个GPU设备")
    
    # 3. 打印每个GPU的详细信息
    for i in range(gpu_count):
        print(f"\nGPU {i} 详细信息:")
        print(f"  名称: {torch.cuda.get_device_name(i)}")
        print(f"  CUDA计算能力: {torch.cuda.get_device_capability(i)}")
        print(f"  总显存: {torch.cuda.get_device_properties(i).total_memory/1024**3:.2f} GB")
        
        # 检查当前是否被选中
        current_device = torch.cuda.current_device()
        print(f"  {'⭐' if i == current_device else ' '} 当前使用设备: {'是' if i == current_device else '否'}")
    
    # 4. 执行实际GPU计算测试
    try:
        print("\n正在执行GPU计算测试...")
        x = torch.randn(1000, 1000).cuda()
        y = torch.randn(1000, 1000).cuda()
        z = (x @ y).mean()  # 矩阵乘法
        print(f"✅ GPU计算测试成功完成! 结果: {z.item():.4f}")
        return True
    except Exception as e:
        print(f"❌ GPU计算测试失败! 错误信息:")
        print(str(e))
        return False

if __name__ == "__main__":
    if check_gpu_availability():
        print("\n🎉 您的GPU已准备好用于PyTorch训练!")
    else:
        print("\n⚠️ 您的GPU不可用，请检查以上错误信息")