#!/bin/bash
# 用法说明：
#   将跟踪后的聚合 JSON（如 roi0.json）拆分为逐帧 JSON 文件，便于接入评测工具。
#
#   默认对 RESULTS_ROOT/MODEL_NAME 下所有 case 批量处理 roi0.json：
#     {case_dir}/roi0.json  ->  {case_dir}/json_results/*.json
#
# 支持两种模式：
#   1) 批量模式（默认）
#        bash tools/temporal_analysis/split_tracking_json.sh
#   2) 单 case 模式
#        bash tools/temporal_analysis/split_tracking_json.sh /path/to/case
#
# 可选第二个参数指定输入 JSON 文件名：
#        bash tools/temporal_analysis/split_tracking_json.sh /path/to/case roi1.json

set -u

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

# -----------------------------------------------------------------------
# 配置区
# -----------------------------------------------------------------------
RESULTS_ROOT="/data1/dongying/Mono3d/G1M3/cases_regular"
MODEL_NAME="model_20260317"

DEFAULT_INPUT_JSON="merge.json"
DEFAULT_OUTPUT_DIR="json_results_merge"
KEEP_ALL_FIELDS=0
NUM_WORKERS_PER_CASE=8
NUM_CASE_WORKERS=8

split_one_case() {
    local case_dir="$1"
    local input_json_name="$2"
    local output_dir_name="$3"

    if [[ ! -d "${case_dir}" ]]; then
        echo "Error: case directory does not exist: ${case_dir}"
        return 1
    fi

    local input_json="${case_dir}/${input_json_name}"
    local output_dir="${case_dir}/${output_dir_name}"

    if [[ ! -f "${input_json}" ]]; then
        echo "Warning: input tracking JSON not found, skipping: ${input_json}"
        return 0
    fi

    echo ""
    echo "--- Split tracking JSON ---"
    echo "Case          : ${case_dir}"
    echo "Input JSON    : ${input_json}"
    echo "Output dir    : ${output_dir}"

    local cmd=(
        python3 "${PROJECT_ROOT}/tools/temporal_analysis/split_tracking_json.py"
        --input "${input_json}"
        --output-dir "${output_dir}"
        --num-workers "${NUM_WORKERS_PER_CASE}"
    )

    if [[ "${KEEP_ALL_FIELDS}" == "1" ]]; then
        cmd+=(--keep-all-fields)
    fi

    "${cmd[@]}"
}

run_with_limited_jobs() {
    local max_jobs="$1"
    while true; do
        local current_jobs
        current_jobs=$(jobs -pr | wc -l)
        if [[ "${current_jobs}" -lt "${max_jobs}" ]]; then
            break
        fi
        sleep 0.2
    done
}

echo ""
echo "######################################################################"
echo "# Split tracking JSON into per-frame evaluator-friendly JSON files"
echo "######################################################################"

if [[ -n "${1:-}" ]]; then
    CASE_DIR="$1"
    INPUT_JSON_NAME="${2:-${DEFAULT_INPUT_JSON}}"
    echo "Single-case mode: ${CASE_DIR}"
    split_one_case "${CASE_DIR}" "${INPUT_JSON_NAME}" "${DEFAULT_OUTPUT_DIR}"
    exit $?
fi

MODEL_DIR="${RESULTS_ROOT}/${MODEL_NAME}"
if [[ ! -d "${MODEL_DIR}" ]]; then
    echo "Error: model directory not found: ${MODEL_DIR}"
    exit 1
fi

mapfile -t CASE_DIRS < <(find "${MODEL_DIR}" -mindepth 1 -maxdepth 1 -type d | sort)

if [[ ${#CASE_DIRS[@]} -eq 0 ]]; then
    echo "Warning: no case directories found under ${MODEL_DIR}"
    exit 0
fi

echo "Batch mode: found ${#CASE_DIRS[@]} case(s) under ${MODEL_DIR}"
echo "Input JSON name : ${DEFAULT_INPUT_JSON}"
echo "Output dir name : ${DEFAULT_OUTPUT_DIR}"
echo "Case workers    : ${NUM_CASE_WORKERS}"
echo "Frame workers   : ${NUM_WORKERS_PER_CASE}"

declare -a PIDS=()
for case_dir in "${CASE_DIRS[@]}"; do
    run_with_limited_jobs "${NUM_CASE_WORKERS}"
    split_one_case "${case_dir}" "${DEFAULT_INPUT_JSON}" "${DEFAULT_OUTPUT_DIR}" &
    PIDS+=("$!")
done

for pid in "${PIDS[@]}"; do
    if ! wait "${pid}"; then
        echo "Error: one or more split jobs failed."
        exit 1
    fi
done

echo ""
echo "Split complete."