#!/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 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')"

# Download SAM2 checkpoints
cd models
echo "📥 Downloading SAM2 models..."

# Try different SAM2 model versions
if [ ! -f "sam2_hiera_large.pt" ]; then
    echo "Trying SAM2 checkpoint version 1..."
    wget -q --show-progress https://dl.fbaipublicfiles.com/segment_anything_2/sam2_hiera_large.pt || true
fi

if [ ! -f "sam2.1_hiera_large.pt" ]; then
    echo "Trying SAM2.1 checkpoint (latest)..."
    wget -q --show-progress https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt || true
fi

# Download SAM2 config files
cd ..
mkdir -p sam2_configs
cd sam2_configs
echo "📥 Downloading SAM2 configuration files..."
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2_hiera_b+.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2_hiera_l.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2_hiera_s.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2_hiera_t.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2.1_hiera_b+.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2.1_hiera_l.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2.1_hiera_s.yaml
wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/sam2_configs/sam2.1_hiera_t.yaml

cd ..

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

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

# Check which SAM2 model is available
echo ""
echo "📊 SAM2 Models available:"
if [ -f "models/sam2_hiera_large.pt" ]; then
    echo "  ✅ sam2_hiera_large.pt"
    echo "     Use in config: sam2_checkpoint: 'models/sam2_hiera_large.pt'"
fi
if [ -f "models/sam2.1_hiera_large.pt" ]; then
    echo "  ✅ sam2.1_hiera_large.pt (recommended)"
    echo "     Use in config: sam2_checkpoint: 'models/sam2.1_hiera_large.pt'"
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"