#!/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] [DRIVE] [WORLD]
#     N     : number of active sheep (1..10), default 10
#     MODE  : "bc" | "rl" | "strombom" | "sequential", default "bc"
#     DRIVE : "differential" | "mecanum", default "differential"
#     WORLD : base world name (without .wbt), default "field"
#             Supported: "field" (rectangular), "field_round" (circular)
#
# Examples:
#   tools/run_webots.sh 10 bc                     # behaviour-cloned MLP, diff drive
#   tools/run_webots.sh 10 rl mecanum             # KL-PPO fine-tune, mecanum wheels
#   tools/run_webots.sh 5 sequential field_round  # analytic baseline, round field
#   tools/run_webots.sh 3 strombom mecanum field_round  # Strömbom, mecanum, round
#
# Notes:
# * bc loads training/runs/bc/policy.zip, rl loads training/runs/rl.
#   Override via HERDING_POLICY_DIR=/path/to/run env var.
# * Conda env "tir" must be active (provides stable-baselines3 + torch).
#
# Headless-ish (no 3D view, fast sim, no modal dialogs):
#   WEBOTS_HEADLESS=1 make webots N=10 MODE=rl DRIVE=mecanum
#   WEBOTS_HEADLESS=1 tools/run_webots.sh 10 rl mecanum
# This passes --no-rendering --minimize --mode=fast --batch to webots.
# Webots still needs a display (Qt); on a machine without one use e.g.:
#   xvfb-run -a env WEBOTS_HEADLESS=1 tools/run_webots.sh 10 rl mecanum
# Optional extra CLI tokens (space-separated):
#   WEBOTS_EXTRA_ARGS="--stdout --stderr" WEBOTS_HEADLESS=1 tools/run_webots.sh 10 rl

set -e
N=${1:-10}
MODE=${2:-bc}
DRIVE=${3:-differential}
WORLD=${4:-field}

if (( N < 0 || N > 10 )); then
    echo "N must be 0..10, got $N" >&2; exit 1
fi
case "$MODE" in
    bc|rl|strombom|sequential|universal|calibrate) ;;
    *) echo "MODE must be bc|rl|strombom|sequential|universal|calibrate, got '$MODE'" >&2; exit 1 ;;
esac
case "$DRIVE" in
    differential|mecanum) ;;
    *) echo "DRIVE must be differential|mecanum, got '$DRIVE'" >&2; exit 1 ;;
esac

ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )/.." && pwd )"
SRC="$ROOT/worlds/${WORLD}.wbt"
if [[ ! -f "$SRC" ]]; then
    echo "World file not found: $SRC" >&2; exit 1
fi
DST="$ROOT/worlds/${WORLD}_test.wbt"

if [[ -n "${HERDING_POLICY_DIR:-}" ]]; then
    RESOLVED_POLICY_DIR="$HERDING_POLICY_DIR"
else
    # Try drive-mode-specific path first, then legacy path.
    if [[ "$MODE" == "rl" ]]; then
        DRIVED="$ROOT/training/runs/rl_${DRIVE}"
        LEGACY="$ROOT/training/runs/rl"
    else
        DRIVED="$ROOT/training/runs/bc_${DRIVE}"
        LEGACY="$ROOT/training/runs/bc"
    fi
    if [[ -d "$DRIVED" ]]; then
        RESOLVED_POLICY_DIR="$DRIVED"
    else
        RESOLVED_POLICY_DIR="$LEGACY"
    fi
fi

cp "$SRC" "$DST"

# Swap robot proto based on drive mode.
# Base worlds reference ShepherdDog (diff-drive). For mecanum we swap in
# ShepherdDogMecanum and inject mecanum contact properties.
if [[ "$DRIVE" == "mecanum" ]]; then
    sed -i 's|"../protos/ShepherdDog.proto"|"../protos/ShepherdDogMecanum.proto"|' "$DST"
    sed -i 's|^ShepherdDog {|ShepherdDogMecanum {|' "$DST"
    # Inject mecanum contact properties into the contactProperties array.
    # Strategy: find the closing '  ]' that ends the contactProperties block
    # (it sits at 2-space indent, immediately before the WorldInfo closing brace)
    # and insert just before it.
    python3 -c "
with open('$DST', 'r') as f:
    txt = f.read()
mec = '''    ContactProperties {
      material1 \"MecanumWheel\"
      coulombFriction [
        1.0
      ]
      bounce 0
      forceDependentSlip [
        0.01
      ]
      softCFM 0.0001
    }
'''
# The contactProperties array closes with '  ]\n}' (2-space indent ] then WorldInfo }).
# Insert the new block just before that closing ].
txt = txt.replace('\n  ]\n}', '\n' + mec + '  ]\n}', 1)
with open('$DST', 'w') as f:
    f.write(txt)
"
fi

# 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" || true)
echo "------------------------------------------------------------"
echo "World      : $DST"
echo "Mode       : $MODE"
echo "Drive      : $DRIVE"
echo "Sheep      : $active active"
echo "Policy dir : $RESOLVED_POLICY_DIR"
echo "------------------------------------------------------------"

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

export HERDING_MODE="$MODE"
export HERDING_POLICY_DIR="$RESOLVED_POLICY_DIR"
export HERDING_DRIVE="$DRIVE"
export HERDING_WORLD="$WORLD"

# 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/.run_done"
mkdir -p "$(dirname "$DONE_FILE")"
rm -f "$DONE_FILE"

if [[ "${WEBOTS_HEADLESS:-}" == "1" ]]; then
    echo "[run_webots] headless flags: --no-rendering --minimize --mode=fast --batch"
    # shellcheck disable=SC2086
    webots --no-rendering --minimize --mode=fast --batch ${WEBOTS_EXTRA_ARGS:-} "$DST" &
else
    # shellcheck disable=SC2086
    webots ${WEBOTS_EXTRA_ARGS:-} "$DST" &
fi
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