#!/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 scene event-json batch inference.

PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
EVENT_JSON_FILE="${EVENT_JSON_FILE:-${MODEL_INFERENCE_DIR}/examples/events/G1Q3_场地评测数据集.json}"
SCENE="${SCENE:-}"
EVENT_ID_FIELD="${EVENT_ID_FIELD:-data_path}"
EVENT_CLIP_IDS_FIELD="${EVENT_CLIP_IDS_FIELD:-clips}"
CONDITION_FIELDS="${CONDITION_FIELDS:-}"
MAX_RECORDS_PER_CONDITION="${MAX_RECORDS_PER_CONDITION:-0}"
CONDITION_SELECT_STRATEGY="${CONDITION_SELECT_STRATEGY:-first}"
SELECTION_ONLY="${SELECTION_ONLY:-0}"
MAX_EVENTS="${MAX_EVENTS:-0}"
EVENT_CACHE_FILE="${EVENT_CACHE_FILE:-${MODEL_INFERENCE_DIR}/.cache/event_clip_cache.json}"
EVENT_RESOLVE_WORKERS="${EVENT_RESOLVE_WORKERS:-4}"
EVENT_REQUEST_TIMEOUT="${EVENT_REQUEST_TIMEOUT:-60}"
EVENT_REQUEST_RETRIES="${EVENT_REQUEST_RETRIES:-3}"
EVENT_REQUEST_RETRY_BACKOFF_SEC="${EVENT_REQUEST_RETRY_BACKOFF_SEC:-2.0}"
EXPORT_ROOT="${EXPORT_ROOT:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2021_with_cls_ego/event_exports}"
OUTPUT_PREFIX="${OUTPUT_PREFIX:-clip}"
CAMERA_TOPIC="${CAMERA_TOPIC:-camera4}"
MAX_FRAMES_PER_CLIP="${MAX_FRAMES_PER_CLIP:-0}"
LIMIT_CLIPS="${LIMIT_CLIPS:-0}"
SKIP_DONE="${SKIP_DONE:-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_2021_with_cls_ego/model_20260506_keep_fake3d_validate}"
ENABLE_ATTR="${ENABLE_ATTR:-1}"
ENABLE_CROSS_CLASS_MERGE_PRIOR="${ENABLE_CROSS_CLASS_MERGE_PRIOR:-1}"
ENABLE_VRU_MERGE="${ENABLE_VRU_MERGE:-0}"
VIS_CLASSES="${VIS_CLASSES:-}"
VIS_CLASS_NAMES="${VIS_CLASS_NAMES:-}"
SHOW_DISTANCE_LABEL="${SHOW_DISTANCE_LABEL:-0}"
DISTANCE_LABEL_MODE="${DISTANCE_LABEL_MODE:-depth}"
DISTANCE_LABEL_PANELS="${DISTANCE_LABEL_PANELS:-3d}"
EXTRA_ARGS="${EXTRA_ARGS:-}"

# ---------------------------------------------------------------------------
# Auto-tracking: run tracking immediately after inference completes, or set
# TRACK_ONLY=1 to skip clip export/inference and run tracking on OUTPUT_DIR.
# Set CONVERT_ONLY=1 to skip clip export/inference and tracking, then convert.
# Set AUTO_TRACK=1 to enable.  All TRACK_* variables mirror the defaults in
# track_objects_exported_onnx_infer_event_json.sh and can be overridden.
# ---------------------------------------------------------------------------
TRACK_ONLY="${TRACK_ONLY:-0}"
CONVERT_ONLY="${CONVERT_ONLY:-0}"
AUTO_TRACK="${AUTO_TRACK:-1}"
if [[ "${TRACK_ONLY}" == "1" ]] && [[ "${CONVERT_ONLY}" != "1" ]]; then
  AUTO_TRACK="1"
fi
if [[ "${CONVERT_ONLY}" == "1" ]]; then
  AUTO_TRACK="0"
fi
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_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}"
TRACK_MANIFEST_NAME="${TRACK_MANIFEST_NAME:-frame_order_manifest.json}"
TRACK_EVENT_TRACKING_DIRNAME="${TRACK_EVENT_TRACKING_DIRNAME:-}"

# ---------------------------------------------------------------------------
# Post-tracking protocol conversion.
# Only merge.json is converted by default; override CONVERT_MERGE_JSON_NAME if
# another tracking result needs conversion.
# ---------------------------------------------------------------------------
ENABLE_CONVERT="${ENABLE_CONVERT:-0}"
if [[ "${CONVERT_ONLY}" == "1" ]]; then
  ENABLE_CONVERT="1"
fi
CONVERT_STRICT="${CONVERT_STRICT:-1}"
CONVERT_RESULTS_ROOT="${CONVERT_RESULTS_ROOT:-${OUTPUT_DIR}}"
CONVERT_OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT:-${OUTPUT_DIR%/}_objectlist}"
CONVERT_OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT:-parallel_root}"
CONVERT_OUTPUT_DIR_NAME="${CONVERT_OUTPUT_DIR_NAME:-objectlist}"
CONVERT_MERGE_JSON_NAME="${CONVERT_MERGE_JSON_NAME:-merge.json}"
CONVERT_PARALLEL_JOBS="${CONVERT_PARALLEL_JOBS:-4}"
CONVERT_CAM_ID="${CONVERT_CAM_ID:-}"
CONVERT_WRAPPER="${PROJECT_ROOT}/tools/convert_merge_tracking_bundle/convert_merge_tracking_exported_onnx_infer_case.sh"

CMD=(
  "${PYTHON_BIN}" "${MODEL_INFERENCE_DIR}/core/run_two_roi_exported_onnx_infer.py"
  --event-json-file "${EVENT_JSON_FILE}"
  --event-id-field "${EVENT_ID_FIELD}"
  --event-clip-ids-field "${EVENT_CLIP_IDS_FIELD}"
  --event-cache-file "${EVENT_CACHE_FILE}"
  --event-resolve-workers "${EVENT_RESOLVE_WORKERS}"
  --event-request-timeout "${EVENT_REQUEST_TIMEOUT}"
  --event-request-retries "${EVENT_REQUEST_RETRIES}"
  --event-request-retry-backoff-sec "${EVENT_REQUEST_RETRY_BACKOFF_SEC}"
  --export-root "${EXPORT_ROOT}"
  --output-prefix "${OUTPUT_PREFIX}"
  --camera-topic "${CAMERA_TOPIC}"
  --max-frames-per-clip "${MAX_FRAMES_PER_CLIP}"
  --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)
  if [[ -n "${DISTANCE_LABEL_MODE}" ]]; then
    CMD+=(--distance-label-mode "${DISTANCE_LABEL_MODE}")
  fi
  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 [[ -n "${SCENE}" ]]; then
  CMD+=(--scene "${SCENE}")
fi

if [[ -n "${CONDITION_FIELDS}" ]]; then
  # shellcheck disable=SC2206
  CONDITION_FIELDS_ARR=(${CONDITION_FIELDS})
  CMD+=(--condition-fields "${CONDITION_FIELDS_ARR[@]}")
fi

if [[ "${MAX_RECORDS_PER_CONDITION}" != "0" ]]; then
  CMD+=(--max-records-per-condition "${MAX_RECORDS_PER_CONDITION}")
fi

if [[ -n "${CONDITION_SELECT_STRATEGY}" ]]; then
  CMD+=(--condition-select-strategy "${CONDITION_SELECT_STRATEGY}")
fi

if [[ "${SELECTION_ONLY}" == "1" ]]; then
  CMD+=(--selection-only)
fi

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

if [[ "${LIMIT_CLIPS}" != "0" ]]; then
  CMD+=(--limit-clips "${LIMIT_CLIPS}")
fi

if [[ "${SKIP_DONE}" == "1" ]]; then
  CMD+=(--skip-done)
fi

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

CMD+=("$@")

if [[ "${CONVERT_ONLY}" == "1" ]]; then
  echo ""
  echo "######################################################################"
  echo "# CONVERT_ONLY=1, skipping event-json clip export/inference and tracking"
  echo "######################################################################"
  echo "Conversion target: ${CONVERT_RESULTS_ROOT}"
elif [[ "${TRACK_ONLY}" == "1" ]]; then
  echo ""
  echo "######################################################################"
  echo "# TRACK_ONLY=1, skipping event-json clip export/inference"
  echo "######################################################################"
  echo "Tracking root: ${OUTPUT_DIR}"
else
  "${CMD[@]}"
fi

# ---------------------------------------------------------------------------
# Auto-tracking step
# ---------------------------------------------------------------------------
if [[ "${AUTO_TRACK}" == "1" ]]; then
  TRACKING_SCRIPT="${PROJECT_ROOT}/tools/temporal_analysis/track_objects_exported_onnx_infer_event_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 event-level 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}"
    MANIFEST_NAME="${TRACK_MANIFEST_NAME}"
  )

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

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

  # Propagate SCENE as SCENE_FILTER so tracking processes only the same scene.
  if [[ -n "${SCENE}" ]]; then
    TRACK_ENV+=(SCENE_FILTER="${SCENE}")
  fi

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

# ---------------------------------------------------------------------------
# Post-tracking protocol conversion step
# ---------------------------------------------------------------------------
if [[ "${ENABLE_CONVERT}" == "1" ]]; then
  if [[ ! -f "${CONVERT_WRAPPER}" ]]; then
    echo "[ENABLE_CONVERT] ERROR: conversion script not found: ${CONVERT_WRAPPER}" >&2
    exit 1
  fi

  echo ""
  echo "######################################################################"
  echo "# ENABLE_CONVERT=1 — converting tracking ${CONVERT_MERGE_JSON_NAME}"
  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

  CONVERT_ENV=(
    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}"
    CONVERT_PARALLEL_JOBS="${CONVERT_PARALLEL_JOBS}"
    CAM_ID="${CONVERT_CAM_ID}"
  )

  if [[ "${CONVERT_STRICT}" == "1" ]]; then
    env "${CONVERT_ENV[@]}" bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"
  else
    if ! env "${CONVERT_ENV[@]}" bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"; then
      echo "[WARN] Protocol conversion failed, but tracking outputs were kept under ${OUTPUT_DIR}" >&2
    fi
  fi
fi
