import importlib
import sys

def get_cuda_tag(version):
    if '+cu' in version:
        return version.split('+cu')[1]
    return None

# Verificar torch
try:
    import torch
    torch_version = torch.__version__
    torch_cuda = get_cuda_tag(torch_version)
    print(f"✅ PyTorch {torch_version} (CUDA {torch_cuda})")
except ImportError:
    print("❌ PyTorch não está instalado.")
    sys.exit(1)

# Verificar torchvision
try:
    import torchvision
    tv_version = torchvision.__version__
    tv_cuda = get_cuda_tag(torchvision.__version__)
    print(f"✅ Torchvision {tv_version} (CUDA {tv_cuda})")
except ImportError:
    print("❌ Torchvision não está instalado.")
    sys.exit(1)

# Comparar
if torch_cuda != tv_cuda:
    print("⚠️ Incompatibilidade entre PyTorch e Torchvision!")
    print(f"    PyTorch usa CUDA {torch_cuda}, mas Torchvision usa CUDA {tv_cuda}")
    print(f"🔁 Recomendo reinstalar Torchvision com:")
    print(f"\n    pip install torchvision --index-url https://download.pytorch.org/whl/cu{torch_cuda}")
else:
    print("✅ PyTorch e Torchvision estão compatíveis com CUDA.")