#!/usr/bin/env python3
"""Test script to verify installation and GPU setup"""

import sys
import torch
import cv2
import numpy as np

print("VR180 Matting Installation Test")
print("=" * 50)

# Check Python version
print(f"Python version: {sys.version}")

# Check PyTorch and CUDA
print(f"\nPyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

# Check OpenCV
print(f"\nOpenCV version: {cv2.__version__}")

# Test imports
try:
    from ultralytics import YOLO
    print("\n✅ YOLO import successful")
except ImportError as e:
    print(f"\n❌ YOLO import failed: {e}")

try:
    from sam2.build_sam import build_sam2_video_predictor
    print("✅ SAM2 import successful")
except ImportError as e:
    print(f"❌ SAM2 import failed: {e}")
    print("   Install with: pip install git+https://github.com/facebookresearch/segment-anything-2.git")

try:
    from vr180_matting.config import VR180Config
    from vr180_matting.detector import YOLODetector
    from vr180_matting.sam2_wrapper import SAM2VideoMatting
    from vr180_matting.memory_manager import VRAMManager
    print("✅ VR180 matting modules import successful")
except ImportError as e:
    print(f"❌ VR180 matting import failed: {e}")
    print("   Make sure to run: pip install -e .")

# Memory check
if torch.cuda.is_available():
    print(f"\n📊 Current GPU Memory Usage:")
    print(f"   Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
    print(f"   Reserved:  {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
    print(f"   Free:      {(torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved()) / 1024**3:.2f} GB")

print("\n" + "=" * 50)
print("Installation test complete!")
print("\nNext steps:")
print("1. If any imports failed, install missing dependencies")
print("2. Download SAM2 model weights if needed")
print("3. Run: vr180-matting --generate-config config.yaml")
print("4. Edit config.yaml with your video path")
print("5. Run: vr180-matting config.yaml")