#!/usr/bin/env bash
set -euo pipefail

# ============================================================
# InsureOS Models — Environment Setup
# Target: 16GB GPU VM (Bytical Audio)
# ============================================================

echo "=== InsureOS Models Setup ==="
echo "GPU check:"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo "WARNING: No GPU detected"

# System packages
echo "=== Installing system dependencies ==="
sudo apt-get update -qq
sudo apt-get install -y -qq python3-pip python3-venv git curl wget

# Create venv
echo "=== Creating Python virtual environment ==="
python3 -m venv .venv
source .venv/bin/activate

# Core ML
echo "=== Installing PyTorch + CUDA ==="
pip install --upgrade pip wheel setuptools
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Training stack
echo "=== Installing training dependencies ==="
pip install \
    transformers>=4.50.0 \
    datasets>=3.0.0 \
    accelerate>=1.2.0 \
    peft>=0.14.0 \
    trl>=0.15.0 \
    bitsandbytes>=0.45.0 \
    flash-attn --no-build-isolation \
    sentencepiece \
    protobuf \
    wandb

# Specialized ML
echo "=== Installing specialized ML packages ==="
pip install \
    scikit-learn>=1.5.0 \
    xgboost>=2.1.0 \
    lightgbm>=4.5.0 \
    interpret>=0.6.0 \
    statsmodels>=0.14.0 \
    scipy>=1.14.0 \
    networkx>=3.4 \
    torch-geometric>=2.6.0 \
    pyg-lib -f https://data.pyg.org/whl/torch-2.5.0+cu121.html

# Serving
echo "=== Installing serving dependencies ==="
pip install \
    fastapi>=0.115.0 \
    uvicorn>=0.32.0 \
    pydantic>=2.10.0

# Evaluation
echo "=== Installing evaluation packages ==="
pip install \
    rouge-score \
    nltk \
    bert-score \
    seqeval

# Data generation
echo "=== Installing data generation packages ==="
pip install \
    faker>=33.0.0 \
    numpy>=1.26.0 \
    pandas>=2.2.0 \
    tqdm>=4.67.0

echo ""
echo "=== Setup complete! ==="
echo "Activate with: source .venv/bin/activate"
echo "Generate data: python -m data.generate_all"
echo "Train all:     bash scripts/train_all.sh"