import os
import sys
import platform
import psutil
import torch

print("=== System Environment Diagnostic ===")
print("OS:", platform.platform())
print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA Available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("CUDA Version (PyTorch):", torch.version.cuda)
    print("GPU Count:", torch.cuda.device_count())
    print("GPU Name:", torch.cuda.get_device_name(0))
    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
    print(f"Total VRAM: {vram_gb:.2f} GB")
    allocated = torch.cuda.memory_allocated(0) / (1024 ** 3)
    reserved = torch.cuda.memory_reserved(0) / (1024 ** 3)
    print(f"Allocated VRAM: {allocated:.2f} GB | Reserved VRAM: {reserved:.2f} GB")
else:
    print("CUDA Version (PyTorch): N/A (CPU-only PyTorch build or no NVIDIA GPU detected)")

mem = psutil.virtual_memory()
print(f"Total System RAM: {mem.total / (1024**3):.2f} GB")
print(f"Available System RAM: {mem.available / (1024**3):.2f} GB")

# Also check nvidia-smi if available
try:
    import subprocess
    res = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if res.returncode == 0:
        print("\n=== nvidia-smi output ===")
        print(res.stdout)
    else:
        print("\n=== nvidia-smi: Not found or failed ===")
except Exception as e:
    print(f"\n=== nvidia-smi error: {e} ===")