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

import sys
import torch
import cv2
import numpy as np
from pathlib import Path
import os

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

# Track all issues
issues = []

# Check Python version
print(f"Python version: {sys.version}")
if sys.version_info < (3, 8):
    issues.append("Python 3.8+ required")

# 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")
else:
    issues.append("No CUDA GPU detected - will run slowly on CPU")

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

# Test imports
print("\n🔍 Testing imports...")
try:
    from ultralytics import YOLO
    print("✅ YOLO import successful")
    
    # Check if YOLO models exist
    yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
    available_yolo = [m for m in yolo_models if Path(m).exists()]
    if available_yolo:
        print(f"   Found YOLO models: {', '.join(available_yolo)}")
    else:
        issues.append("No YOLO models found - will download on first run")
        
except ImportError as e:
    print(f"❌ YOLO import failed: {e}")
    issues.append("Ultralytics YOLO not installed")

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")
    issues.append("SAM2 not installed")

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 .")
    issues.append("VR180 matting package not installed")

# Check SAM2 models
print("\n🔍 Checking SAM2 models...")
models_dir = Path("models")
sam2_models = {
    "sam2_hiera_large.pt": "Original SAM2 Large",
    "sam2.1_hiera_large.pt": "SAM2.1 Large (recommended)",
    "sam2_hiera_base.pt": "SAM2 Base",
    "sam2.1_hiera_base.pt": "SAM2.1 Base"
}

found_models = []
for model_file, model_name in sam2_models.items():
    model_path = models_dir / model_file
    if model_path.exists():
        size_mb = model_path.stat().st_size / (1024 * 1024)
        print(f"✅ {model_name}: {model_file} ({size_mb:.1f} MB)")
        found_models.append(model_file)
    
if not found_models:
    print("❌ No SAM2 models found!")
    issues.append("No SAM2 models found - run setup script or download manually")
else:
    print(f"\n💡 Recommended config for best model found:")
    if "sam2.1_hiera_large.pt" in found_models:
        print("   sam2_model_cfg: 'sam2.1_hiera_l'")
        print("   sam2_checkpoint: 'models/sam2.1_hiera_large.pt'")
    elif "sam2_hiera_large.pt" in found_models:
        print("   sam2_model_cfg: 'sam2_hiera_l'")
        print("   sam2_checkpoint: 'models/sam2_hiera_large.pt'")

# Check SAM2 configs
print("\n🔍 Checking SAM2 config files...")
configs_dir = Path("sam2_configs")
if configs_dir.exists():
    config_files = list(configs_dir.glob("*.yaml"))
    if config_files:
        print(f"✅ Found {len(config_files)} SAM2 config files")
    else:
        print("❌ No SAM2 config files found")
        issues.append("SAM2 config files missing - may cause model loading errors")
else:
    print("❌ sam2_configs directory not found")
    issues.append("SAM2 configs directory missing")

# Test model loading if possible
if not any("SAM2 not installed" in issue for issue in issues):
    print("\n🧪 Testing SAM2 model loading...")
    try:
        # Try to load the default model config
        from vr180_matting.config import VR180Config
        config = VR180Config(
            input=type('obj', (object,), {'video_path': 'test.mp4'})(),
            processing=type('obj', (object,), {'scale_factor': 0.5, 'chunk_size': 900, 'overlap_frames': 60})(),
            detection=type('obj', (object,), {'confidence_threshold': 0.7, 'model': 'yolov8n'})(),
            matting=type('obj', (object,), {
                'use_disparity_mapping': True,
                'memory_offload': True,
                'fp16': True,
                'sam2_model_cfg': 'sam2.1_hiera_l' if Path('models/sam2.1_hiera_large.pt').exists() else 'sam2_hiera_l',
                'sam2_checkpoint': 'models/sam2.1_hiera_large.pt' if Path('models/sam2.1_hiera_large.pt').exists() else 'models/sam2_hiera_large.pt'
            })(),
            output=type('obj', (object,), {'path': 'output/', 'format': 'alpha', 'background_color': [0, 255, 0], 'maintain_sbs': True})(),
            hardware=type('obj', (object,), {'device': 'cuda' if torch.cuda.is_available() else 'cpu', 'max_vram_gb': 10})()
        )
        
        # Try loading just the model config check
        model_path = Path(config.matting.sam2_checkpoint)
        if model_path.exists():
            print(f"✅ Found checkpoint: {model_path}")
            # Quick check of checkpoint structure
            checkpoint = torch.load(model_path, map_location='cpu')
            if 'model' in checkpoint:
                print(f"   Model has {len(checkpoint['model'])} parameters")
        else:
            print(f"❌ Checkpoint not found: {model_path}")
            issues.append(f"SAM2 checkpoint missing: {model_path}")
            
    except Exception as e:
        print(f"⚠️  Could not test model loading: {e}")

# Check ffmpeg
print("\n🔍 Checking ffmpeg...")
import subprocess
try:
    result = subprocess.run(['ffmpeg', '-version'], capture_output=True, text=True)
    if result.returncode == 0:
        version_line = result.stdout.split('\n')[0]
        print(f"✅ {version_line}")
    else:
        print("❌ ffmpeg error")
        issues.append("ffmpeg not working properly")
except FileNotFoundError:
    print("❌ ffmpeg not found")
    issues.append("ffmpeg not installed - required for video processing")

# 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")

# Check example config
print("\n🔍 Checking configuration files...")
example_configs = ["config_example.yaml", "config_runpod.yaml", "config.yaml"]
found_configs = [c for c in example_configs if Path(c).exists()]
if found_configs:
    print(f"✅ Found configs: {', '.join(found_configs)}")
else:
    print("❌ No config files found")
    issues.append("No configuration files found - generate with: vr180-matting --generate-config config.yaml")

# Summary
print("\n" + "=" * 50)
if issues:
    print("❌ Issues found:")
    for i, issue in enumerate(issues, 1):
        print(f"   {i}. {issue}")
    print("\n📋 To fix issues:")
    print("   1. Run: ./runpod_setup.sh")
    print("   2. Make sure SAM2 is installed: pip install git+https://github.com/facebookresearch/segment-anything-2.git")
    print("   3. Install package: pip install -e .")
else:
    print("✅ All checks passed! Ready to process videos.")
    print("\n📋 Quick start:")
    print("   1. Copy example config: cp config_runpod.yaml config.yaml")
    print("   2. Edit config.yaml with your video path")
    print("   3. Run: vr180-matting config.yaml")