#!/bin/bash
# RunPod Quick Setup Script

echo "🚀 Setting up VR180 Matting on RunPod..."
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader)"
echo ""

# Update system
echo "📦 Installing system dependencies..."
apt-get update && apt-get install -y ffmpeg git wget nano

# Install Python dependencies
echo "🐍 Installing Python dependencies..."
pip install --upgrade pip
pip install -r requirements.txt

# Install decord for SAM2 video loading
echo "📹 Installing decord for video processing..."
pip install decord

# Install CuPy for GPU acceleration of stereo validation
echo "🚀 Installing CuPy for GPU acceleration..."
# Auto-detect CUDA version and install appropriate CuPy
python -c "
import torch
if torch.cuda.is_available():
    cuda_version = torch.version.cuda
    print(f'CUDA version detected: {cuda_version}')
    if cuda_version.startswith('11.'):
        import subprocess
        subprocess.run(['pip', 'install', 'cupy-cuda11x>=12.0.0'])
        print('Installed CuPy for CUDA 11.x')
    elif cuda_version.startswith('12.'):
        import subprocess
        subprocess.run(['pip', 'install', 'cupy-cuda12x>=12.0.0'])
        print('Installed CuPy for CUDA 12.x')
    else:
        print(f'Unsupported CUDA version: {cuda_version}')
else:
    print('CUDA not available, skipping CuPy installation')
"

# Install SAM2 separately (not on PyPI)
echo "🎯 Installing SAM2..."
pip install git+https://github.com/facebookresearch/segment-anything-2.git

# Install project
echo "📦 Installing VR180 matting package..."
pip install -e .

# Download models
echo "📥 Downloading models..."
mkdir -p models

# Download YOLOv8 models
python -c "from ultralytics import YOLO; YOLO('yolov8n.pt'); YOLO('yolov8m.pt')"

# Clone SAM2 repo for checkpoints
echo "📥 Cloning SAM2 for model checkpoints..."
if [ ! -d "segment-anything-2" ]; then
    git clone https://github.com/facebookresearch/segment-anything-2.git
fi

# Download SAM2 checkpoints using their official script
cd segment-anything-2/checkpoints
if [ ! -f "sam2.1_hiera_large.pt" ]; then
    echo "📥 Downloading SAM2 checkpoints..."
    chmod +x download_ckpts.sh
    bash download_ckpts.sh
fi
cd ../..

# Create working directories
mkdir -p /workspace/data /workspace/output

# Test installation
echo ""
echo "🧪 Testing installation..."
python test_installation.py

# Check which SAM2 models are available
echo ""
echo "📊 SAM2 Models available:"
if [ -f "segment-anything-2/checkpoints/sam2.1_hiera_large.pt" ]; then
    echo "  ✅ sam2.1_hiera_large.pt (recommended)"
    echo "     Config: sam2_model_cfg: 'sam2.1_hiera_l'"
    echo "     Checkpoint: sam2_checkpoint: 'segment-anything-2/checkpoints/sam2.1_hiera_large.pt'"
fi
if [ -f "segment-anything-2/checkpoints/sam2.1_hiera_base_plus.pt" ]; then
    echo "  ✅ sam2.1_hiera_base_plus.pt"
    echo "     Config: sam2_model_cfg: 'sam2.1_hiera_base_plus'"
fi
if [ -f "segment-anything-2/checkpoints/sam2_hiera_large.pt" ]; then
    echo "  ✅ sam2_hiera_large.pt (legacy)"
    echo "     Config: sam2_model_cfg: 'sam2_hiera_l'"
fi

echo ""
echo "✅ Setup complete!"
echo ""
echo "📝 Quick start:"
echo "1. Upload your VR180 video to /workspace/data/"
echo "   wget -O /workspace/data/video.mp4 'your-video-url'"
echo ""
echo "2. Use the RunPod optimized config:"
echo "   cp config_runpod.yaml config.yaml"
echo "   nano config.yaml  # Update video path"
echo ""
echo "3. Run the matting:"
echo "   vr180-matting config.yaml"
echo ""
echo "💡 For A40 GPU, you can use higher quality settings:"
echo "   vr180-matting config.yaml --scale 0.75"