#!/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 single case / eval-dir inference.
#
# Supported modes:
#   1. CASE_DIR=<case-dir>
#   2. EVAL_DIR=<eval-root>
#   3. INPUT_DIR=<case-dir-or-eval-root>
#   4. POSTPROCESS_ONLY=1 OUTPUT_DIR=<inference-output-root>
#   5. ENABLE_PARALLEL=1 GPU_IDS="0 1" INPUT_DIR=<eval-root>

PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
CASE_DIR="${CASE_DIR:-}"
EVAL_DIR="${EVAL_DIR:-}"

# /data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/DL_KPI_SCENE
# /data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/OP_KPI_SCENE
# /data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/OP_KPI_TRAFFIC_SCENARIO
# /data1/xdzhu/Testdata_0129
# /data1/xdzhu/Testdata_0129019b4415-d171-7ae9-88ed-0065c4c6e1e0

INPUT_DIR="${INPUT_DIR:-${EVAL_DIR:-/data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/OP_KPI_SCENE}}"
# INPUT_DIR="${INPUT_DIR:-${CASE_DIR}}"
# The core inference script now auto-detects whether the exported model keeps the edge branch.
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/KPI/OP_KPI_SCENE/model_20260506_all_cases}"
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:-}"
ATTR_DEVICE="${ATTR_DEVICE:-}"
VIS_CLASSES="${VIS_CLASSES:-}"
VIS_CLASS_NAMES="${VIS_CLASS_NAMES:-}"
SAVE_VISUALIZATION="${SAVE_VISUALIZATION:-0}"
EXTRA_ARGS="${EXTRA_ARGS:-}"
POSTPROCESS_ONLY="${POSTPROCESS_ONLY:-1}"
ENABLE_PARALLEL="${ENABLE_PARALLEL:-0}"
GPU_IDS="${GPU_IDS:-}"
NUM_SHARDS="${NUM_SHARDS:-1}"
SHARD_INDEX="${SHARD_INDEX:-0}"
SHARD_STRATEGY="${SHARD_STRATEGY:-round_robin}"
CASE_INDEX_START="${CASE_INDEX_START:-}"
CASE_INDEX_END="${CASE_INDEX_END:-}"
PARALLEL_LOG_DIR="${PARALLEL_LOG_DIR:-${OUTPUT_DIR%/}/parallel_logs}"
FORWARD_ARGS=("$@")

# 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:-0}"
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}"
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:-20260506}"
TRACK_FILE_PATTERN="${TRACK_FILE_PATTERN:-*.json}"
TRACK_MERGE_OUTPUT_NAME="${TRACK_MERGE_OUTPUT_NAME:-combined_tracking.json}"
TRACK_PARALLEL_JOBS="${TRACK_PARALLEL_JOBS:-4}"
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
# case/eval-root output tree. By default, each case writes converted protocol
# artifacts into its own objectlist/ subdirectory. Set CONVERT_OUTPUT_ROOT to
# write only converted artifacts into a separate root while preserving case
# relative paths.
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:-/data1/dongying/Mono3d/G1Q3/model_inference/KPI/OP_KPI_SCENE/model_20260506_objectlist}"
if [[ -z "${CONVERT_OUTPUT_LAYOUT+x}" ]]; then
  if [[ -n "${CONVERT_OUTPUT_ROOT}" ]]; then
    CONVERT_OUTPUT_LAYOUT="parallel_root"
  else
    CONVERT_OUTPUT_LAYOUT="case_subdir"
  fi
else
  CONVERT_OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT:-case_subdir}"
fi
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_CASE_LIST_FILE="${CONVERT_CASE_LIST_FILE:-}"
  CONVERT_PARALLEL_JOBS="${CONVERT_PARALLEL_JOBS:-8}"
CONVERT_WRAPPER="${PROJECT_ROOT}/tools/convert_merge_tracking_bundle/convert_merge_tracking_exported_onnx_infer_case.sh"

is_eval_case_dir() {
  local case_dir="$1"

  [[ -d "${case_dir}/images" ]] || return 1
  [[ -f "${case_dir}/calib/L2_calib/camera4.json" || -f "${case_dir}/calib/camera4.json" ]]
}

collect_eval_case_dirs() {
  local eval_root="${1%/}"
  local first_level_dir
  local second_level_dir

  while IFS= read -r -d '' first_level_dir; do
    if is_eval_case_dir "${first_level_dir}"; then
      printf '%s\n' "${first_level_dir}"
      continue
    fi

    while IFS= read -r -d '' second_level_dir; do
      if is_eval_case_dir "${second_level_dir}"; then
        printf '%s\n' "${second_level_dir}"
      fi
    done < <(find "${first_level_dir}" -mindepth 1 -maxdepth 1 -type d -print0 | sort -z)
  done < <(find "${eval_root}" -mindepth 1 -maxdepth 1 -type d -print0 | sort -z)
}

write_convert_case_list() {
  local target_mode="$1"
  local convert_root="${2%/}"
  local list_file="$3"
  local eval_root="${EVAL_DIR:-${INPUT_DIR}}"
  local -a case_dirs=()
  local case_index_start="${CASE_INDEX_START:-0}"
  local case_index_end="${CASE_INDEX_END:-}"
  local total_cases
  local relative_case_dir
  local case_dir
  local index

  mkdir -p "$(dirname "${list_file}")"
  : > "${list_file}"

  if [[ "${target_mode}" == "case" ]]; then
    printf '%s\n' "${convert_root}" > "${list_file}"
    return 0
  fi

  mapfile -t case_dirs < <(collect_eval_case_dirs "${eval_root}")
  total_cases="${#case_dirs[@]}"
  if [[ -z "${case_index_end}" ]]; then
    case_index_end="${total_cases}"
  fi

  for ((index = case_index_start; index < case_index_end; index++)); do
    case_dir="${case_dirs[${index}]}"
    relative_case_dir="${case_dir#${eval_root%/}/}"
    printf '%s\n' "${convert_root}/${relative_case_dir}" >> "${list_file}"
  done
}

resolve_target_mode() {
  if [[ -n "${CASE_DIR}" ]]; then
    printf 'case\n'
  elif [[ -n "${EVAL_DIR}" ]]; then
    printf 'eval\n'
  elif [[ -d "${INPUT_DIR}/images" ]]; then
    printf 'case\n'
  else
    printf 'eval\n'
  fi
}

launch_parallel_eval_workers() {
  local target_mode="$1"
  local num_shards="${NUM_SHARDS}"
  local failed=0
  local -a gpu_ids_arr=()
  local -a pids=()
  local -a log_files=()

  if [[ "${target_mode}" != "eval" ]]; then
    echo "ENABLE_PARALLEL=1 only supports eval-root inference." >&2
    echo "Please point INPUT_DIR/EVAL_DIR to an eval root instead of a single case." >&2
    exit 1
  fi
  if [[ -z "${GPU_IDS}" ]]; then
    echo "ENABLE_PARALLEL=1 requires GPU_IDS, for example: GPU_IDS=\"0 1 2 3\"." >&2
    exit 1
  fi

  # shellcheck disable=SC2206
  gpu_ids_arr=(${GPU_IDS})
  if [[ "${num_shards}" == "1" ]]; then
    num_shards="${#gpu_ids_arr[@]}"
  fi
  if [[ "${num_shards}" -lt 2 ]]; then
    echo "Parallel mode requires at least 2 shards, got NUM_SHARDS=${num_shards}." >&2
    exit 1
  fi
  if [[ "${#gpu_ids_arr[@]}" -ne "${num_shards}" ]]; then
    echo "GPU_IDS count (${#gpu_ids_arr[@]}) must match NUM_SHARDS (${num_shards}) in ENABLE_PARALLEL mode." >&2
    exit 1
  fi

  mkdir -p "${PARALLEL_LOG_DIR}"

  echo ""
  echo "######################################################################"
  echo "# Parallel eval-dir inference"
  echo "######################################################################"
  echo "Eval root       : ${INPUT_DIR}"
  echo "Output root     : ${OUTPUT_DIR}"
  echo "Shards          : ${num_shards}"
  echo "Shard strategy  : ${SHARD_STRATEGY}"
  echo "GPU ids         : ${GPU_IDS}"
  echo "Worker logs     : ${PARALLEL_LOG_DIR}"

  for shard_index in "${!gpu_ids_arr[@]}"; do
    local gpu_id="${gpu_ids_arr[${shard_index}]}"
    local device="cuda:${gpu_id}"
    local attr_device="${ATTR_DEVICE:-${device}}"
    local log_file="${PARALLEL_LOG_DIR}/shard_${shard_index}_gpu_${gpu_id}.log"

    echo "Launching shard ${shard_index}/${num_shards} on ${device} -> ${log_file}"
    ENABLE_PARALLEL=0 \
    ENABLE_TRACKING=0 \
    ENABLE_CONVERT=0 \
    NUM_SHARDS="${num_shards}" \
    SHARD_INDEX="${shard_index}" \
    SHARD_STRATEGY="${SHARD_STRATEGY}" \
    DEVICE="${device}" \
    ATTR_DEVICE="${attr_device}" \
    bash "${BASH_SOURCE[0]}" "${FORWARD_ARGS[@]}" >"${log_file}" 2>&1 &

    pids+=("$!")
    log_files+=("${log_file}")
  done

  for shard_index in "${!pids[@]}"; do
    if wait "${pids[${shard_index}]}"; then
      echo "Shard ${shard_index}/${num_shards} finished successfully. Log: ${log_files[${shard_index}]}"
    else
      echo "[ERROR] Shard ${shard_index}/${num_shards} failed. Log: ${log_files[${shard_index}]}" >&2
      failed=1
    fi
  done

  if [[ "${failed}" == "1" ]]; then
    echo "Parallel inference did not finish cleanly. Postprocess was skipped." >&2
    exit 1
  fi
}

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

if [[ -n "${DEVICE}" ]]; then
  CMD+=(--device "${DEVICE}")
fi

if [[ -n "${ATTR_DEVICE}" ]]; then
  CMD+=(--attr-device "${ATTR_DEVICE}")
fi

if [[ -n "${CASE_INDEX_START}" ]]; then
  CMD+=(--case-index-start "${CASE_INDEX_START}")
fi

if [[ -n "${CASE_INDEX_END}" ]]; then
  CMD+=(--case-index-end "${CASE_INDEX_END}")
fi

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 [[ "${NUM_SHARDS}" != "1" ]] || [[ "${SHARD_INDEX}" != "0" ]] || [[ "${SHARD_STRATEGY}" != "round_robin" ]]; then
  CMD+=(--num-shards "${NUM_SHARDS}" --shard-index "${SHARD_INDEX}" --shard-strategy "${SHARD_STRATEGY}")
fi

if [[ -n "${CASE_DIR}" ]]; then
  CMD+=(--case-dir "${CASE_DIR}")
elif [[ -n "${EVAL_DIR}" ]]; then
  CMD+=(--eval-dir "${EVAL_DIR}")
elif [[ -d "${INPUT_DIR}/images" ]]; then
  CMD+=(--case-dir "${INPUT_DIR}")
else
  CMD+=(--eval-dir "${INPUT_DIR}")
fi

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

if [[ "${SAVE_VISUALIZATION}" != "1" ]]; then
  CMD+=(--skip-visualizations)
fi

CMD+=("$@")

TARGET_MODE="$(resolve_target_mode)"

if [[ "${ENABLE_PARALLEL}" == "1" ]] && [[ "${POSTPROCESS_ONLY}" != "1" ]]; then
  launch_parallel_eval_workers "${TARGET_MODE}"
  POSTPROCESS_ONLY="1"
fi

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 case/eval-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}" \
    TRACK_PARALLEL_JOBS="${TRACK_PARALLEL_JOBS}" \
    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}" \
      TRACK_PARALLEL_JOBS="${TRACK_PARALLEL_JOBS}" \
      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

CONVERT_OUTPUT_ROOT_EFFECTIVE="${CONVERT_OUTPUT_ROOT}"
CONVERT_OUTPUT_LAYOUT_EFFECTIVE="${CONVERT_OUTPUT_LAYOUT}"
CONVERT_CASE_LIST_FILE_EFFECTIVE="${CONVERT_CASE_LIST_FILE}"

case "${CONVERT_OUTPUT_LAYOUT}" in
  same_dir)
    if [[ -f "${CONVERT_RESULTS_ROOT}" ]] || [[ "$(basename "${CONVERT_RESULTS_ROOT}")" == "${CONVERT_MERGE_JSON_NAME}" ]]; then
      CONVERT_OUTPUT_ROOT_EFFECTIVE="$(dirname "${CONVERT_RESULTS_ROOT}")"
    else
      CONVERT_OUTPUT_ROOT_EFFECTIVE="${CONVERT_RESULTS_ROOT%/}"
    fi
    CONVERT_OUTPUT_LAYOUT_EFFECTIVE="parallel_root"
    ;;
  case_subdir|parallel_root)
    ;;
  *)
    echo "Unsupported CONVERT_OUTPUT_LAYOUT: ${CONVERT_OUTPUT_LAYOUT}" >&2
    echo "Expected one of: same_dir, case_subdir, parallel_root" >&2
    exit 1
    ;;
esac

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

if [[ "${CONVERT_OUTPUT_LAYOUT}" == "same_dir" ]]; then
  echo "Conversion output: same directory as tracking results"
elif [[ "${CONVERT_OUTPUT_LAYOUT}" == "case_subdir" ]]; then
  echo "Conversion subdir: ${CONVERT_OUTPUT_DIR_NAME}"
elif [[ -n "${CONVERT_OUTPUT_ROOT_EFFECTIVE}" ]]; then
  echo "Conversion output root: ${CONVERT_OUTPUT_ROOT_EFFECTIVE}"
fi

if [[ -z "${CONVERT_CASE_LIST_FILE_EFFECTIVE}" ]]; then
  CONVERT_CASE_LIST_FILE_EFFECTIVE="${OUTPUT_DIR%/}/_status/convert_case_list.txt"
fi

if [[ -d "${CONVERT_RESULTS_ROOT}" ]]; then
  write_convert_case_list "${TARGET_MODE}" "${CONVERT_RESULTS_ROOT}" "${CONVERT_CASE_LIST_FILE_EFFECTIVE}"
  echo "Conversion case list: ${CONVERT_CASE_LIST_FILE_EFFECTIVE}"
fi

if [[ "${CONVERT_STRICT}" == "1" ]]; then
  PYTHON_BIN="${PYTHON_BIN}" \
  CASE_LIST_FILE="${CONVERT_CASE_LIST_FILE_EFFECTIVE}" \
  OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT_EFFECTIVE}" \
  OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT_EFFECTIVE}" \
  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}" \
  bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"
else
  if ! PYTHON_BIN="${PYTHON_BIN}" \
    CASE_LIST_FILE="${CONVERT_CASE_LIST_FILE_EFFECTIVE}" \
    OUTPUT_ROOT="${CONVERT_OUTPUT_ROOT_EFFECTIVE}" \
    OUTPUT_LAYOUT="${CONVERT_OUTPUT_LAYOUT_EFFECTIVE}" \
    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}" \
    bash "${CONVERT_WRAPPER}" "${CONVERT_RESULTS_ROOT}"; then
    echo "[WARN] Protocol conversion failed, but tracking outputs were kept under ${OUTPUT_DIR}" >&2
  fi
fi
