#!/usr/bin/env bash
# Final GPU setup script with correct library paths

set -eo pipefail

echo "=== Local GPU Training Setup ==="
echo ""

# Check GPU
echo "GPU detected:"
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv
echo ""

VENV_DIR="$HOME/.venv/spies-gpu"

# Create venv
echo "Creating virtual environment..."
python3 -m venv "$VENV_DIR"

echo "Activating venv..."
source "$VENV_DIR/bin/activate"

echo ""
echo "Installing PyTorch with CUDA..."
pip install --quiet torch torchvision --index-url https://download.pytorch.org/whl/cu121

echo ""
echo "Installing training dependencies..."
pip install --quiet accelerate diffusers transformers pillow opencv-python imageio imageio-ffmpeg einops safetensors omegaconf huggingface-hub numpy

echo ""
echo "=== Setup Complete ==="
echo ""
echo "To activate GPU environment:"
echo "  source tools/activate-gpu-env.sh"
echo ""
echo "Or manually:"
echo "  source ~/.venv/spies-gpu/bin/activate"
echo "  export LD_LIBRARY_PATH=/nix/store/ab3753m6i7isgvzphlar0a8xb84gl96i-gcc-15.2.0-lib/lib:/nix/store/755d2m4ix0gm2k0nn4l85cziaccn4dr2-nvidia-x11-580.142-6.12.85/lib:\$LD_LIBRARY_PATH"