#!/bin/bash
# 用法说明：
#   批量对 test_merged_model.sh 推理结果进行跟踪后处理，覆盖三路输入：
#     ROI0:       {case_dir}/roi0/        →  {case_dir}/roi0.json
#     ROI1:       {case_dir}/roi1/        →  {case_dir}/roi1.json
#     Merge:      {case_dir}/merge_json/  →  {case_dir}/merge.json
#   并将三路跟踪结果合并为：
#                                            {case_dir}/combined_tracking.json
#
# 支持两种模式：
#   1) 批量模式（默认）：处理 RESULTS_ROOT/MODEL_NAME 下所有 case
#        bash track_objects.sh
#   2) 单 case 模式：指定 case 目录路径作为第一个参数
#        bash track_objects.sh /data1/dongying/Mono3d/G1M3/cases_regular/model_20260228/some_case

# Set PYTHONPATH to project root for module imports
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH}"

# -----------------------------------------------------------------------
# 配置区：与 test_merged_model.sh 保持一致
# -----------------------------------------------------------------------
# RESULTS_ROOT="/data1/dongying/Mono3d/G1M3/cases_feishu_results"
# MODEL_NAME="torchscript_python_300w_newdata-cncap-v6-ego"
# /data1/dongying/Mono3d/G1Q3/tmp/20251018215336_results/predictions

RESULTS_ROOT="/data1/dongying/Mono3d/G1M3/cases_regular"
RESULTS_ROOT="/data1/dongying/Mono3d/G1M3/pdcl_batch_infer/cases_aeb_v2/CBNA/"
MODEL_NAME="50_15_60"
NUM_WORKERS=16

# 公共跟踪参数
TRACK_CLASSES="${TRACK_CLASSES:-0 1 2 3 4 5 6 7 8 9 10 11 12}"
IOU_THRESH=0.3
MAX_AGE=5
DIST_THRESH=100
MAX_FRAMES="${MAX_FRAMES:-200}"
MODEL_VERSION="${MODEL_VERSION:-model_20260403}"

echo ""
echo "######################################################################"
echo "# Batch tracking and merge: single scan + case-level parallel execution"
echo "######################################################################"

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

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

if [[ -n "$1" ]]; then
    # ===================================================================
    # 单 case 模式：对指定 case 目录进行跟踪 + 合并
    # ===================================================================
    CASE_DIR="$1"
    if [[ ! -d "${CASE_DIR}" ]]; then
        echo "Error: case directory does not exist: ${CASE_DIR}"
        exit 1
    fi
    echo "Single-case mode: ${CASE_DIR}"

    # 逐源跟踪
    for SOURCE in roi0 roi1 merge; do
        INPUT_DIR="${CASE_DIR}/${SOURCE}"
        if [[ "${SOURCE}" == "merge" ]]; then
            OUTPUT_FILE="${CASE_DIR}/merge.json"
        else
            OUTPUT_FILE="${CASE_DIR}/${SOURCE}.json"
        fi
        if [[ -d "${INPUT_DIR}" ]]; then
            echo ""
            echo "--- Tracking ${SOURCE} ---"
            python "${PROJECT_ROOT}/tools/temporal_analysis/track_objects.py" \
                --input-dir "${INPUT_DIR}" \
                --output "${OUTPUT_FILE}" \
                --file-pattern "*.json" \
                --classes ${TRACK_CLASSES} \
                --iou-threshold ${IOU_THRESH} \
                --max-age ${MAX_AGE} \
                --distance-threshold ${DIST_THRESH} \
                "${FRAME_LIMIT_ARGS[@]}" \
                "${MODEL_VERSION_ARGS[@]}"
        else
            echo "Warning: ${INPUT_DIR} not found, skipping ${SOURCE}"
        fi
    done

    # 合并三路跟踪结果
    echo ""
    echo "--- Merging tracking results ---"
    python "${PROJECT_ROOT}/tools/temporal_analysis/merge_tracking_results.py" \
        --roi0  "${CASE_DIR}/roi0.json" \
        --roi1  "${CASE_DIR}/roi1.json" \
        --merge "${CASE_DIR}/merge.json" \
        --output "${CASE_DIR}/combined_tracking.json"

else
    # ===================================================================
    # 批量模式：扫描 RESULTS_ROOT/MODEL_NAME 下所有 case
    # ===================================================================
    python "${PROJECT_ROOT}/tools/temporal_analysis/track_objects.py" \
        --results-root "${RESULTS_ROOT}" \
        --model-name "${MODEL_NAME}" \
        --num-workers "${NUM_WORKERS}" \
        --file-pattern "*.json" \
        --classes ${TRACK_CLASSES} \
        --iou-threshold ${IOU_THRESH} \
        --max-age ${MAX_AGE} \
        --distance-threshold ${DIST_THRESH} \
        "${FRAME_LIMIT_ARGS[@]}" \
        "${MODEL_VERSION_ARGS[@]}"
fi

# -----------------------------------------------------------------------
# 附加3D距离增强匹配（注释状态，按需启用）
# -----------------------------------------------------------------------
# while IFS= read -r -d '' merge_json_dir; do
#     case_dir=$(dirname "$merge_json_dir")
#     output_file="${case_dir}/merge_tracking_3d.json"
#     python "${PROJECT_ROOT}/tools/temporal_analysis/track_objects.py" \
#         --input-dir "$merge_json_dir" \
#         --output "$output_file" \
#         --classes 0 1 2 3 4 5 6 7 8 9 10 11 12 \
#         --iou-threshold 0.3 \
#         --max-age 5 \
#         --distance-threshold 100 \
#         --use-3d \
#         --max-3d-distance 2.0
# done < <(find "${RESULTS_ROOT}/${MODEL_NAME}" -type d -name "merge_json" -print0 | sort -z)