#!/bin/bash
set -e  # Exit on any error

echo "============================================================================"
echo "VINE Model - Complete Setup Script"
echo "This script sets up everything needed to use the VINE model from HuggingFace"
echo "Model: https://huggingface.co/video-fm/vine"
echo "============================================================================"

# Get the directory where this script is located
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$SCRIPT_DIR"

# ============================================================================
# Step 2: Install PyTorch
# ============================================================================
echo ""
echo "=========================================="
echo "Step 2: Installing PyTorch 2.7.1 with CUDA 12.6 support..."
echo "=========================================="
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu126

# Upgrade pip
pip install --upgrade pip

# ============================================================================
# Step 3: Install Core Dependencies
# ============================================================================
echo ""
echo "=========================================="
echo "Step 3: Installing core dependencies..."
echo "=========================================="

# Install transformers and HuggingFace tools
pip install transformers>=4.40.0
pip install huggingface-hub
pip install safetensors
pip install accelerate

# Install video processing dependencies
pip install opencv-python
pip install pillow
pip install matplotlib
pip install seaborn
pip install pandas
pip install numpy
pip install tqdm
pip install scikit-learn

# Install Gradio for demos (optional)
pip install gradio

echo "✓ Core dependencies installed"

# ============================================================================
# Step 4: Clone Required Repositories
# ============================================================================
echo ""
echo "=========================================="
echo "Step 4: Cloning required repositories..."
echo "=========================================="

mkdir -p src
cd src

# Clone SAM2
if [ ! -d "video-sam2" ]; then
    echo "Cloning video-sam2..."
    git clone https://github.com/video-fm/video-sam2.git
else
    echo "video-sam2 already exists, skipping clone"
fi

# Clone GroundingDINO
if [ ! -d "GroundingDINO" ]; then
    echo "Cloning GroundingDINO..."
    git clone https://github.com/video-fm/GroundingDINO.git
else
    echo "GroundingDINO already exists, skipping clone"
fi

# Clone LASER
if [ ! -d "LASER" ]; then
    echo "Cloning LASER..."
    git clone https://github.com/kevinxuez/LASER.git
else
    echo "LASER already exists, skipping clone"
fi

# Clone vine_hf
if [ ! -d "vine_hf" ]; then
    echo "Cloning vine_hf..."
    git clone https://github.com/kevinxuez/vine_hf.git
else
    echo "vine_hf already exists, skipping clone"
fi

echo "✓ Repositories cloned"

# ============================================================================
# Step 5: Install Packages in Editable Mode
# ============================================================================
echo ""
echo "=========================================="
echo "Step 5: Installing packages in editable mode..."
echo "=========================================="

echo "Installing video-sam2..."
pip install --no-cache-dir -e ./video-sam2

echo "Installing GroundingDINO..."
pip install --no-cache-dir --use-pep517 -e ./GroundingDINO

echo "Installing LASER..."
pip install --no-cache-dir -e ./LASER

echo "Installing vine_hf..."
pip install --no-cache-dir -e ./vine_hf

echo "✓ Packages installed"

# ============================================================================
# Step 6: Build GroundingDINO Extensions
# ============================================================================
echo ""
echo "=========================================="
echo "Step 6: Building GroundingDINO native extensions..."
echo "=========================================="
cd GroundingDINO
python setup.py build_ext --force --inplace
cd ..

# Return to main directory
cd "$SCRIPT_DIR"

echo "✓ Extensions built"

# ============================================================================
# Step 7: Download Model Checkpoints
# ============================================================================
echo ""
echo "=========================================="
echo "Step 7: Downloading model checkpoints..."
echo "=========================================="

# Create checkpoints directory
mkdir -p checkpoints
cd checkpoints

# Download SAM2 checkpoint (~149 MB)
if [ ! -f "sam2_hiera_tiny.pt" ]; then
    echo "Downloading SAM2 checkpoint (sam2_hiera_tiny.pt ~149 MB)..."
    wget -q --show-progress https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt
    echo "✓ SAM2 checkpoint downloaded"
else
    echo "✓ SAM2 checkpoint already exists"
fi

# Download SAM2 config
if [ ! -f "sam2_hiera_t.yaml" ]; then
    echo "Downloading SAM2 config (sam2_hiera_t.yaml)..."
    wget -q --show-progress https://raw.githubusercontent.com/facebookresearch/sam2/main/sam2/configs/sam2.1/sam2.1_hiera_t.yaml -O sam2_hiera_t.yaml
    echo "✓ SAM2 config downloaded"
else
    echo "✓ SAM2 config already exists"
fi

# Download GroundingDINO checkpoint (~662 MB)
if [ ! -f "groundingdino_swint_ogc.pth" ]; then
    echo "Downloading GroundingDINO checkpoint (groundingdino_swint_ogc.pth ~662 MB)..."
    wget -q --show-progress https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
    echo "✓ GroundingDINO checkpoint downloaded"
else
    echo "✓ GroundingDINO checkpoint already exists"
fi

# Download GroundingDINO config
if [ ! -f "GroundingDINO_SwinT_OGC.py" ]; then
    echo "Downloading GroundingDINO config (GroundingDINO_SwinT_OGC.py)..."
    wget -q --show-progress https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/main/groundingdino/config/GroundingDINO_SwinT_OGC.py
    echo "✓ GroundingDINO config downloaded"
else
    echo "✓ GroundingDINO config already exists"
fi

# Return to main directory
cd "$SCRIPT_DIR"

echo ""
echo "✓ All checkpoints downloaded to: $SCRIPT_DIR/checkpoints/"

# ============================================================================
# Step 8: Create Test Script
# ============================================================================
echo ""
echo "=========================================="
echo "Step 8: Creating test script..."
echo "=========================================="

cat > test_vine.py << 'TESTEOF'
"""
Test script for VINE model loaded from HuggingFace Hub
"""
import os
import sys
from pathlib import Path

os.environ['OPENAI_API_KEY'] = "dummy-key"

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

print("=" * 80)
print("Testing VINE Model from video-fm/vine")
print("=" * 80)

# Load VINE from HuggingFace
print("\n1. Loading VINE model from HuggingFace Hub...")
from transformers import AutoModel
model = AutoModel.from_pretrained('video-fm/vine', trust_remote_code=True)
print("✓ Model loaded successfully")

# Verify checkpoint files
print("\n2. Verifying checkpoint files...")
checkpoint_dir = Path(__file__).parent / "checkpoints"
checkpoints = {
    "SAM2 config": checkpoint_dir / "sam2_hiera_t.yaml",
    "SAM2 checkpoint": checkpoint_dir / "sam2_hiera_tiny.pt",
    "GroundingDINO config": checkpoint_dir / "GroundingDINO_SwinT_OGC.py",
    "GroundingDINO checkpoint": checkpoint_dir / "groundingdino_swint_ogc.pth",
}

all_exist = True
for name, path in checkpoints.items():
    if path.exists():
        size_mb = path.stat().st_size / (1024 * 1024)
        print(f"✓ {name}: {path.name} ({size_mb:.1f} MB)")
    else:
        print(f"✗ {name}: NOT FOUND at {path}")
        all_exist = False

# Create pipeline
print("\n3. Creating VINE pipeline...")
from vine_hf import VinePipeline

pipeline = VinePipeline(
    model=model,
    tokenizer=None,
    sam_config_path=str(checkpoints["SAM2 config"]),
    sam_checkpoint_path=str(checkpoints["SAM2 checkpoint"]),
    gd_config_path=str(checkpoints["GroundingDINO config"]),
    gd_checkpoint_path=str(checkpoints["GroundingDINO checkpoint"]),
    device="cuda",
    trust_remote_code=True
)
print("✓ Pipeline created successfully")

print("\n" + "=" * 80)
print("✅ VINE Setup Complete and Working!")
print("=" * 80)
print("\nYou can now use the model for video understanding:")
print("""
from transformers import AutoModel
from vine_hf import VinePipeline

model = AutoModel.from_pretrained('video-fm/vine', trust_remote_code=True)
pipeline = VinePipeline(model=model, ...)
results = pipeline('video.mp4', categorical_keywords=['person', 'dog'], ...)
""")
TESTEOF

echo "✓ Test script created: test_vine.py"

# ============================================================================
# Step 9: Test the Installation
# ============================================================================
echo ""
echo "=========================================="
echo "Step 9: Testing installation..."
echo "=========================================="

echo "Checking PyTorch and CUDA..."
python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}'); print(f'CUDA version: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}')"

echo ""
echo "Running VINE model test..."
python test_vine.py

# ============================================================================
# Final Summary
# ============================================================================
echo ""
echo "============================================================================"
echo "✅ VINE Setup Complete!"
echo "============================================================================"
echo ""
echo "What was installed:"
echo "  ✓ Conda environment: vine_demo"
echo "  ✓ PyTorch 2.7.1 with CUDA 12.6"
echo "  ✓ Required packages: laser, sam2, groundingdino, vine_hf"
echo "  ✓ Model checkpoints downloaded to: checkpoints/"
echo ""
echo "Checkpoint files:"
echo "  ✓ checkpoints/sam2_hiera_tiny.pt (~149 MB)"
echo "  ✓ checkpoints/sam2_hiera_t.yaml"
echo "  ✓ checkpoints/groundingdino_swint_ogc.pth (~662 MB)"
echo "  ✓ checkpoints/GroundingDINO_SwinT_OGC.py"
echo ""
echo "To use VINE:"
echo "  1. Activate environment: conda activate vine_demo"
echo "  2. Run your script or test_vine.py"
echo ""
echo "Example usage:"
echo "  python test_vine.py"
echo ""
echo "Model URL: https://huggingface.co/video-fm/vine"
echo "Documentation: See README.md on HuggingFace Hub"
echo ""
echo "🎉 Happy video understanding with VINE!"
echo "============================================================================"