#!/usr/bin/env bash
# Launch one ScienceWorld training run. Run from anywhere; paths are resolved
# relative to the bundle root.
#
#     bash scripts/run.sh <method>
#
#     <method> in: ftb | opd | guided_opd | tcod_b2f | tcod_f2b
#
# One method = one full run on 8 GPUs (4 student rollout + 2 teacher + 2 FSDP
# trainer), 200 training steps, a checkpoint every 50. Run methods one at a
# time: each one wants all 8 GPUs and its own Ray cluster.
#
# Console output is tee'd to outputs/logs/<method>_<timestamp>.log.

set -euo pipefail

METHOD=${1:-}
BUNDLE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)

case "$METHOD" in
    ftb|opd|guided_opd|tcod_b2f|tcod_f2b) ;;
    *) echo "usage: $0 <ftb|opd|guided_opd|tcod_b2f|tcod_f2b>" >&2; exit 2 ;;
esac
CFG="$BUNDLE/configs/$METHOD.yaml"

# Trinity resolves the configs' relative paths (data/, models/, outputs/)
# against the process CWD, so the bundle root has to be the CWD.
cd "$BUNDLE"

# ------------------------------------------------------------ pre-flight ----
for p in tcod/trinity models/Qwen3-1.7B/config.json models/Qwen3-32B/config.json \
         data/scienceworld/train.jsonl data/scienceworld/test.jsonl "$CFG"; do
    [ -e "$p" ] || { echo "[run] missing $p -- run 'bash scripts/setup.sh' first" >&2; exit 1; }
done
command -v trinity >/dev/null 2>&1 \
    || { echo "[run] 'trinity' CLI not on PATH -- activate the environment used in setup" >&2; exit 1; }
command -v java >/dev/null 2>&1 \
    || { echo "[run] no 'java' on PATH -- ScienceWorld needs Java 17+" >&2; exit 1; }

# The overlay lives in tcod/trinity; re-apply it every launch so an edited
# overlay/ never silently disagrees with what actually runs.
cp -R overlay/trinity/. tcod/trinity/

mkdir -p outputs/tmp outputs/ray_tmp outputs/buffers outputs/checkpoints outputs/logs

# --------------------------------------------------------------- monitor ----
# The four baseline configs use tensorboard. ftb.yaml is the upstream release
# file and uses wandb; without a login, wandb.init() hard-fails the explorer
# Ray actor after a 90s timeout (WANDB_MODE=offline does not reach the actor).
if [ "$METHOD" = "ftb" ] && grep -q 'monitor_type: wandb' "$CFG" \
        && [ -z "${WANDB_API_KEY:-}" ] && [ ! -f "$HOME/.netrc" ]; then
    echo "[run] WARNING: configs/ftb.yaml uses wandb but no WANDB_API_KEY / ~/.netrc found." >&2
    echo "[run]          Run 'wandb login', or change its monitor_type to tensorboard." >&2
fi

# ------------------------------------------------------------------- ray ----
if ! ray status >/dev/null 2>&1; then
    echo "[run] starting a local ray head"
    ray start --head
fi

# ------------------------------------------------------------------- run ----
LOG="outputs/logs/${METHOD}_$(date +%Y%m%d_%H%M%S).log"
echo "[run] method=$METHOD"
echo "[run] config=$CFG"
echo "[run] cwd=$PWD"
echo "[run] log=$LOG"
trinity run --config "$CFG" 2>&1 | tee "$LOG"