#!/bin/bash
# 批量将跟踪结果可视化到视频帧上。
#
# 流程（每个 case）：
#   直接读取 camera4.bin，将 merge_tracking.json 叠加到每帧，保存可视化图片目录或 mp4。
#   无需保存中间帧，节省磁盘空间。
#
# 目录约定（与 test_merged_model.sh 保持一致）：
#   CASES_ROOT/.../{case_id}/{sub}/sigmastar.1/camera4.bin   ← 原始视频
#   RESULTS_ROOT/MODEL_NAME/{run_name}/merge_tracking.json   ← 跟踪结果
#   RESULTS_ROOT/MODEL_NAME/{run_name}/tracking_vis.mp4      ← 可视化输出（SAVE_VIDEO=true）
#   RESULTS_ROOT/MODEL_NAME/{run_name}/tracking_vis/         ← 可视化图片（SAVE_VIDEO=false）

# 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 / track_objects.sh 保持一致
# -----------------------------------------------------------------------
CASES_ROOT="/data1/dongying/Mono3d/G1M3/cases_coding"
RESULTS_ROOT="/data1/dongying/Mono3d/G1M3/cases_coding_results"
MODEL_NAME="torchscript_python_300w_newdata-cncap-v2"

# 可视化选项
SAVE_VIDEO=false           # true=保存 mp4；false=保存图片目录
VIS_FPS=25                # 输出视频帧率
VIS_CLASS_ID=""           # 过滤类别，例如 "--class-id 0"；留空则显示全部类别

# -----------------------------------------------------------------------
# 批量循环：与 test_merged_model.sh 相同的路径推导逻辑
# -----------------------------------------------------------------------
while IFS= read -r -d '' camera4_bin; do
    sigmastar_dir=$(dirname "$camera4_bin")

    rel_path="${sigmastar_dir#"$CASES_ROOT"/}"
    rel_path="${rel_path%/sigmastar.1}"
    run_name=$(echo "$rel_path" | tr '/ ' '__')

    case_dir="${RESULTS_ROOT}/${MODEL_NAME}/${run_name}"
    tracking_file="${case_dir}/merge_tracking.json"

    # 跟踪结果不存在则跳过（该 case 尚未跑 track_objects.sh）
    if [[ ! -f "$tracking_file" ]]; then
        echo "[SKIP] No merge_tracking.json: $case_dir"
        continue
    fi

    echo ""
    echo "=========================================="
    echo "Case   : $run_name"
    echo "Video  : $camera4_bin"
    echo "Track  : $tracking_file"
    echo "=========================================="

    # ------------------------------------------------------------------
    # 直接读取视频并可视化（不保存中间帧）
    # ------------------------------------------------------------------
    if [[ "$SAVE_VIDEO" == true ]]; then
        output_arg="${case_dir}/tracking_vis.mp4"
        save_video_flag="--save-video"
        echo "[INFO] Visualizing → $output_arg"
    else
        output_arg="${case_dir}/tracking_vis"
        save_video_flag=""
        echo "[INFO] Visualizing → $output_arg/"
    fi

    python "${PROJECT_ROOT}/tools/temporal_analysis/visualize_tracking_boxes.py" \
        --tracking "$tracking_file" \
        --images "$camera4_bin" \
        --output "$output_arg" \
        --fps "$VIS_FPS" \
        $save_video_flag \
        $VIS_CLASS_ID

    if [[ $? -ne 0 ]]; then
        echo "[ERROR] Visualization failed for case: $run_name"
        continue
    fi

    echo "[DONE] $run_name"

done < <(find "$CASES_ROOT" -name "camera4.bin" -print0 | sort -z)

echo ""
echo "All cases processed."