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

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

# Launcher for CNCAP JSON batch video inference.

PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
CNCAP_JSON_FILE="${CNCAP_JSON_FILE:-${MODEL_INFERENCE_DIR}/examples/cncap/G1M3_AFS1616_CNCAP-202411.json}"
CNCAP_VALUES_KEY="${CNCAP_VALUES_KEY:-values}"
CNCAP_PATH_PREFIX_SRC="${CNCAP_PATH_PREFIX_SRC:-/mnt/hfs/project-G1M3}"
CNCAP_PATH_PREFIX_DST="${CNCAP_PATH_PREFIX_DST:-/mnt/G1M3}"
VIDEO_STRIDE="${VIDEO_STRIDE:-1}"
EXPORTED_MODEL="${EXPORTED_MODEL:-${PROJECT_ROOT}/runs/export/train_mono3d_two_roi_20260506-keep_fake_3d_branch/merged_model.torchscript}"
OUTPUT_DIR="${OUTPUT_DIR:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2024_with_cls_ego/model_20260506_keep_fake3d_validate}"
MAX_IMAGES="${MAX_IMAGES:-0}"
ENABLE_ATTR="${ENABLE_ATTR:-1}"
ENABLE_CROSS_CLASS_MERGE_PRIOR="${ENABLE_CROSS_CLASS_MERGE_PRIOR:-1}"
ENABLE_VRU_MERGE="${ENABLE_VRU_MERGE:-0}"
DEVICE="${DEVICE:-cuda:2}"
VIS_CLASSES="${VIS_CLASSES:-}"
VIS_CLASS_NAMES="${VIS_CLASS_NAMES:-}"
EXTRA_ARGS="${EXTRA_ARGS:-}"

# ---------------------------------------------------------------------------
# Auto-tracking: run tracking immediately after inference completes.
# Set AUTO_TRACK=1 to enable.  All TRACK_* variables mirror the defaults in
# track_objects_exported_onnx_infer_cncap_json.sh and can be overridden.
# ---------------------------------------------------------------------------
AUTO_TRACK="${AUTO_TRACK:-1}"
ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING:-0}"
TRACK_CLASSES="${TRACK_CLASSES:-0 1 2 3 4 5 6 7 8 9 10 11 12 17 18 19}"
TRACK_IOU_THRESH="${TRACK_IOU_THRESH:-0.3}"
TRACK_MAX_AGE="${TRACK_MAX_AGE:-5}"
TRACK_MIN_HITS="${TRACK_MIN_HITS:-1}"
TRACK_DIST_THRESH="${TRACK_DIST_THRESH:-100}"
TRACK_MAX_3D_DISTANCE="${TRACK_MAX_3D_DISTANCE:-10.0}"
TRACK_MAX_FRAMES="${TRACK_MAX_FRAMES:-}"
TRACK_MODEL_VERSION="${TRACK_MODEL_VERSION:-20260506}"
TRACK_FILE_PATTERN="${TRACK_FILE_PATTERN:-*.json}"
TRACK_ENABLE_USE_3D="${TRACK_ENABLE_USE_3D:-0}"
TRACK_MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME:-combined_tracking.json}"

if [[ -z "${CNCAP_JSON_FILE}" ]]; then
  echo "CNCAP_JSON_FILE is required." >&2
  exit 1
fi

# /data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2024_with_cls_ego/model_20260407/20251115/CVE/CPLA_RL_AEB_20_5_1_20251115140740/predictions/roi0/camera4_345630.json

CMD=(
  "${PYTHON_BIN}" "${MODEL_INFERENCE_DIR}/core/run_two_roi_exported_onnx_infer.py"
  --cncap-json-file "${CNCAP_JSON_FILE}"
  --cncap-values-key "${CNCAP_VALUES_KEY}"
  --cncap-path-prefix-src "${CNCAP_PATH_PREFIX_SRC}"
  --cncap-path-prefix-dst "${CNCAP_PATH_PREFIX_DST}"
  --video-stride "${VIDEO_STRIDE}"
  --exported-model "${EXPORTED_MODEL}"
  --output-dir "${OUTPUT_DIR}"
  --device "${DEVICE}"
)

if [[ "${ENABLE_ATTR}" == "1" ]]; then
  CMD+=(--enable-attr)
fi

if [[ "${ENABLE_CROSS_CLASS_MERGE_PRIOR}" == "1" ]]; then
  CMD+=(--enable-cross-class-merge-prior)
fi

if [[ "${ENABLE_VRU_MERGE}" == "1" ]]; then
  CMD+=(--enable-vru-merge)
fi

if [[ -n "${VIS_CLASSES}" ]]; then
  # shellcheck disable=SC2206
  VIS_CLASSES_ARR=(${VIS_CLASSES})
  CMD+=(--vis-classes "${VIS_CLASSES_ARR[@]}")
fi

if [[ -n "${VIS_CLASS_NAMES}" ]]; then
  # shellcheck disable=SC2206
  VIS_CLASS_NAMES_ARR=(${VIS_CLASS_NAMES})
  CMD+=(--vis-class-names "${VIS_CLASS_NAMES_ARR[@]}")
fi

if [[ "${MAX_IMAGES}" != "0" ]]; then
  CMD+=(--max-images "${MAX_IMAGES}")
fi

if [[ -n "${EXTRA_ARGS}" ]]; then
  # shellcheck disable=SC2206
  EXTRA_ARR=(${EXTRA_ARGS})
  CMD+=("${EXTRA_ARR[@]}")
fi

CMD+=("$@")
"${CMD[@]}"

# ---------------------------------------------------------------------------
# Auto-tracking step
# ---------------------------------------------------------------------------
if [[ "${AUTO_TRACK}" == "1" ]]; then
  TRACKING_SCRIPT="${PROJECT_ROOT}/tools/temporal_analysis/track_objects_exported_onnx_infer_cncap_json.sh"
  if [[ ! -f "${TRACKING_SCRIPT}" ]]; then
    echo "[AUTO_TRACK] ERROR: tracking script not found: ${TRACKING_SCRIPT}" >&2
    exit 1
  fi

  echo ""
  echo "######################################################################"
  echo "# AUTO_TRACK=1 — starting CNCAP tracking on ${OUTPUT_DIR}"
  echo "######################################################################"

  TRACK_ENV=(
    PYTHON_BIN="${PYTHON_BIN}"
    RESULTS_ROOT="${OUTPUT_DIR}"
    TRACK_CLASSES="${TRACK_CLASSES}"
    IOU_THRESH="${TRACK_IOU_THRESH}"
    MAX_AGE="${TRACK_MAX_AGE}"
    MIN_HITS="${TRACK_MIN_HITS}"
    DIST_THRESH="${TRACK_DIST_THRESH}"
    MAX_3D_DISTANCE="${TRACK_MAX_3D_DISTANCE}"
    FILE_PATTERN="${TRACK_FILE_PATTERN}"
    ENABLE_USE_3D="${TRACK_ENABLE_USE_3D}"
    MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME}"
    ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING}"
  )

  if [[ -n "${TRACK_MAX_FRAMES}" ]]; then
    TRACK_ENV+=(MAX_FRAMES="${TRACK_MAX_FRAMES}")
  fi

  if [[ -n "${TRACK_MODEL_VERSION}" ]]; then
    TRACK_ENV+=(MODEL_VERSION="${TRACK_MODEL_VERSION}")
  fi

  env "${TRACK_ENV[@]}" bash "${TRACKING_SCRIPT}"
fi
