#!/usr/bin/env bash
# The zero-thought kickoff for the training box: pulls latest, names the next
# run (highest existing runNN + 1), resumes from the newest committed
# final.zip, and applies the standing flag set below. Idempotent: rerunning
# while training is live just attaches to the tmux session; rerunning after a
# finished run starts the next one, continuing from that run's final.zip.
#
# Usage: ./next_run.sh [extra train.py args...]
# Extras are appended after the standing flags, so they win on conflict
# (e.g. ./next_run.sh --timesteps 5000000, or a fresh start by editing out
# --resume below for one run).
set -euo pipefail
cd "$(dirname "$0")"

# Standing flags for every run. reset-std re-opens exploration each run —
# we mostly start runs because the reward changed. ent-coef holds it open:
# run05 lesson — SB3's 0.0001 default let the action std collapse to ~0.1
# within 4M steps of a resume, re-converging before learning anything new.
TIMESTEPS=20000000
N_PARALLEL=14
SPEEDUP=16
RESET_STD=0.3
ENT_COEF=0.001

SESSION="cosmic-train"

# Live session: attach, never start a second trainer or compute a new name.
if tmux has-session -t "$SESSION" 2>/dev/null; then
    echo "Session '$SESSION' already running — attaching (detach with Ctrl-B then D)."
    exec tmux attach -t "$SESSION"
fi

# Name and resume point come from committed state, so pull before deciding.
git pull --rebase

LAST=$( { ls -d checkpoints/run[0-9]* logs/run[0-9]* 2>/dev/null || true; } \
    | sed -E 's|.*run0*([0-9]+).*|\1|' | sort -n | tail -1)
NEXT=$(printf 'run%02d' $(( ${LAST:-0} + 1 )))

RESUME=$( { ls checkpoints/run[0-9]*/final.zip 2>/dev/null || true; } | sort -V | tail -1)

ARGS=(--timesteps "$TIMESTEPS" --n-parallel "$N_PARALLEL" --speedup "$SPEEDUP" \
    --reset-std "$RESET_STD" --ent-coef "$ENT_COEF")
if [ -n "$RESUME" ]; then
    ARGS+=(--resume "$RESUME")
    echo "Starting $NEXT, resuming from $RESUME"
else
    echo "Starting $NEXT from scratch — no final.zip found in checkpoints/"
fi

exec ./start_training.sh "$NEXT" "${ARGS[@]}" "$@"