#!/bin/bash
# Setup script for vast.ai GPU instance (8x A100 80GB)
# Run this after SSHing into a fresh vast.ai instance
set -e

echo "========================================="
echo "ArcisVLM Training Environment Setup"
echo "========================================="

# Install system dependencies
echo "=== Installing system packages ==="
apt-get update -qq
apt-get install -y -qq git git-lfs htop tmux 2>&1 | tail -3

# Install Python packages
echo "=== Installing Python packages ==="
pip install --upgrade pip -q
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -q 2>&1 | tail -3
pip install datasets transformers pillow pyyaml tqdm -q 2>&1 | tail -3

# Setup Git LFS
git lfs install

# Clone repo (token should be set as GH_TOKEN env var)
echo "=== Cloning repository ==="
if [ -d "/workspace/arcisvlm" ]; then
    echo "Repo already exists, pulling latest..."
    cd /workspace/arcisvlm && git pull
else
    if [ -n "$GH_TOKEN" ]; then
        git clone https://x-access-token:${GH_TOKEN}@github.com/hardiksa/arcisvlm.git /workspace/arcisvlm
    else
        git clone https://github.com/hardiksa/arcisvlm.git /workspace/arcisvlm
    fi
fi

cd /workspace/arcisvlm

# Pull LFS files (existing checkpoints)
echo "=== Pulling LFS files ==="
git lfs pull

# Download training datasets
echo "=== Downloading datasets ==="
python3 scripts/download_datasets.py --config configs/scale_1.3b.yaml --stage all

# Verify GPU setup
echo ""
echo "========================================="
echo "Environment Verification"
echo "========================================="
python3 -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'GPU count: {torch.cuda.device_count()}')
for i in range(torch.cuda.device_count()):
    name = torch.cuda.get_device_name(i)
    mem = torch.cuda.get_device_properties(i).total_mem / 1e9
    print(f'  GPU {i}: {name} ({mem:.1f} GB)')
print(f'NCCL available: {torch.distributed.is_nccl_available()}')
"

echo ""
echo "========================================="
echo "Setup complete! Run training with:"
echo "  bash scripts/launch_training.sh"
echo "========================================="