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

# 用法说明：
#   适配 run_two_roi_exported_onnx_infer_cncap_json.sh 的输出目录，对每个 case 执行：
#     {case_dir}/predictions/roi0/   -> {case_dir}/roi0.json
#     {case_dir}/predictions/roi1/   -> {case_dir}/roi1.json
#     {case_dir}/predictions/merge/  -> {case_dir}/merge.json
#   并继续合并为：
#     {case_dir}/combined_tracking.json
#
# 支持三种模式：
#   1) 批量模式（默认）：遍历 RESULTS_ROOT 下所有包含 predictions/{roi0,roi1,merge} 的 case
#        bash track_objects_exported_onnx_infer_cncap_json.sh
#   2) 指定根目录模式：第一个参数传推理输出根目录
#        bash track_objects_exported_onnx_infer_cncap_json.sh /path/to/inference_output_root
#   3) 单 case 模式：第一个参数传 case 目录，或直接传 case_dir/predictions
#        bash track_objects_exported_onnx_infer_cncap_json.sh /path/to/one_case
#        bash track_objects_exported_onnx_infer_cncap_json.sh /path/to/one_case/predictions

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_cncap_json.sh 保持一致
# -----------------------------------------------------------------------
RESULTS_ROOT="${RESULTS_ROOT:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2024_with_cls_ego/model_20260407}"
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}"
MAX_FRAMES="${MAX_FRAMES:-}"
MODEL_VERSION="${MODEL_VERSION:-}"
MERGE_OUTPUT_NAME="${MERGE_OUTPUT_NAME:-combined_tracking.json}"
FILE_PATTERN="${FILE_PATTERN:-*.json}"
ENABLE_USE_3D="${ENABLE_USE_3D:-0}"
ENABLE_VRU_TRACKING="${ENABLE_VRU_TRACKING:-1}"
VRU_TRACKING_OUTPUT_NAME="${VRU_TRACKING_OUTPUT_NAME:-merge_vru.json}"
VRU_COMBINED_OUTPUT_NAME="${VRU_COMBINED_OUTPUT_NAME:-combined_tracking_with_vru.json}"

echo ""
echo "######################################################################"
echo "# Tracking launcher for exported CNCAP inference outputs"
echo "######################################################################"

MAX_FRAMES_ARGS=()
if [[ -n "${MAX_FRAMES}" ]]; then
    MAX_FRAMES_ARGS+=(--max-frames "${MAX_FRAMES}")
fi

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_predictions_dir() {
    local dir_path="$1"
    [[ -d "${dir_path}" ]] && [[ -d "${dir_path}/roi0" || -d "${dir_path}/roi1" || -d "${dir_path}/merge" ]]
}

is_case_dir() {
    local dir_path="$1"
    is_predictions_dir "${dir_path}/predictions"
}

has_json_files() {
    local dir_path="$1"
    [[ -d "${dir_path}" ]] && compgen -G "${dir_path}/${FILE_PATTERN}" > /dev/null
}

resolve_case_dir() {
    local target_path="$1"
    if is_case_dir "${target_path}"; then
        printf '%s\n' "${target_path}"
        return 0
    fi
    if is_predictions_dir "${target_path}"; then
        dirname "${target_path}"
        return 0
    fi
    return 1
}

run_tracking_for_source() {
    local input_dir="$1"
    local output_file="$2"
    local source_name="$3"
    shift 3
    local extra_tracking_args=("$@")

    if [[ ! -d "${input_dir}" ]]; then
        echo "Warning: ${input_dir} not found, skipping ${source_name}"
        return 0
    fi

    echo ""
    echo "--- Tracking ${source_name} ---"
    "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/temporal_analysis/track_objects.py" \
        --input-dir "${input_dir}" \
        --output "${output_file}" \
        --file-pattern "${FILE_PATTERN}" \
        --classes ${TRACK_CLASSES} \
        --iou-threshold "${IOU_THRESH}" \
        --max-age "${MAX_AGE}" \
        --min-hits "${MIN_HITS}" \
        --distance-threshold "${DIST_THRESH}" \
        "${MAX_FRAMES_ARGS[@]}" \
        "${MODEL_VERSION_ARGS[@]}" \
        "${USE_3D_ARGS[@]}" \
        "${extra_tracking_args[@]}"
}

run_single_case() {
    local case_dir="$1"
    local predictions_dir="${case_dir}/predictions"
    local has_input_source=0
    local merge_args=()

    echo ""
    echo "Case: ${case_dir}"

    if [[ -d "${predictions_dir}/roi0" ]]; then
        run_tracking_for_source "${predictions_dir}/roi0" "${case_dir}/roi0.json" "roi0"
        merge_args+=(--roi0 "${case_dir}/roi0.json")
        has_input_source=1
    else
        echo "Warning: ${predictions_dir}/roi0 not found, roi0 will not participate in merge"
    fi

    if [[ -d "${predictions_dir}/roi1" ]]; then
        run_tracking_for_source "${predictions_dir}/roi1" "${case_dir}/roi1.json" "roi1"
        merge_args+=(--roi1 "${case_dir}/roi1.json")
        has_input_source=1
    else
        echo "Warning: ${predictions_dir}/roi1 not found, roi1 will not participate in merge"
    fi

    if [[ -d "${predictions_dir}/merge" ]]; then
        run_tracking_for_source "${predictions_dir}/merge" "${case_dir}/merge.json" "merge"
        merge_args+=(--merge "${case_dir}/merge.json")
        has_input_source=1
    else
        echo "Warning: ${predictions_dir}/merge not found, merge will not participate in merge"
    fi

    if [[ "${has_input_source}" -eq 0 ]]; then
        echo "Warning: no tracking input directories were found for ${case_dir}, skipping merge"
        return 0
    fi

    echo ""
    echo "--- Merging tracking results ---"
    "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/temporal_analysis/merge_tracking_results.py" \
        "${merge_args[@]}" \
        --output "${case_dir}/${MERGE_OUTPUT_NAME}"

    if [[ "${ENABLE_VRU_TRACKING}" == "1" ]] && has_json_files "${predictions_dir}/merge_vru"; then
        run_tracking_for_source \
            "${predictions_dir}/merge_vru" \
            "${case_dir}/${VRU_TRACKING_OUTPUT_NAME}" \
            "merge_vru" \
            --association-mode vru

        if [[ -f "${case_dir}/${VRU_TRACKING_OUTPUT_NAME}" ]]; then
            echo ""
            echo "--- Merging tracking results with VRU source ---"
            "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/temporal_analysis/merge_tracking_results.py" \
                "${merge_args[@]}" \
                --extra-source "merge_vru=${case_dir}/${VRU_TRACKING_OUTPUT_NAME}" \
                --output "${case_dir}/${VRU_COMBINED_OUTPUT_NAME}"
        else
            echo "Warning: VRU tracking output was not created, skipping ${VRU_COMBINED_OUTPUT_NAME}"
        fi
    elif [[ "${ENABLE_VRU_TRACKING}" == "1" ]]; then
        echo "Info: ${predictions_dir}/merge_vru has no ${FILE_PATTERN} files, skipping VRU tracking"
    fi
}

run_batch_root() {
    local batch_root="$1"
    local total_cases=0
    local success_cases=0
    local failed_cases=0

    echo "Batch-root mode: ${batch_root}"

    while IFS= read -r -d '' predictions_dir; do
        local case_dir
        if ! is_predictions_dir "${predictions_dir}"; then
            continue
        fi
        case_dir="$(dirname "${predictions_dir}")"
        ((total_cases+=1))

        if run_single_case "${case_dir}"; then
            ((success_cases+=1))
        else
            ((failed_cases+=1))
            printf '[FAIL] case=%s\n' "${case_dir}" >&2
        fi
    done < <(
        find "${batch_root}" -type d -name predictions -print0 | sort -z
    )

    if [[ "${total_cases}" -eq 0 ]]; then
        echo "[ERROR] No case directories containing predictions/{roi0,roi1,merge} were found under: ${batch_root}" >&2
        return 1
    fi

    echo ""
    printf '[DONE] cases=%d success=%d failed=%d\n' \
        "${total_cases}" "${success_cases}" "${failed_cases}"

    [[ "${failed_cases}" -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 RESOLVED_CASE_DIR="$(resolve_case_dir "${TARGET_PATH}")"; then
    run_single_case "${RESOLVED_CASE_DIR}"
else
    run_batch_root "${TARGET_PATH}"
fi
