#!/bin/bash
# =============================================================================
# train_baseline.sh — Karman Re100 Baseline (scratch, multi-seed)
#
# Trains kar_re100 from scratch with N seeds on a single GPU.
# Output: output/kar_re100_seed{N}/
#
# Usage:
#   bash train_baseline.sh [--seeds 41,42,43,44,45] [--gpu 0] [--episodes 500]
# =============================================================================
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TRAIN_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
REPO_DIR="$(cd "$TRAIN_DIR/../../.." && pwd)"
CONDA_ENV="pycuda_3_10"
GPU=0; EPISODES=500; DELAY=120
SEEDS=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --seeds) SEEDS="$2"; shift 2 ;;
        --gpu) GPU="$2"; shift 2 ;;
        --episodes) EPISODES="$2"; shift 2 ;;
        *) echo "Unknown: $1"; exit 1 ;;
    esac
done

registry() { PYTHONPATH="$REPO_DIR/src" python3 -m drl_pinball.train.case_registry kar_re100 "$1"; }
CAL="$(registry calibration)"
CFG="$(registry config)"
TRAINER="$(registry trainer)"
[[ -n "$SEEDS" ]] || SEEDS="$(registry seeds)"
IFS=',' read -ra SEED_ARR <<< "$SEEDS"

echo "=== Baseline: kar_re100, seeds=${SEED_ARR[*]}, ep=$EPISODES, GPU=$GPU ==="

for seed in "${SEED_ARR[@]}"; do
    echo "  [$(date '+%H:%M:%S')] Training seed=$seed..."
    conda run --no-capture-output -n "$CONDA_ENV" python -u \
        "$TRAINER" \
        --case-name kar_re100 --device-id "$GPU" --seed "$seed" \
        --config "$CFG" --calibration "$CAL" \
        --total-episodes "$EPISODES" --lr 0.0003
    echo "  [$(date '+%H:%M:%S')] Seed $seed done. Sleeping ${DELAY}s..."
    sleep "$DELAY"
done
echo "=== All done ==="
