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

# 用法说明：
#   批量对 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) 批量模式（默认）：遍历 BATCH_ROOT 下所有 场景/模型 目录
#        bash track_objects_cncap.sh
#   2) 批量模式（指定根目录）：第一个参数传场景根目录
#        bash track_objects_cncap.sh /data1/dongying/Mono3d/G1M3/pdcl_batch_infer/cases_aeb_v2
#   3) 单 case 模式：第一个参数传具体 case 目录路径
#        bash track_objects_cncap.sh /path/to/one_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:-}"
PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/yolov5/bin/python}"

# -----------------------------------------------------------------------
# 配置区
# -----------------------------------------------------------------------
# 批量根目录，目录结构约定为：
#   BATCH_ROOT/<scene_name>/<model_name>/<case_name>/{roi0,roi1,merge_json}
BATCH_ROOT="${BATCH_ROOT:-/data1/dongying/Mono3d/G1M3/pdcl_batch_infer/cases_aeb_v2}"

# 可选过滤，留空表示不过滤：
#   SCENE_FILTER="CBNA"
#   MODEL_FILTER="50_15_60"
SCENE_FILTER="${SCENE_FILTER:-}"
MODEL_FILTER="${MODEL_FILTER:-}"
NUM_WORKERS="${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
MODEL_VERSION="${MODEL_VERSION:-}"
MERGE_OUTPUT_NAME="${MERGE_OUTPUT_NAME:-combined_tracking.json}"

echo ""
echo "######################################################################"
echo "# Batch tracking and merge: auto scene/model traversal"
echo "######################################################################"

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

is_case_dir() {
    local dir_path="$1"
    [[ -d "${dir_path}/roi0" || -d "${dir_path}/roi1" || -d "${dir_path}/merge_json" ]]
}

run_single_case() {
    local case_dir="$1"

    echo "Single-case mode: ${case_dir}"

    for source in roi0 roi1 merge; do
        local input_dir="${case_dir}/${source}"
        local output_file
        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_BIN}" "${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} \
                "${MODEL_VERSION_ARGS[@]}"
        else
            echo "Warning: ${input_dir} not found, skipping ${source}"
        fi
    done

    echo ""
    echo "--- Merging tracking results ---"
    "${PYTHON_BIN}" "${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}/${MERGE_OUTPUT_NAME}"
}

run_one_model() {
    local results_root="$1"
    local model_name="$2"

    echo ""
    printf '[RUN] scene=%s model=%s\n' "$(basename "${results_root}")" "${model_name}"

    "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/temporal_analysis/track_objects.py" \
        --results-root "${results_root}" \
        --model-name "${model_name}" \
        --num-workers "${NUM_WORKERS}" \
        --merge-output-name "${MERGE_OUTPUT_NAME}" \
        --file-pattern "*.json" \
        --classes ${TRACK_CLASSES} \
        --iou-threshold ${IOU_THRESH} \
        --max-age ${MAX_AGE} \
        --distance-threshold ${DIST_THRESH} \
        "${MODEL_VERSION_ARGS[@]}"
}

run_batch_root() {
    local batch_root="$1"
    local total_models=0
    local success_models=0
    local failed_models=0
    local matched_scenes=0

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

    for scene_dir in "${batch_root}"/*; do
        [[ -d "${scene_dir}" ]] || continue
        local scene_name
        scene_name="$(basename "${scene_dir}")"

        if [[ -n "${SCENE_FILTER}" ]] && [[ ! "${scene_name}" =~ ${SCENE_FILTER} ]]; then
            continue
        fi
        ((matched_scenes+=1))

        local found_model=0
        for model_dir in "${scene_dir}"/*; do
            [[ -d "${model_dir}" ]] || continue
            local model_name
            model_name="$(basename "${model_dir}")"

            if [[ -n "${MODEL_FILTER}" ]] && [[ ! "${model_name}" =~ ${MODEL_FILTER} ]]; then
                continue
            fi

            found_model=1
            ((total_models+=1))
            if run_one_model "${scene_dir}" "${model_name}"; then
                ((success_models+=1))
            else
                ((failed_models+=1))
                printf '[FAIL] scene=%s model=%s\n' "${scene_name}" "${model_name}" >&2
            fi
        done

        if [[ "${found_model}" == "0" ]]; then
            printf '[SKIP] empty scene dir or no matched model: %s\n' "${scene_name}"
        fi
    done

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

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

    echo ""
    printf '[DONE] models=%d success=%d failed=%d\n' \
        "${total_models}" "${success_models}" "${failed_models}"

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

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

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

if is_case_dir "${TARGET_PATH}"; then
    run_single_case "${TARGET_PATH}"
else
    run_batch_root "${TARGET_PATH}"
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_BIN}" "${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 "${BATCH_ROOT}" -mindepth 3 -maxdepth 3 -type d -name "merge_json" -print0 | sort -z)