#!/usr/bin/env bash
# Interactive Webots launcher. Prompts for the experiment knobs
# (mode, drive, world, LiDAR FOV, number of dogs, flock size, GT
# bypass) and then dispatches to tools/run_webots.sh with the
# selected configuration.
#
# Usage:  bash tools/webots_menu.sh

set -e
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
ROOT="$( cd "$SCRIPT_DIR/.." && pwd )"

# Resolve HERDING_PYTHON the same way every other launcher does.
source "$SCRIPT_DIR/setup_env.sh"

# ----- Cosmetics ----------------------------------------------------
if [[ -t 1 ]]; then
    B=$'\e[1m'; D=$'\e[2m'; R=$'\e[0m'
    G=$'\e[32m'; Y=$'\e[33m'; C=$'\e[36m'
else
    B=""; D=""; R=""; G=""; Y=""; C=""
fi

banner () {
    cat <<EOF
${B}${C}┌──────────────────────────────────────────────────────────────────┐
│              Shepherd-dog Webots launcher (interactive)          │
└──────────────────────────────────────────────────────────────────┘${R}

${Y}⚠  Python interpreter${R}
  This script and the Webots controllers read ${B}\$HERDING_PYTHON${R} to
  decide which interpreter to start. Current value:
      ${G}$HERDING_PYTHON${R}
  ${D}If that path is wrong on your machine, edit ${R}${B}tools/setup_env.sh${R}${D}
  or  export HERDING_PYTHON=/path/to/python3  in your shell.${R}

EOF
}

ask_choice () {
    # ask_choice "Prompt" "default" "label1:val1" "label2:val2" ...
    local prompt="$1" default="$2"; shift 2
    local i=1 labels=() values=()
    for opt in "$@"; do
        labels+=("${opt%%:*}")
        values+=("${opt#*:}")
    done
    while true; do
        echo "${B}$prompt${R}"
        for i in "${!labels[@]}"; do
            local marker=" "
            [[ "${values[$i]}" == "$default" ]] && marker="${G}*${R}"
            printf "  %s %d) ${B}%s${R}\n" "$marker" "$((i+1))" "${labels[$i]}"
        done
        printf "  Choice [${G}1-${#labels[@]}${R}, default ${G}%s${R}]: " "$default"
        local raw; read -r raw || true
        raw="${raw:-}"
        if [[ -z "$raw" ]]; then
            CHOICE="$default"; return
        fi
        if [[ "$raw" =~ ^[0-9]+$ ]] && (( raw >= 1 && raw <= ${#labels[@]} )); then
            CHOICE="${values[$((raw-1))]}"; return
        fi
        echo "  ${Y}invalid — try again${R}"
    done
}

ask_int () {
    # ask_int "Prompt" default min max
    local prompt="$1" default="$2" lo="$3" hi="$4"
    while true; do
        printf "${B}%s${R} [${G}%s${R}-${G}%s${R}, default ${G}%s${R}]: " "$prompt" "$lo" "$hi" "$default"
        local raw; read -r raw || true
        raw="${raw:-$default}"
        if [[ "$raw" =~ ^[0-9]+$ ]] && (( raw >= lo && raw <= hi )); then
            CHOICE="$raw"; return
        fi
        echo "  ${Y}must be an integer in [$lo, $hi]${R}"
    done
}

# ----- Prompts ------------------------------------------------------
banner

ask_choice "Mode" "bc" \
    "BC (behaviour-cloned MLP):bc" \
    "RL (KL-PPO fine-tune):rl" \
    "Strömbom (analytic):strombom" \
    "Sequential (analytic):sequential" \
    "Universal teacher (BC source):universal"
MODE="$CHOICE"
echo

ask_choice "Drive" "differential" \
    "Differential (2-wheel):differential" \
    "Mecanum (4-wheel, omnidirectional):mecanum"
DRIVE="$CHOICE"
echo

ask_choice "World" "field" \
    "Rectangular (field):field" \
    "Round (field_round):field_round"
WORLD="$CHOICE"
echo

# LiDAR ablation only applies to differential (mecanum proto has its
# own 140° sensor that we don't fork).
if [[ "$DRIVE" == "differential" ]]; then
    ask_choice "LiDAR FOV" "140" \
        "140° (canonical, ShepherdDog.proto):140" \
        "360° (FOV ablation, ShepherdDog360.proto):360"
    LIDAR="$CHOICE"
else
    LIDAR="140"
    echo "${D}LiDAR: 140° (mecanum drive — no 360° proto variant available)${R}"
fi
echo

ask_choice "Number of shepherd dogs" "1" \
    "1 — solo:1" \
    "2 — axis-split (X-dog + Y-dog):2"
NDOGS="$CHOICE"
echo

if [[ "$NDOGS" == "2" ]]; then
    ask_choice "Axis-split leak (soft mask gain on the off-axis)" "0.3" \
        "0.0 — strict (each dog only moves on its axis; tends to deadlock):0.0" \
        "0.3 — default (off-axis at 30% gain; verified to pen):0.3" \
        "0.5 — softer:0.5" \
        "1.0 — no mask (both dogs run full policy):1.0"
    AXIS_LEAK="$CHOICE"
    echo
fi

ask_int "Flock size (number of sheep)" 5 1 10
N_SHEEP="$CHOICE"
echo

ask_choice "Perception" "lidar" \
    "LiDAR (canonical):lidar" \
    "Ground-truth bypass (HERDING_USE_GT=1):gt"
if [[ "$CHOICE" == "gt" ]]; then USE_GT=1; else USE_GT=0; fi
echo

ask_choice "Headless?" "no" \
    "No — show the Webots window:no" \
    "Yes — headless, fast simulation (xvfb-run):yes"
HEADLESS="$CHOICE"
echo

# ----- Summary ------------------------------------------------------
cat <<EOF
${B}${C}── Launch configuration ──────────────────────────────────────────${R}
  Mode             : ${B}$MODE${R}
  Drive            : ${B}$DRIVE${R}
  World            : ${B}$WORLD${R}
  LiDAR FOV        : ${B}${LIDAR}°${R}
  Dogs             : ${B}$NDOGS${R}$( [[ "$NDOGS" == "2" ]] && echo "  (axis_leak=${B}$AXIS_LEAK${R})" )
  Sheep            : ${B}$N_SHEEP${R}
  Perception       : ${B}$( [[ "$USE_GT" == "1" ]] && echo "GT bypass" || echo "LiDAR" )${R}
  Headless         : ${B}$HEADLESS${R}
${C}──────────────────────────────────────────────────────────────────${R}

EOF
printf "${B}Launch? [Y/n] ${R}"
read -r confirm || true
if [[ "$confirm" =~ ^[Nn] ]]; then
    echo "Aborted."; exit 0
fi

# ----- Dispatch -----------------------------------------------------
export HERDING_LIDAR="$LIDAR"
export HERDING_NDOGS="$NDOGS"
export HERDING_USE_GT="$USE_GT"
[[ -n "${AXIS_LEAK:-}" ]] && export HERDING_AXIS_LEAK="$AXIS_LEAK"
if [[ "$HEADLESS" == "yes" ]]; then
    export WEBOTS_HEADLESS=1
    export WEBOTS_EXTRA_ARGS="--stdout --stderr"
    exec xvfb-run -a bash "$SCRIPT_DIR/run_webots.sh" \
        "$N_SHEEP" "$MODE" "$DRIVE" "$WORLD"
else
    exec bash "$SCRIPT_DIR/run_webots.sh" \
        "$N_SHEEP" "$MODE" "$DRIVE" "$WORLD"
fi