#!/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 batch video-root inference.
#
# Supported postprocess-only mode:
#   POSTPROCESS_ONLY=1 OUTPUT_DIR=<inference-output-root> bash run_two_roi_exported_onnx_infer_video_root.sh

PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
VIDEO_ROOT_DIR="${VIDEO_ROOT_DIR:-/data1/dongying/Mono3d/G1Q3/tmp/20260507}"
# VIDEO_ROOT_DIR="${VIDEO_ROOT_DIR:-/data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/CNCAP/CSTA_LN}"
VIDEO_STRIDE="${VIDEO_STRIDE:-1}"
EXPORTED_MODEL="${EXPORTED_MODEL:-${PROJECT_ROOT}/runs/export/train_mono3d_two_roi_20260506-drop_fake_3d_branch/merged_model.torchscript}"
# OUTPUT_DIR="${OUTPUT_DIR:-/data1/dongying/Mono3d/G1M3/test_outputs/video_case_model_20260416_with_cls}"
OUTPUT_DIR="${OUTPUT_DIR:-/data1/dongying/Mono3d/G1Q3/tmp/20260507_results/model_20260507_drop_fake3d}"
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:-1}"
VIS_CLASSES="${VIS_CLASSES:-0 1 2 3 4 5 6 7 8 9 10 11 12 17 18 19}"
VIS_CLASS_NAMES="${VIS_CLASS_NAMES:-}"
SHOW_DISTANCE_LABEL="${SHOW_DISTANCE_LABEL:-1}"
# depth_lateral keeps the existing longitudinal z label and adds signed
# camera-frame lateral x offset. Override with depth/lateral/xz/etc if needed.
DISTANCE_LABEL_MODE="${DISTANCE_LABEL_MODE:-depth_lateral}"
DISTANCE_LABEL_PANELS="${DISTANCE_LABEL_PANELS:-3d}"
EXTRA_ARGS="${EXTRA_ARGS:-}"
POSTPROCESS_ONLY="${POSTPROCESS_ONLY:-0}"

# Optional post-inference tracking stage. When enabled, the launcher reuses
# the exported-inference tracking wrapper, which runs track_objects.py over
# predictions/{roi0,roi1,merge} and merges the results per case.
ENABLE_TRACKING="${ENABLE_TRACKING:-1}"
ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING:-0}"
TRACK_STRICT="${TRACK_STRICT:-1}"
TRACK_RESULTS_ROOT="${TRACK_RESULTS_ROOT:-${OUTPUT_DIR}}"
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_ENABLE_USE_3D="${TRACK_ENABLE_USE_3D:-0}"
TRACK_MAX_3D_DISTANCE="${TRACK_MAX_3D_DISTANCE:-10.0}"
TRACK_MAX_FRAMES="${TRACK_MAX_FRAMES:-}"
TRACK_MODEL_VERSION="${TRACK_MODEL_VERSION:-20260427}"
TRACK_FILE_PATTERN="${TRACK_FILE_PATTERN:-*.json}"
TRACK_MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME:-combined_tracking.json}"
TRACKING_WRAPPER="${PROJECT_ROOT}/tools/temporal_analysis/track_objects_exported_onnx_infer_case.sh"

# Optional post-tracking protocol-conversion stage.
# Runs convert_merge_tracking_exported_onnx_infer_case.sh on the tracked
# video-root output tree. By default, each case writes converted protobuf
# artifacts into its own objectlist/ subdirectory.
ENABLE_CONVERT="${ENABLE_CONVERT:-1}"
CONVERT_STRICT="${CONVERT_STRICT:-1}"
CONVERT_RESULTS_ROOT="${CONVERT_RESULTS_ROOT:-${TRACK_RESULTS_ROOT}}"
CONVERT_OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT:-}"
CONVERT_OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT:-case_subdir}"
CONVERT_OUTPUT_DIR_NAME="${CONVERT_OUTPUT_DIR_NAME:-objectlist}"
CONVERT_MERGE_JSON_NAME="${CONVERT_MERGE_JSON_NAME:-${TRACK_MERGE_OUTPUT_NAME}}"
CONVERT_CAM_ID="${CONVERT_CAM_ID:-}"
CONVERT_WRAPPER="${PROJECT_ROOT}/tools/convert_merge_tracking_bundle/convert_merge_tracking_exported_onnx_infer_case.sh"
ENABLE_CONVERT_VRU="${ENABLE_CONVERT_VRU:-0}"
CONVERT_VRU_OUTPUT_ROOT="${CONVERT_VRU_OUTPUT_ROOT:-}"
CONVERT_VRU_OUTPUT_LAYOUT="${CONVERT_VRU_OUTPUT_LAYOUT:-case_subdir}"
CONVERT_VRU_OUTPUT_DIR_NAME="${CONVERT_VRU_OUTPUT_DIR_NAME:-objectlist_vru}"
CONVERT_VRU_MERGE_JSON_NAME="${CONVERT_VRU_MERGE_JSON_NAME:-merge_vru.json}"

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

CMD=(
  "${PYTHON_BIN}" "${MODEL_INFERENCE_DIR}/core/run_two_roi_exported_onnx_infer.py"
  --video-root-dir "${VIDEO_ROOT_DIR}"
  --video-stride "${VIDEO_STRIDE}"
  --exported-model "${EXPORTED_MODEL}"
  --output-dir "${OUTPUT_DIR}"
)

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 [[ "${SHOW_DISTANCE_LABEL}" == "1" ]]; then
  CMD+=(--show-distance-label)
  CMD+=(--distance-label-mode "${DISTANCE_LABEL_MODE}")

  if [[ -n "${DISTANCE_LABEL_PANELS}" ]]; then
    # shellcheck disable=SC2206
    DISTANCE_LABEL_PANELS_ARR=(${DISTANCE_LABEL_PANELS})
    CMD+=(--distance-label-panels "${DISTANCE_LABEL_PANELS_ARR[@]}")
  fi
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+=("$@")

if [[ "${POSTPROCESS_ONLY}" == "1" ]]; then
  echo ""
  echo "######################################################################"
  echo "# POSTPROCESS_ONLY=1, skipping inference"
  echo "######################################################################"
  echo "Postprocess root: ${OUTPUT_DIR}"
else
  "${CMD[@]}"
fi

if [[ "${ENABLE_TRACKING}" == "1" ]]; then
  echo ""
  echo "######################################################################"
  echo "# Post-inference tracking for exported video-root outputs"
  echo "######################################################################"
  echo "Tracking root: ${TRACK_RESULTS_ROOT}"

  if [[ "${TRACK_STRICT}" == "1" ]]; then
    PYTHON_BIN="${PYTHON_BIN}" \
    RESULTS_ROOT="${TRACK_RESULTS_ROOT}" \
    TRACK_CLASSES="${TRACK_CLASSES}" \
    IOU_THRESH="${TRACK_IOU_THRESH}" \
    MAX_AGE="${TRACK_MAX_AGE}" \
    MIN_HITS="${TRACK_MIN_HITS}" \
    DIST_THRESH="${TRACK_DIST_THRESH}" \
    ENABLE_USE_3D="${TRACK_ENABLE_USE_3D}" \
    MAX_3D_DISTANCE="${TRACK_MAX_3D_DISTANCE}" \
    MAX_FRAMES="${TRACK_MAX_FRAMES}" \
    MODEL_VERSION="${TRACK_MODEL_VERSION}" \
    FILE_PATTERN="${TRACK_FILE_PATTERN}" \
    MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME}" \
    ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING}" \
    bash "${TRACKING_WRAPPER}" "${TRACK_RESULTS_ROOT}"
  else
    if ! PYTHON_BIN="${PYTHON_BIN}" \
      RESULTS_ROOT="${TRACK_RESULTS_ROOT}" \
      TRACK_CLASSES="${TRACK_CLASSES}" \
      IOU_THRESH="${TRACK_IOU_THRESH}" \
      MAX_AGE="${TRACK_MAX_AGE}" \
      MIN_HITS="${TRACK_MIN_HITS}" \
      DIST_THRESH="${TRACK_DIST_THRESH}" \
      ENABLE_USE_3D="${TRACK_ENABLE_USE_3D}" \
      MAX_3D_DISTANCE="${TRACK_MAX_3D_DISTANCE}" \
      MAX_FRAMES="${TRACK_MAX_FRAMES}" \
      MODEL_VERSION="${TRACK_MODEL_VERSION}" \
      FILE_PATTERN="${TRACK_FILE_PATTERN}" \
      MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME}" \
      ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING}" \
      bash "${TRACKING_WRAPPER}" "${TRACK_RESULTS_ROOT}"; then
      echo "[WARN] Tracking failed, but inference outputs were kept under ${OUTPUT_DIR}" >&2
    fi
  fi
elif [[ "${ENABLE_CONVERT}" != "1" ]]; then
  exit 0
fi

if [[ "${ENABLE_CONVERT}" != "1" ]]; then
  exit 0
fi

echo ""
echo "######################################################################"
echo "# Post-tracking protocol conversion for exported video-root outputs"
echo "######################################################################"
echo "Conversion target: ${CONVERT_RESULTS_ROOT}"
echo "Conversion layout: ${CONVERT_OUTPUT_LAYOUT}"

if [[ "${CONVERT_OUTPUT_LAYOUT}" == "case_subdir" ]]; then
  echo "Conversion subdir: ${CONVERT_OUTPUT_DIR_NAME}"
elif [[ -n "${CONVERT_OUTPUT_ROOT}" ]]; then
  echo "Conversion output root: ${CONVERT_OUTPUT_ROOT}"
fi

if [[ "${CONVERT_STRICT}" == "1" ]]; then
  PYTHON_BIN="${PYTHON_BIN}" \
  OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT}" \
  OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT}" \
  OUTPUT_DIR_NAME="${CONVERT_OUTPUT_DIR_NAME}" \
  MERGE_JSON_NAME="${CONVERT_MERGE_JSON_NAME}" \
  CAM_ID="${CONVERT_CAM_ID}" \
  bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"
else
  if ! PYTHON_BIN="${PYTHON_BIN}" \
    OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT}" \
    OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT}" \
    OUTPUT_DIR_NAME="${CONVERT_OUTPUT_DIR_NAME}" \
    MERGE_JSON_NAME="${CONVERT_MERGE_JSON_NAME}" \
    CAM_ID="${CONVERT_CAM_ID}" \
    bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"; then
    echo "[WARN] Protocol conversion failed, but tracking outputs were kept under ${OUTPUT_DIR}" >&2
  fi
fi

if [[ "${ENABLE_CONVERT_VRU}" == "1" ]]; then
  echo ""
  echo "######################################################################"
  echo "# Post-tracking protocol conversion for VRU tracking outputs"
  echo "######################################################################"
  echo "VRU conversion target: ${CONVERT_RESULTS_ROOT}"
  echo "VRU conversion layout: ${CONVERT_VRU_OUTPUT_LAYOUT}"
  if [[ "${CONVERT_VRU_OUTPUT_LAYOUT}" == "case_subdir" ]]; then
    echo "VRU conversion subdir: ${CONVERT_VRU_OUTPUT_DIR_NAME}"
  elif [[ -n "${CONVERT_VRU_OUTPUT_ROOT}" ]]; then
    echo "VRU conversion output root: ${CONVERT_VRU_OUTPUT_ROOT}"
  fi

  if [[ -d "${CONVERT_RESULTS_ROOT}" ]] && [[ -n "$(find "${CONVERT_RESULTS_ROOT}" -type f -name "${CONVERT_VRU_MERGE_JSON_NAME}" -print -quit)" ]]; then
    if [[ "${CONVERT_STRICT}" == "1" ]]; then
      PYTHON_BIN="${PYTHON_BIN}" \
      OUTPUT_ROOT="${CONVERT_VRU_OUTPUT_ROOT}" \
      OUTPUT_LAYOUT="${CONVERT_VRU_OUTPUT_LAYOUT}" \
      OUTPUT_DIR_NAME="${CONVERT_VRU_OUTPUT_DIR_NAME}" \
      MERGE_JSON_NAME="${CONVERT_VRU_MERGE_JSON_NAME}" \
      CAM_ID="${CONVERT_CAM_ID}" \
      bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"
    else
      if ! PYTHON_BIN="${PYTHON_BIN}" \
        OUTPUT_ROOT="${CONVERT_VRU_OUTPUT_ROOT}" \
        OUTPUT_LAYOUT="${CONVERT_VRU_OUTPUT_LAYOUT}" \
        OUTPUT_DIR_NAME="${CONVERT_VRU_OUTPUT_DIR_NAME}" \
        MERGE_JSON_NAME="${CONVERT_VRU_MERGE_JSON_NAME}" \
        CAM_ID="${CONVERT_CAM_ID}" \
        bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"; then
        echo "[WARN] VRU protocol conversion failed, but tracking outputs were kept under ${OUTPUT_DIR}" >&2
      fi
    fi
  else
    echo "Info: no ${CONVERT_VRU_MERGE_JSON_NAME} files found under ${CONVERT_RESULTS_ROOT}, skipping VRU protocol conversion"
  fi
fi
