#!/usr/bin/env bash
# launch_training.sh — Run GRPO training on HF Spaces GPU (or locally).
#
# Usage:
#   chmod +x scripts/launch_training.sh
#   ./scripts/launch_training.sh

set -euo pipefail

export USE_UNSLOTH=1
export MODEL_NAME="${MODEL_NAME:-Qwen/Qwen2.5-1.5B-Instruct}"
export OUTPUT_DIR="${OUTPUT_DIR:-checkpoints/}"
export PYTHONUTF8=1

# Activate venv if present
if [ -f ".venv/bin/activate" ]; then
    source .venv/bin/activate
elif [ -f "venv/bin/activate" ]; then
    source venv/bin/activate
fi

echo "=== Tetris GRPO Training ==="
echo "Model:      ${MODEL_NAME}"
echo "Output:     ${OUTPUT_DIR}"
echo "Unsloth:    ${USE_UNSLOTH}"
echo "WANDB:      ${WANDB_PROJECT:-disabled}"
echo "============================"

mkdir -p "${OUTPUT_DIR}"

python training/grpo_train.py 2>&1 | tee "${OUTPUT_DIR}/training.log"

echo ""
echo "Training complete. Logs at ${OUTPUT_DIR}/training.log"