#!/bin/bash
# ─────────────────────────────────────────────────────────────
# Setup Python environment on PC with GPU support
#
# Usage:
#   cd spy_optimizer
#   ./setup_pc.sh
# ─────────────────────────────────────────────────────────────

set -euo pipefail

echo "══════════════════════════════════════════════════════"
echo "  🖥️ Setting up GPU training environment"
echo "══════════════════════════════════════════════════════"

# Check Python
python3 --version || { echo "Python3 not found!"; exit 1; }

# Check CUDA
if command -v nvidia-smi &> /dev/null; then
    echo ""
    echo "  🎮 GPU detected:"
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
else
    echo "  ⚠️  nvidia-smi not found. Make sure CUDA drivers are installed."
fi

# Create venv
echo ""
echo "  📦 Creating virtual environment..."
python3 -m venv .venv 2>/dev/null || echo "  (venv already exists)"
source .venv/bin/activate

# Install PyTorch with CUDA
echo ""
echo "  🔧 Installing PyTorch with CUDA support..."
pip install --upgrade pip
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

# Install other dependencies
echo ""
echo "  📚 Installing dependencies..."
pip install pandas numpy scikit-learn pyarrow lightgbm xgboost optuna

# Verify CUDA
echo ""
echo "  🔍 Verifying CUDA..."
python3 -c "
import torch
print(f'  PyTorch: {torch.__version__}')
print(f'  CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'  GPU: {torch.cuda.get_device_name(0)}')
    print(f'  Memory: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB')
    print(f'  CUDA version: {torch.version.cuda}')
"

echo ""
echo "══════════════════════════════════════════════════════"
echo "  ✅ Setup complete!"
echo ""
echo "  Next steps:"
echo "  1. ./sync_to_pc.sh user@server_ip"
echo "  2. source .venv/bin/activate"
echo "  3. python train_gpu.py"
echo "══════════════════════════════════════════════════════"