#!/bin/bash
# =============================================================================
# train_illusion.sh — Hydrodynamic Illusion Training (scratch)
#
# Trains all 4 target cylinder sizes on a single GPU, serially.
# Output: output/ill_{1L,15L,075L,2L}_seed43/
#
# Usage:
#   bash train_illusion.sh [--gpu 0] [--episodes 500]
#   bash train_illusion.sh --only 2L
# =============================================================================
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"
registry() { PYTHONPATH="$REPO_DIR/src" python3 -m drl_pinball.train.case_registry "$CASE" "$1"; }
GPU=0; EPISODES=500; DELAY=120; ONLY=""

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

echo "=== Illusion: GPU=$GPU, ep=$EPISODES ==="

for size in 1L 075L 15L 2L; do
    [[ -n "$ONLY" && "$ONLY" != "$size" ]] && continue
    CASE="ill_${size}"
    CAL="$(registry calibration)"
    CFG="$(registry config)"
    TRAINER="$(registry trainer)"
    SEED="$(registry seeds)"
    echo "  [$(date '+%H:%M:%S')] Training $CASE..."
    conda run --no-capture-output -n "$CONDA_ENV" python -u \
        "$TRAINER" \
        --case-name "$CASE" --device-id "$GPU" --seed "$SEED" \
        --config "$CFG" --calibration "$CAL" \
        --total-episodes "$EPISODES"
    echo "  [$(date '+%H:%M:%S')] $CASE done. Sleeping ${DELAY}s..."
    sleep "$DELAY"
done
echo "=== All done ==="
