#!/bin/bash
# ============================================================
#  Установка окружения MLflow на GPU-сервере (RTX 3090)
#  Запускать НА СЕРВЕРЕ: bash setup_server.sh
# ============================================================
set -e

echo "============================================"
echo "  MLflow Practice — установка на сервере"
echo "============================================"

# --- Проверка GPU ---
echo ""
echo "📡 Проверка GPU..."
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
else
    echo "⚠️  nvidia-smi не найден — CUDA может быть не установлена"
fi

# --- Проверка Python ---
echo ""
echo "🐍 Проверка Python..."
python3 --version

# --- Создание виртуального окружения ---
echo ""
echo "📦 Создание виртуального окружения venv..."
if [ ! -d ".venv" ]; then
    python3 -m venv .venv
fi
source .venv/bin/activate
echo "   Активировано: $(which python)"

# --- Обновление pip ---
echo ""
echo "⬆️  Обновление pip..."
pip install --upgrade pip wheel setuptools

# --- Установка MLflow и базовых библиотек ---
echo ""
echo "📥 Установка MLflow + scikit-learn + matplotlib..."
pip install mlflow scikit-learn pandas numpy matplotlib

# --- Установка PyTorch под CUDA ---
echo ""
echo "🔥 Установка PyTorch (CUDA 12.1)..."
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# --- Проверка CUDA ---
echo ""
echo "✅ Проверка PyTorch + 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'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
"

# --- Проверка MLflow ---
echo ""
echo "✅ Проверка MLflow..."
python3 -c "import mlflow; print(f'MLflow: {mlflow.__version__}')"

echo ""
echo "============================================"
echo "  ✅ Установка завершена!"
echo "============================================"
echo ""
echo "Следующие шаги:"
echo "  1. Активируйте окружение:  source .venv/bin/activate"
echo "  2. Запустите MLflow UI:    ./start_ui.sh"
echo "  3. В другом терминале:"
echo "     source .venv/bin/activate"
echo "     python src/train_simple.py --n-estimators 100 --max-depth 8"
echo "     python src/train_gpu.py --epochs 10 --batch-size 256 --lr 0.001"
echo ""