#!/usr/bin/env bash
set -euo pipefail
shopt -s nullglob

# 用法说明：
#   适配 run_two_roi_exported_onnx_infer_event_json.sh 的输出目录。
#   与按 clip 跟踪不同，这里会对同一个 event_id 目录下所有 clip 的逐帧 JSON
#   先按文件名中的时间戳进行全局排序，再执行一次 event 级跟踪。
#
# 输出目录约定：
#   默认自动解析每个 event 对应的 datename（优先取 _status/event_manifest.json，
#   再回退 scene 级 _status/scene_event_manifest.json），输出为：
#   {event_dir}/{datename}/roi0.json
#   {event_dir}/{datename}/roi1.json
#   {event_dir}/{datename}/merge.json
#   {event_dir}/{datename}/combined_tracking.json
#   {event_dir}/{datename}/frame_order_manifest.json
#
#   如需强制使用固定目录名，可设置：
#     EVENT_TRACKING_DIRNAME=event_tracking
#
# 支持三种模式：
#   1) 批量模式（默认）：遍历 RESULTS_ROOT 下所有 event 目录
#        bash track_objects_exported_onnx_infer_event_json.sh
#   2) 指定根目录或 scene 目录：
#        bash track_objects_exported_onnx_infer_event_json.sh /path/to/event_results_root
#        bash track_objects_exported_onnx_infer_event_json.sh /path/to/event_results_root/scene_xxx
#   3) 单 event 模式：
#        bash track_objects_exported_onnx_infer_event_json.sh /path/to/event_results_root/scene_xxx/event_id_xxx
#        bash track_objects_exported_onnx_infer_event_json.sh /path/to/event_results_root/scene_xxx/rawid_xxx

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"

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

# -----------------------------------------------------------------------
# 配置区：默认与 run_two_roi_exported_onnx_infer_event_json.sh 保持一致
# -----------------------------------------------------------------------
RESULTS_ROOT="${RESULTS_ROOT:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2021_with_cls_ego/model_20260416}"
TRACK_CLASSES="${TRACK_CLASSES:-0 1 2 3 4 5 6 7 8 9 10 11 12}"
IOU_THRESH="${IOU_THRESH:-0.3}"
MAX_AGE="${MAX_AGE:-5}"
MIN_HITS="${MIN_HITS:-1}"
DIST_THRESH="${DIST_THRESH:-100}"
MAX_3D_DISTANCE="${MAX_3D_DISTANCE:-10.0}"
MODEL_VERSION="${MODEL_VERSION:-20260416}"
FILE_PATTERN="${FILE_PATTERN:-*.json}"
ENABLE_USE_3D="${ENABLE_USE_3D:-0}"
MERGE_OUTPUT_NAME="${MERGE_OUTPUT_NAME:-combined_tracking.json}"
MANIFEST_NAME="${MANIFEST_NAME:-frame_order_manifest.json}"
EVENT_TRACKING_DIRNAME="${EVENT_TRACKING_DIRNAME:-}"
SCENE_FILTER="${SCENE_FILTER:-}"
EVENT_FILTER="${EVENT_FILTER:-}"

echo ""
echo "######################################################################"
echo "# Event-level tracking launcher for exported event-json inference"
echo "######################################################################"

MODEL_VERSION_ARGS=()
if [[ -n "${MODEL_VERSION}" ]]; then
    MODEL_VERSION_ARGS+=(--model-version "${MODEL_VERSION}")
fi

USE_3D_ARGS=()
if [[ "${ENABLE_USE_3D}" == "1" ]]; then
    USE_3D_ARGS+=(--use-3d --max-3d-distance "${MAX_3D_DISTANCE}")
fi

is_event_dir() {
    local dir_path="$1"
    [[ -d "${dir_path}" ]] || return 1

    if [[ -f "${dir_path}/_status/event_manifest.json" ]]; then
        return 0
    fi

    local predictions_dir
    for predictions_dir in "${dir_path}"/*/predictions; do
        [[ -d "${predictions_dir}" ]] || continue
        return 0
    done
    return 1
}

emit_batch_event_dirs() {
    local batch_root="$1"
    local manifest_path

    {
        while IFS= read -r -d '' manifest_path; do
            printf '%s\0' "$(dirname "$(dirname "${manifest_path}")")"
        done < <(find "${batch_root}" -type f -path "*/_status/event_manifest.json" -print0)

        # Backward-compatible fallback for older event-json outputs that may
        # not have event manifests yet but still use the historical prefix.
        find "${batch_root}" -type d -name "event_id_*" -print0
    } | sort -zu
}

run_single_event() {
    local event_dir="$1"
    local output_dir_args=()

    echo ""
    echo "Event: ${event_dir}"

    if [[ -n "${EVENT_TRACKING_DIRNAME}" ]]; then
        output_dir_args=(--output-dir "${event_dir}/${EVENT_TRACKING_DIRNAME}")
    fi

    "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/temporal_analysis/track_event_objects.py" \
        --event-dir "${event_dir}" \
        "${output_dir_args[@]}" \
        --file-pattern "${FILE_PATTERN}" \
        --classes ${TRACK_CLASSES} \
        --iou-threshold "${IOU_THRESH}" \
        --max-age "${MAX_AGE}" \
        --min-hits "${MIN_HITS}" \
        --distance-threshold "${DIST_THRESH}" \
        --merge-output-name "${MERGE_OUTPUT_NAME}" \
        --manifest-name "${MANIFEST_NAME}" \
        "${MODEL_VERSION_ARGS[@]}" \
        "${USE_3D_ARGS[@]}"
}

run_batch_root() {
    local batch_root="$1"
    local total_events=0
    local success_events=0
    local failed_events=0

    echo "Batch-root mode: ${batch_root}"
    [[ -n "${SCENE_FILTER}" ]] && echo "Scene filter: ${SCENE_FILTER}"
    [[ -n "${EVENT_FILTER}" ]] && echo "Event filter: ${EVENT_FILTER}"

    while IFS= read -r -d '' event_dir; do
        local scene_name
        local event_name
        scene_name="$(basename "$(dirname "${event_dir}")")"
        event_name="$(basename "${event_dir}")"

        if [[ -n "${SCENE_FILTER}" ]] && [[ ! "${scene_name}" =~ ${SCENE_FILTER} ]]; then
            continue
        fi
        if [[ -n "${EVENT_FILTER}" ]] && [[ ! "${event_name}" =~ ${EVENT_FILTER} ]]; then
            continue
        fi
        if ! is_event_dir "${event_dir}"; then
            continue
        fi

        ((total_events+=1))
        if run_single_event "${event_dir}"; then
            ((success_events+=1))
        else
            ((failed_events+=1))
            printf '[FAIL] scene=%s event=%s\n' "${scene_name}" "${event_name}" >&2
        fi
    done < <(emit_batch_event_dirs "${batch_root}")

    if [[ "${total_events}" -eq 0 ]]; then
        echo "[ERROR] No event directories were found under: ${batch_root}" >&2
        return 1
    fi

    echo ""
    printf '[DONE] events=%d success=%d failed=%d\n' \
        "${total_events}" "${success_events}" "${failed_events}"

    [[ "${failed_events}" -eq 0 ]]
}

TARGET_PATH="${1:-${RESULTS_ROOT}}"

if [[ ! -d "${TARGET_PATH}" ]]; then
    echo "Error: target directory does not exist: ${TARGET_PATH}" >&2
    exit 1
fi

if is_event_dir "${TARGET_PATH}"; then
    run_single_event "${TARGET_PATH}"
else
    run_batch_root "${TARGET_PATH}"
fi
