import torch
import sys

print("=" * 40)
print("  TitanAudio Pro - GPU System Check")
print("=" * 40)

cuda_ok = torch.cuda.is_available()
print(f"PyTorch Version : {torch.__version__}")
print(f"CUDA Available  : {cuda_ok}")

if cuda_ok:
    print(f"CUDA Version    : {torch.version.cuda}")
    print(f"GPU Device      : {torch.cuda.get_device_name(0)}")
    print(f"VRAM Total      : {round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2)} GB")
    print(f"VRAM Free       : {round((torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)) / 1024**3, 2)} GB")
    print("\n  STATUS: GPU READY - 100% GPU pipeline active!")
else:
    print("\n  WARNING: CUDA not available. CPU mode only.")
    sys.exit(1)

print("=" * 40)

# Also check key libraries
print("\n  Checking installed libraries...")
libs = [
    ("demucs", "demucs"),
    ("librosa", "librosa"),
    ("pedalboard", "pedalboard"),
    ("soundfile", "soundfile"),
    ("audiomentations", "audiomentations"),
    ("onnxruntime", "onnxruntime"),
]

for name, module in libs:
    try:
        m = __import__(module)
        ver = getattr(m, "__version__", "installed")
        print(f"  [OK] {name} v{ver}")
    except ImportError:
        print(f"  [MISSING] {name} - needs install")

print("\nCheck complete!")