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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

POLICY="../policy_35k.onnx"
ROUND="35k"
NAME=""
HISTORY_LEN=5
FLAT_DIM=""
SAMPLES=64
MIN_SAMPLES=32
LOG_PREFIX="rlgym_go1_deploy"
DOCKER_IMAGE="openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8"
COMPARE_LIMIT=64
RUN_CHECKER=1
QUANT="int16"

usage() {
  cat <<'EOF'
Usage:
  ./quantize_policy_x5.sh [options]

Default: quantize Gym policy_35k.onnx as 5-frame/225-dim int16 Gemm BPU model.

Options:
  --policy PATH          ONNX policy path, relative to this directory or absolute
  --round NAME           round label used in output paths, e.g. 15k/25k/30k/35k
  --name NAME            model basename; default is policy filename without .onnx
  --history-len N        observation history length; Gym=5, RobotLab=10
  --flat-dim N           flat input dim; default 45 * history-len
  --samples N            calibration sample count; default 64 for faster mapping
  --min-samples N        minimum valid samples required; default 32
  --log-prefix PREFIX    log dir prefix below logs/, default rlgym_go1_deploy
  --docker-image IMAGE   D-Robotics CPU toolchain image
  --compare-limit N      float ONNX equivalence sample count, default 64
  --quant int16|int8     int16 keeps current conservative path; int8 omits set_all_nodes_int16
  --skip-checker         skip hb_mapper checker before makertbin
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --policy) POLICY="$2"; shift 2 ;;
    --round) ROUND="$2"; shift 2 ;;
    --name) NAME="$2"; shift 2 ;;
    --history-len) HISTORY_LEN="$2"; shift 2 ;;
    --flat-dim) FLAT_DIM="$2"; shift 2 ;;
    --samples) SAMPLES="$2"; shift 2 ;;
    --min-samples) MIN_SAMPLES="$2"; shift 2 ;;
    --log-prefix) LOG_PREFIX="$2"; shift 2 ;;
    --docker-image) DOCKER_IMAGE="$2"; shift 2 ;;
    --compare-limit) COMPARE_LIMIT="$2"; shift 2 ;;
    --quant) QUANT="$2"; shift 2 ;;
    --skip-checker) RUN_CHECKER=0; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done

if [[ -z "${FLAT_DIM}" ]]; then
  FLAT_DIM=$((45 * HISTORY_LEN))
fi

case "${QUANT}" in
  int16|int8) ;;
  *) echo "--quant must be int16 or int8, got: ${QUANT}" >&2; exit 2 ;;
esac

if [[ "${POLICY}" = /* ]]; then
  POLICY_ABS="${POLICY}"
else
  POLICY_ABS="${SCRIPT_DIR}/${POLICY}"
fi
POLICY_ABS="$(cd "$(dirname "${POLICY_ABS}")" && pwd)/$(basename "${POLICY_ABS}")"

if [[ ! -f "${POLICY_ABS}" ]]; then
  echo "Policy not found: ${POLICY_ABS}" >&2
  exit 2
fi
case "${POLICY_ABS}" in
  "${REPO_ROOT}"/*) POLICY_REL="${POLICY_ABS#${REPO_ROOT}/}" ;;
  *) echo "Policy must be inside repo root ${REPO_ROOT}: ${POLICY_ABS}" >&2; exit 2 ;;
esac

if [[ -z "${NAME}" ]]; then
  NAME="$(basename "${POLICY_ABS}" .onnx)"
fi

CAL_DIR="calibration_data_${ROUND}_gym_fast${SAMPLES}"
if [[ "${QUANT}" = "int16" ]]; then
  OUTPUT_DIR="mapper_output_${ROUND}_gemm"
  OUTPUT_PREFIX="${NAME}_int16_gemm"
else
  OUTPUT_DIR="mapper_output_${ROUND}_int8_gemm"
  OUTPUT_PREFIX="${NAME}_int8_gemm"
fi
YAML_FILE="${OUTPUT_PREFIX}.yaml"

echo "[INFO] repo          : ${REPO_ROOT}"
echo "[INFO] policy        : ${POLICY_REL}"
echo "[INFO] name/round    : ${NAME} / ${ROUND}"
echo "[INFO] history/shape : ${HISTORY_LEN} / 1x1x1x${FLAT_DIM}"
echo "[INFO] calibration   : ${CAL_DIR} (${SAMPLES} samples, prefix ${LOG_PREFIX})"
echo "[INFO] output        : ${OUTPUT_DIR}/${OUTPUT_PREFIX}.bin"
echo "[INFO] quant         : ${QUANT}"

docker run --rm --platform linux/amd64 \
  -e POLICY_REL="${POLICY_REL}" \
  -e NAME="${NAME}" \
  -e HISTORY_LEN="${HISTORY_LEN}" \
  -e FLAT_DIM="${FLAT_DIM}" \
  -e SAMPLES="${SAMPLES}" \
  -e MIN_SAMPLES="${MIN_SAMPLES}" \
  -e LOG_PREFIX="${LOG_PREFIX}" \
  -e CAL_DIR="${CAL_DIR}" \
  -e OUTPUT_DIR="${OUTPUT_DIR}" \
  -e OUTPUT_PREFIX="${OUTPUT_PREFIX}" \
  -e YAML_FILE="${YAML_FILE}" \
  -e COMPARE_LIMIT="${COMPARE_LIMIT}" \
  -e RUN_CHECKER="${RUN_CHECKER}" \
  -e QUANT="${QUANT}" \
  -v "${REPO_ROOT}:/workspace/deploy_go1_pro" \
  "${DOCKER_IMAGE}" \
  bash -lc '
    set -euo pipefail
    cd /workspace/deploy_go1_pro/deploy_45dim_rl_gym/bpu_quantization

    POLICY="/workspace/deploy_go1_pro/${POLICY_REL}"
    ACTIONS_ONNX="${NAME}_actions.onnx"
    OPSET_ONNX="${NAME}_opset11.onnx"
    BPU4D_ONNX="${NAME}_bpu4d.onnx"
    GEMM_ONNX="${NAME}_bpu4d_gemm.onnx"

    python3 make_calibration_data.py \
      --logs-root ../../logs \
      --log-prefix "${LOG_PREFIX}" \
      --history-len "${HISTORY_LEN}" \
      --output-dir "${CAL_DIR}" \
      --max-samples "${SAMPLES}" \
      --min-samples "${MIN_SAMPLES}" \
      --overwrite

    python3 keep_actions_output.py \
      --input "${POLICY}" \
      --output "${ACTIONS_ONNX}"

    python3 downgrade_policy_to_opset11.py \
      --input "${ACTIONS_ONNX}" \
      --output "${OPSET_ONNX}"

    python3 make_bpu_4d_onnx.py \
      --input "${OPSET_ONNX}" \
      --output "${BPU4D_ONNX}" \
      --flat-dim "${FLAT_DIM}"

    python3 replace_group_conv_with_gemm.py \
      --input "${BPU4D_ONNX}" \
      --output "${GEMM_ONNX}"

    python3 compare_4d_onnx.py \
      --flat-onnx "${ACTIONS_ONNX}" \
      --bpu4d-onnx "${GEMM_ONNX}" \
      --calibration-dir "${CAL_DIR}" \
      --flat-dim "${FLAT_DIM}" \
      --limit "${COMPARE_LIMIT}"

    if [[ "${QUANT}" = "int16" ]]; then
      OPTIMIZATION_LINE="  optimization: \"set_all_nodes_int16\""
    else
      OPTIMIZATION_LINE=""
    fi

    cat > "${YAML_FILE}" <<YAML
model_parameters:
  onnx_model: "./${GEMM_ONNX}"
  march: "bayes-e"
  layer_out_dump: false
  working_dir: "${OUTPUT_DIR}"
  output_model_file_prefix: "${OUTPUT_PREFIX}"

input_parameters:
  input_name: "obs_4d"
  input_shape: "1x1x1x${FLAT_DIM}"
  input_type_rt: "featuremap"
  input_layout_rt: "NCHW"
  input_type_train: "featuremap"
  input_layout_train: "NCHW"
  norm_type: "no_preprocess"

calibration_parameters:
  cal_data_dir: "./${CAL_DIR}"
  cal_data_type: "float32"
  calibration_type: "default"
${OPTIMIZATION_LINE}
  per_channel: true

compiler_parameters:
  compile_mode: "latency"
  debug: false
  optimize_level: "O3"
YAML

    if [[ "${RUN_CHECKER}" = "1" ]]; then
      hb_mapper checker \
        --model "${GEMM_ONNX}" \
        --model-type onnx \
        --march bayes-e \
        --input-shape obs_4d "1x1x1x${FLAT_DIM}"
    fi

    hb_mapper makertbin \
      --config "${YAML_FILE}" \
      --model-type onnx

    ls -lh "${OUTPUT_DIR}/${OUTPUT_PREFIX}.bin"
  '

echo "[INFO] Done: deploy_45dim_rl_gym/bpu_quantization/${OUTPUT_DIR}/${OUTPUT_PREFIX}.bin"
