#!/usr/bin/env bash
# startup.sh — Environment setup for HRM_optimised
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

echo "=========================================="
echo "  HRM_optimised — Environment Setup"
echo "=========================================="

# 1. Create virtual environment
if [ ! -d "venv" ]; then
    echo "[1/3] Creating virtual environment (venv/)..."
    python3 -m venv venv
else
    echo "[1/3] Virtual environment already exists."
fi

# 2. Activate venv
echo "[2/3] Activating virtual environment..."
source venv/bin/activate

# 3. Install dependencies
echo "[3/3] Installing dependencies from requirements.txt..."
pip install --upgrade pip
pip install -r requirements.txt

# Try to install triton if available (needed for tiered models)
if [[ $(python3 -c "import torch; print(torch.cuda.is_available())") == "True" ]]; then
    echo "CUDA detected. Attempting to install triton..."
    pip install triton --no-cache-dir || echo "Warning: Triton installation failed. Tiered models may use PyTorch fallback."
fi

echo ""
echo "=========================================="
echo "  ✅ Setup complete!"
echo "  To start, run: source venv/bin/activate"
echo "=========================================="