#!/bin/bash
# Resume a canonical V5 case from an episode checkpoint.
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; CASE=""; SEED=43; RESUME_EP=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        --case) CASE="$2"; shift 2 ;;
        --seed) SEED="$2"; shift 2 ;;
        --resume) RESUME_EP="$2"; shift 2 ;;
        --episodes) EPISODES="$2"; shift 2 ;;
        --gpu) GPU="$2"; shift 2 ;;
        *) echo "Unknown: $1"; exit 1 ;;
    esac
done
[[ -z "$CASE" || "$RESUME_EP" -eq 0 ]] && { echo "Usage: resume.sh --case CASE --resume EP [--seed 43]"; exit 1; }
registry() { PYTHONPATH="$REPO_DIR/src" python3 -m drl_pinball.train.case_registry "$CASE" "$1"; }
CFG="$(registry config)" || exit 1
CAL="$(registry calibration)" || exit 1
TRAINER="$(registry trainer)" || exit 1
[[ -f "$CFG" ]] || { echo "ERROR: config not found: $CFG"; exit 1; }
[[ -f "$CAL" ]] || { echo "ERROR: calibration not found: $CAL"; exit 1; }
[[ -f "$TRAINER" ]] || { echo "ERROR: trainer not found: $TRAINER"; exit 1; }
echo "=== Resume: $CASE (seed=$SEED, resume=$RESUME_EP, total=$EPISODES, GPU=$GPU) ==="
echo "  Trainer: $TRAINER"
echo "  Config: $CFG"
echo "  Calibration: $CAL"
EXTRA=()
[[ "$TRAINER" == *train_karman.py ]] && EXTRA=(--lr 0.0003)
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" "${EXTRA[@]}" --resume-from "$RESUME_EP"
echo "=== Done ==="
