#!/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="flat"
level="0"
export_only=false

usage() {
    cat <<'EOF'
Export a DreamWaQ checkpoint to ONNX and run MuJoCo sim2sim.

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: flat.
  -l, --level N          Terrain difficulty level passed to sim2sim. Default: 0.
      --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 -t dreamwaq -l 3
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
            ;;
        --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}"
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}"
exec uv run scripts/dreamwaq_sim2sim_mujoco.py \
    --onnx "$output" \
    --terrain "$terrain" \
    --level "$level"