#!/bin/bash
# Launch Webots with N sheep enabled and the chosen controller mode.
# Generates a temporary world file in worlds/field_test.wbt with sheep
# beyond N commented out, sets the env vars the dog controller reads,
# then execs Webots on it.
#
# Usage:
#   tools/run_webots.sh [N] [MODE]
#     N    : number of active sheep (1..10), default 10
#     MODE : "bc" | "rl" | "strombom" | "sequential" | "dagger", default "bc"
#
# Examples:
#   tools/run_webots.sh 10 bc         # BC-trained policy, 10 sheep
#   tools/run_webots.sh 10 rl         # KL-PPO fine-tune of bc, 10 sheep
#   tools/run_webots.sh 5 sequential  # the analytic teacher, 5 sheep
#   tools/run_webots.sh 3 strombom    # canonical baseline, 3 sheep
#
# Notes:
# * The RL mode loads the latest BC policy by default — priority
#   the BC policy (bc/policy.zip) (the controller resolves it).
#   (LiDAR-perception, frame-stack K=4). Override via
#   HERDING_POLICY_DIR=/path/to/run env var.
# * Conda env "tir" must be active (provides stable-baselines3 + torch).

set -e
N=${1:-10}
MODE=${2:-bc}

if (( N < 1 || N > 10 )); then
    echo "N must be 1..10, got $N" >&2; exit 1
fi
case "$MODE" in
    bc|rl|strombom|sequential|dagger) ;;
    *) echo "MODE must be bc|rl|strombom|sequential|dagger, got '$MODE'" >&2; exit 1 ;;
esac
DAGGER_DRIVER=${HERDING_DAGGER_DRIVER:-teacher}

ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )/.." && pwd )"
SRC="$ROOT/worlds/field.wbt"
DST="$ROOT/worlds/field_test.wbt"

cp "$SRC" "$DST"
# Comment out sheep N+1..10 by prefixing the matching Sheep { ... } line.
for i in $(seq $((N+1)) 10); do
    sed -i "s|^Sheep .* \"sheep${i}\".*|# &|" "$DST"
done

active=$(grep -c '^Sheep' "$DST")
echo "------------------------------------------------------------"
echo "World      : $DST"
echo "Mode       : $MODE"
echo "Sheep      : $active active"
echo "Policy dir : ${HERDING_POLICY_DIR:-$ROOT/training/runs/bc}"
echo "------------------------------------------------------------"

# Webots strips HERDING_* env vars from controller subprocesses in some
# setups, so we also write a runtime config file the controller reads.
RESOLVED_POLICY_DIR="${HERDING_POLICY_DIR:-$ROOT/training/runs/bc}"
cat > "$ROOT/herding_runtime.cfg" <<EOF
HERDING_MODE=$MODE
HERDING_POLICY_DIR=$RESOLVED_POLICY_DIR
HERDING_DAGGER_DRIVER=$DAGGER_DRIVER
EOF

export HERDING_MODE="$MODE"
export HERDING_POLICY_DIR="$RESOLVED_POLICY_DIR"

# The controller writes this sentinel when all GT sheep are penned. We
# poll for it and kill Webots so the run finishes cleanly instead of
# idling for minutes after the task is done.
DONE_FILE="$ROOT/training/dagger/.DONE"
mkdir -p "$(dirname "$DONE_FILE")"
rm -f "$DONE_FILE"

webots "$DST" &
WEBOTS_PID=$!

cleanup() {
    kill "$WEBOTS_PID" 2>/dev/null || true
    wait "$WEBOTS_PID" 2>/dev/null || true
    exit 0
}
trap cleanup INT TERM

# Poll for the sentinel; bail when Webots exits on its own or when the
# user closes the window.
while kill -0 "$WEBOTS_PID" 2>/dev/null; do
    if [[ -f "$DONE_FILE" ]]; then
        echo "[run_webots] all sheep penned — closing Webots"
        sleep 1                       # let the controller print its line
        kill "$WEBOTS_PID" 2>/dev/null || true
        break
    fi
    sleep 1
done
wait "$WEBOTS_PID" 2>/dev/null || true