#!/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

# Make sure HERDING_PYTHON is resolved and on PATH so Webots inherits
# the right interpreter (controllers/{shepherd_dog,sheep}/runtime.ini
# both read $HERDING_PYTHON via env-var expansion).
source "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )/setup_env.sh"

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
    # The training pipeline writes policies to:
    #   training/runs/{bc,rl}_<drive>_<world>
    # Try that first; fall back to the drive-only and finally the
    # bare-mode legacy paths so older policy checkouts still load.
    if [[ "$MODE" == "rl" ]]; then
        BASE="rl"
    else
        BASE="bc"
    fi
    for CAND in \
        "$ROOT/training/runs/${BASE}_${DRIVE}_${WORLD}" \
        "$ROOT/training/runs/${BASE}_${DRIVE}" \
        "$ROOT/training/runs/${BASE}"
    do
        if [[ -d "$CAND" ]]; then
            RESOLVED_POLICY_DIR="$CAND"
            break
        fi
    done
    : "${RESOLVED_POLICY_DIR:=$ROOT/training/runs/${BASE}_${DRIVE}_${WORLD}}"
fi

cp "$SRC" "$DST"

# LiDAR FOV variant. Mecanum defaults to 360° (the trained mecanum
# target); diff defaults to 140°. Override with HERDING_LIDAR=140 or
# HERDING_LIDAR=360 for ablations.
if [[ -z "${HERDING_LIDAR:-}" ]]; then
    if [[ "$DRIVE" == "mecanum" ]]; then
        LIDAR_VARIANT="360"
    else
        LIDAR_VARIANT="140"
    fi
else
    LIDAR_VARIANT="$HERDING_LIDAR"
fi
if [[ "$LIDAR_VARIANT" != "140" && "$LIDAR_VARIANT" != "360" ]]; then
    echo "HERDING_LIDAR must be 140 or 360, got '$LIDAR_VARIANT'" >&2; exit 1
fi
export HERDING_LIDAR="$LIDAR_VARIANT"

# Swap robot proto based on drive mode + LiDAR variant.
# Base worlds reference ShepherdDog (diff-drive 140°). The four
# combinations the launcher supports:
#   diff   + 140°  → ShepherdDog.proto          (default)
#   diff   + 360°  → ShepherdDog360.proto       (FOV ablation for diff)
#   mecanum+ 140°  → ShepherdDogMecanum.proto
#   mecanum+ 360°  → ShepherdDogMecanum360.proto (the trained mecanum target)
if [[ "$DRIVE" == "mecanum" && "$LIDAR_VARIANT" == "360" ]]; then
    sed -i 's|"../protos/ShepherdDog.proto"|"../protos/ShepherdDogMecanum360.proto"|' "$DST"
    sed -i 's|^ShepherdDog {|ShepherdDogMecanum360 {|' "$DST"
elif [[ "$DRIVE" == "mecanum" ]]; then
    sed -i 's|"../protos/ShepherdDog.proto"|"../protos/ShepherdDogMecanum.proto"|' "$DST"
    sed -i 's|^ShepherdDog {|ShepherdDogMecanum {|' "$DST"
elif [[ "$LIDAR_VARIANT" == "360" ]]; then
    sed -i 's|"../protos/ShepherdDog.proto"|"../protos/ShepherdDog360.proto"|' "$DST"
    sed -i 's|^ShepherdDog {|ShepherdDog360 {|' "$DST"
fi
if [[ "$DRIVE" == "mecanum" ]]; then
    # Wheel-ground friction. The chassis is driven kinematically by
    # the Supervisor (see drive_mecanum in controllers/shepherd_dog),
    # so these properties only affect wheel-spin visuals, not the
    # robot's motion. coulombFriction 2.0 plus a small
    # forceDependentSlip keeps the rollers from locking up against
    # the ground.
    python3 -c "
with open('$DST', 'r') as f:
    txt = f.read()
mec = '''    ContactProperties {
      material1 \"MecanumWheelA\"
      coulombFriction [
        2.0
      ]
      bounce 0
      forceDependentSlip [
        0.005
      ]
      softCFM 0.0001
    }
    ContactProperties {
      material1 \"MecanumWheelB\"
      coulombFriction [
        2.0
      ]
      bounce 0
      forceDependentSlip [
        0.005
      ]
      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

# Dual-dog axis split. When HERDING_NDOGS=2 the launcher replaces the
# single dog node in the world with two named dogs whose customData
# carries the axis assignment (x or y); the controller masks the
# off-axis component of every action.
NDOGS="${HERDING_NDOGS:-1}"
if [[ "$NDOGS" != "1" && "$NDOGS" != "2" ]]; then
    echo "HERDING_NDOGS must be 1 or 2, got '$NDOGS'" >&2; exit 1
fi
if [[ "$NDOGS" == "2" ]]; then
    DOG_NODE_NAME="ShepherdDog"
    if [[ "$DRIVE" == "mecanum" ]]; then
        DOG_NODE_NAME="ShepherdDogMecanum"
    elif [[ "$LIDAR_VARIANT" == "360" ]]; then
        DOG_NODE_NAME="ShepherdDog360"
    fi
    python3 - "$DST" "$DOG_NODE_NAME" <<'PY'
import re, sys
path, node = sys.argv[1], sys.argv[2]
with open(path) as f:
    txt = f.read()
# Match the single existing dog block from "ShepherdDog{,360,Mecanum} {"
# up to its closing "}" on a line by itself.
pattern = re.compile(rf"^{re.escape(node)} \{{\n(.*?\n)^\}}\n", re.MULTILINE | re.DOTALL)
m = pattern.search(txt)
if m is None:
    sys.exit(f"[run_webots] could not locate single-dog block ({node}) for split")
two_dogs = (
    f"{node} {{\n"
    f"  translation -4 -10 0.5\n"
    f"  rotation 0 0 1 1.5708\n"
    f'  name "ShepherdDogX"\n'
    f'  customData "axis=x"\n'
    f'  controller "shepherd_dog"\n'
    f"}}\n"
    f"{node} {{\n"
    f"  translation 4 -10 0.5\n"
    f"  rotation 0 0 1 1.5708\n"
    f'  name "ShepherdDogY"\n'
    f'  customData "axis=y"\n'
    f'  controller "shepherd_dog"\n'
    f"}}\n"
)
with open(path, 'w') as f:
    f.write(txt[:m.start()] + two_dogs + txt[m.end():])
PY
fi
export HERDING_NDOGS="$NDOGS"

active=$(grep -c '^Sheep' "$DST" || true)
ndog=$(grep -cE '^(ShepherdDog|ShepherdDog360|ShepherdDogMecanum) \{' "$DST" || true)
echo "------------------------------------------------------------"
echo "World      : $DST"
echo "Mode       : $MODE"
echo "Drive      : $DRIVE"
echo "LiDAR      : ${LIDAR_VARIANT}°"
echo "Dogs       : $ndog (axis-split=${NDOGS})"
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_LIDAR=$LIDAR_VARIANT
HERDING_NDOGS=$NDOGS
HERDING_AXIS_LEAK=${HERDING_AXIS_LEAK:-0.3}
HERDING_USE_GT=${HERDING_USE_GT:-0}
HERDING_SEED=${HERDING_SEED:-}
EOF

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

# 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