#!/usr/bin/env bash
set -Eeuo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
RUNS_DIR="${PROJECT_DIR}/runs/go1-dreamwaq-walk/rslrl"

checkpoint=""
output=""
terrain="stairs_box"
level="0"
forward_speed="0.5"
timeout="30"
realtime_factor="1.0"
export_only=false
manual=false
headless=false
no_validation=false

usage() {
    cat <<'EOF'
Export a DreamWaQ checkpoint and run the MuJoCo stair-course validation.

Usage:
  scripts/run_dreamwaq_sim2sim.sh [options]

Options:
  -c, --checkpoint PATH  Checkpoint to export. Defaults to the newest model_*.pt.
  -o, --output PATH      ONNX output path. Defaults to policy.onnx beside checkpoint.
  -t, --terrain NAME     flat, rough, stairs, dreamwaq, stairs_test,
                         stairs_box, or flat_stairs. Default: stairs_box.
  -l, --level N          Terrain difficulty level passed to sim2sim. Default: 0.
      --speed MPS        Autonomous forward command. Default: 0.5.
      --timeout SEC      Validation timeout in simulation seconds. Default: 30.
      --realtime-factor X
                         Viewer playback speed. Default: 1.0 (real time).
      --manual           Start at zero velocity and use keyboard commands.
      --headless         Run validation without a viewer or real-time delay.
      --no-validation    Run until the viewer closes instead of returning PASS/FAIL.
      --export-only      Export and validate ONNX without starting MuJoCo.
  -h, --help             Show this help.

Examples:
  scripts/run_dreamwaq_sim2sim.sh
  scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt
  scripts/run_dreamwaq_sim2sim.sh --headless
  scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt --speed 0.6
EOF
}

die() {
    echo "[ERROR] $*" >&2
    exit 1
}

while (($#)); do
    case "$1" in
        -c|--checkpoint)
            (($# >= 2)) || die "$1 requires a path"
            checkpoint="$2"
            shift 2
            ;;
        -o|--output)
            (($# >= 2)) || die "$1 requires a path"
            output="$2"
            shift 2
            ;;
        -t|--terrain)
            (($# >= 2)) || die "$1 requires a terrain name"
            terrain="$2"
            shift 2
            ;;
        -l|--level)
            (($# >= 2)) || die "$1 requires an integer"
            level="$2"
            shift 2
            ;;
        --speed)
            (($# >= 2)) || die "$1 requires a number"
            forward_speed="$2"
            shift 2
            ;;
        --timeout)
            (($# >= 2)) || die "$1 requires a number"
            timeout="$2"
            shift 2
            ;;
        --realtime-factor)
            (($# >= 2)) || die "$1 requires a number"
            realtime_factor="$2"
            shift 2
            ;;
        --manual)
            manual=true
            shift
            ;;
        --headless)
            headless=true
            shift
            ;;
        --no-validation)
            no_validation=true
            shift
            ;;
        --export-only)
            export_only=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            die "unknown argument: $1 (use --help)"
            ;;
    esac
done

case "$terrain" in
    flat|rough|stairs|dreamwaq|stairs_test|stairs_box|flat_stairs) ;;
    *) die "unsupported terrain: ${terrain}" ;;
esac
[[ "$level" =~ ^[0-9]+$ ]] || die "level must be a non-negative integer: ${level}"
[[ "$forward_speed" =~ ^[0-9]+([.][0-9]+)?$ ]] || die "speed must be a non-negative number: ${forward_speed}"
[[ "$timeout" =~ ^[0-9]+([.][0-9]+)?$ ]] || die "timeout must be a non-negative number: ${timeout}"
[[ "$realtime_factor" =~ ^[0-9]+([.][0-9]+)?$ ]] || die "realtime-factor must be a positive number: ${realtime_factor}"
[[ "$realtime_factor" != "0" && "$realtime_factor" != "0.0" ]] || die "realtime-factor must be positive"
command -v uv >/dev/null 2>&1 || die "uv is not available in PATH"

if [[ -z "$checkpoint" ]]; then
    [[ -d "$RUNS_DIR" ]] || die "run directory not found: ${RUNS_DIR}"
    checkpoint="$({
        find "$RUNS_DIR" -type f -name 'model_*.pt' -printf '%T@ %p\n'
    } | sort -nr | sed -n '1{s/^[^ ]* //;p;}')"
    [[ -n "$checkpoint" ]] || die "no model_*.pt checkpoint found under ${RUNS_DIR}"
fi

if [[ "$checkpoint" != /* ]]; then
    checkpoint="${PROJECT_DIR}/${checkpoint}"
fi
[[ -f "$checkpoint" ]] || die "checkpoint not found: ${checkpoint}"

if [[ -z "$output" ]]; then
    output="$(dirname -- "$checkpoint")/policy.onnx"
elif [[ "$output" != /* ]]; then
    output="${PROJECT_DIR}/${output}"
fi
mkdir -p -- "$(dirname -- "$output")"

echo "[Pipeline] checkpoint: ${checkpoint}"
echo "[Pipeline] ONNX:       ${output}"

cd -- "$PROJECT_DIR"
uv run scripts/export_dreamwaq_onnx_new.py \
    --checkpoint "$checkpoint" \
    --output "$output"

[[ -s "$output" ]] || die "ONNX export did not produce a non-empty file: ${output}"
echo "[Pipeline] ONNX export complete"

if [[ "$export_only" == true ]]; then
    exit 0
fi

echo "[Pipeline] starting MuJoCo: terrain=${terrain}, level=${level}"
sim_args=(
    --onnx "$output"
    --terrain "$terrain"
    --level "$level"
    --forward-speed "$forward_speed"
    --timeout "$timeout"
    --realtime-factor "$realtime_factor"
)
[[ "$manual" == true ]] && sim_args+=(--manual)
[[ "$headless" == true ]] && sim_args+=(--headless)
[[ "$no_validation" == true ]] && sim_args+=(--no-validation)
exec uv run scripts/dreamwaq_sim2sim_mujoco.py "${sim_args[@]}"