#!/usr/bin/env bash
# Full evaluation suite for a trained checkpoint.
#
#   scripts/benchmark.sh runs/jevon-final [arcade-data-dir] [weights]
#
# `weights` selects which checkpoint the whole suite is run against, and
# defaults to balanced.pt -- the one chosen on min(maze_action, snake_action)
# rather than on aggregate eval loss, which the boolean questions dominate.
# Falls back to best.pt for runs trained before that existed. Every number in
# RESULTS.md comes from the single checkpoint named here, so the evaluation,
# the probe and all the play configs cannot silently disagree about which
# model they measured.
#
# Writes eval.json, probe.json and play/*.json under the run directory, then
# regenerates RESULTS.md from them.
set -euo pipefail

RUN="${1:-runs/jevon-final}"
ARCADE="${2:-}"
WEIGHTS="${3:-}"
PY="${PYTHON:-$(dirname "$0")/../../.venv/bin/python}"
[ -x "$PY" ] || PY=python3
PLAY="$RUN/play"
mkdir -p "$PLAY"

if [ -z "$WEIGHTS" ]; then
  if [ -f "$RUN/balanced.pt" ]; then WEIGHTS=balanced.pt; else WEIGHTS=best.pt; fi
fi
[ -f "$RUN/$WEIGHTS" ] || { echo "no such checkpoint: $RUN/$WEIGHTS" >&2; exit 1; }
echo "==> benchmarking $RUN/$WEIGHTS"

echo "==> question-level metrics and calibration"
"$PY" scripts/evaluate.py --checkpoint "$RUN" --weights "$WEIGHTS" \
  --splits data/frozen/val.jsonl,data/frozen/test.jsonl,data/frozen/ood.jsonl \
  --out "$RUN/eval.json"

echo "==> planner field probe"
"$PY" scripts/probe.py --checkpoint "$RUN" --weights "$WEIGHTS" \
  --out "$RUN/probe.json"

play () {  # play <name> <game> <controller> <extra...>
  local name="$1" game="$2" ctrl="$3"; shift 3
  echo "==> play: $name"
  local trace=()
  # --trace-max-steps shortens the replay, never the episode, so the summary is
  # still measured over the full run. Without it a random walk on a 31x31 board
  # contributes its whole 3844-step cap and the arcade data directory grows to
  # megabytes of JSON nobody will scrub through. record.sh in the arcade repo
  # uses the same cap; they must agree or the two repos ship different replays.
  [ -n "$ARCADE" ] && trace=(--trace-out "$ARCADE/${name}.json" --trace-max-steps 400)
  "$PY" scripts/play.py --checkpoint "$RUN" --weights "$WEIGHTS" \
    --game "$game" --controller "$ctrl" \
    --out "$PLAY/${name}_summary.json" "${trace[@]}" "$@"
  [ -n "$ARCADE" ] && cp "$PLAY/${name}_summary.json" "$ARCADE/" || true
}

MAZE_ARGS=(--episodes 3 --maze-sizes 11,21,31)
SNAKE_ARGS=(--episodes 6 --snake-sizes 12)

play maze_model        maze  model        "${MAZE_ARGS[@]}"
# Same network, same absence of search or memory, but sampling from the
# distribution instead of taking its argmax. Now that only legal moves are
# offered, a greedy controller cannot bang into a wall, but it can still cycle:
# the planner is stateless, so a two-cell oscillation repeats forever. Sampling
# breaks that with no memory at all, if the policy is calibrated enough for its
# second choice to be worth taking. Whether it is, is the point of measuring.
play maze_model_sampled maze model "${MAZE_ARGS[@]}" --temperature 1.0
play maze_model_field  maze  model-field  "${MAZE_ARGS[@]}"
play maze_model_memory maze  model+memory "${MAZE_ARGS[@]}"
play maze_reference    maze  reference    "${MAZE_ARGS[@]}"
play maze_random       maze  random       "${MAZE_ARGS[@]}"
# The control for maze_model_memory, and the reason that row cannot be read on
# its own. Identical bookkeeping, no network: whatever this solves, the memory
# solved. On a small board near-exhaustive exploration is enough to match the
# model on solve rate, so the two separate on path length, not on arrival.
play maze_random_memory maze random+memory "${MAZE_ARGS[@]}"
play snake_model       snake model        "${SNAKE_ARGS[@]}"
play snake_reference   snake reference    "${SNAKE_ARGS[@]}"
play snake_random      snake random       "${SNAKE_ARGS[@]}"

echo "==> RESULTS.md"
"$PY" scripts/make_results.py --run "$RUN" --eval "$RUN/eval.json" \
  --probe "$RUN/probe.json" --play-dir "$PLAY"

# The README headline is generated from the same JSON, so it cannot disagree
# with RESULTS.md. Only refreshed when benchmarking the run the README is
# about; benchmarking an experimental arm must not rewrite the front page.
if [ "$RUN" = "runs/jevon-final" ]; then
  echo "==> README headline"
  "$PY" scripts/readme_table.py --run "$RUN"
fi