"""
批量可视化脚本

用途：批量读取多帧图像和对应的真值标签文件，生成可视化结果

使用方法：
    python scripts_for_gt/visualize_batch.py \
        --image-dir /path/to/images \
        --label-dir /path/to/labels \
        --output output_dir \
        --roi 704 352 \
        --virtual-fx 500 \
        --max-samples 100
"""

import argparse
from pathlib import Path
import sys

# Add project root to path
FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from visualize_single_frame import (
    DEFAULT_CLASS_NAMES,
    format_target_classes,
    resolve_target_class_ids,
    visualize_single_frame,
)


def resolve_label_path(label_dir, image_stem, preferred_extension=None):
    """按同名文件自动查找标签，支持 .txt / .json。"""
    candidate_extensions = []
    if preferred_extension:
        normalized_extension = preferred_extension if preferred_extension.startswith(".") else f".{preferred_extension}"
        candidate_extensions.append(normalized_extension)

    for extension in (".txt", ".json"):
        if extension not in candidate_extensions:
            candidate_extensions.append(extension)

    for extension in candidate_extensions:
        candidate_path = label_dir / f"{image_stem}{extension}"
        if candidate_path.exists():
            return candidate_path

    return None


def visualize_batch(
    image_dir,
    label_dir,
    output_dir,
    calib_dir=None,
    roi_size=None,
    virtual_fx=None,
    ori_img_size=None,
    max_samples=None,
    sample_interval=1,
    file_extension=".jpg",
    label_extension=None,
    target_class_ids=None,
):
    """批量可视化图像
    
    Args:
        image_dir: 图像目录
        label_dir: 标签目录
        output_dir: 输出目录
        calib_dir: 标定文件目录（包含L2_calib/camera4.json）
        roi_size: ROI尺寸 (width, height)
        virtual_fx: 虚拟焦距
        ori_img_size: 原始图像尺寸 (width, height)
        max_samples: 最大处理样本数
        sample_interval: 采样间隔，每隔 N 帧处理一帧
        file_extension: 图像文件扩展名
        label_extension: 标签文件扩展名，默认自动在 .txt/.json 中查找
        target_class_ids: 目标类别ID集合，None表示不过滤
    """
    image_dir = Path(image_dir)
    label_dir = Path(label_dir)
    output_dir = Path(output_dir)
    
    # 获取所有图像文件
    image_files = sorted(image_dir.glob(f"*{file_extension}"))
    
    if sample_interval > 1:
        image_files = image_files[::sample_interval]
    
    if max_samples is not None:
        image_files = image_files[:max_samples]
    
    print(f"找到 {len(image_files)} 个图像文件（采样间隔={sample_interval}）")
    
    # 类别名称
    names = DEFAULT_CLASS_NAMES.copy()

    if target_class_ids is not None:
        print(f"类别过滤: {format_target_classes(target_class_ids, names=names)}")
    
    success_count = 0
    error_count = 0
    
    for i, image_path in enumerate(image_files):
        print(f"\n[{i+1}/{len(image_files)}] 处理: {image_path.name}")
        
        # 查找对应的标签文件
        label_path = resolve_label_path(label_dir, image_path.stem, preferred_extension=label_extension)
        if label_path is None:
            expected_extensions = [label_extension] if label_extension else [".txt", ".json"]
            print(f"  警告：未找到标签文件 {image_path.stem}，已检查扩展名: {expected_extensions}")
            error_count += 1
            continue
        
        # 查找对应的标定文件
        if calib_dir is not None:
            # 兼容 L2_calib 子目录和直接放在 calib/ 下两种结构
            calib_path = Path(calib_dir) / "L2_calib" / "camera4.json"
            if not calib_path.exists():
                calib_path = Path(calib_dir) / "camera4.json"
        else:
            # 自动推断：依次尝试两种常见目录结构
            base = image_path.parent.parent
            calib_path = base / "calib" / "L2_calib" / "camera4.json"
            if not calib_path.exists():
                calib_path = base / "calib" / "camera4.json"

        if not calib_path.exists():
            calib_path = None
            print(f"  警告：未找到标定文件，将仅进行2D可视化")
        
        try:
            # 可视化单帧
            sample_output_dir = output_dir
            visualize_single_frame(
                image_path=str(image_path),
                label_path=str(label_path),
                output_dir=str(sample_output_dir),
                calib_path=str(calib_path) if calib_path else None,
                roi_size=roi_size,
                virtual_fx=virtual_fx,
                ori_img_size=ori_img_size,
                names=names,
                save_types='combined',  # 只保存合成图
                target_class_ids=target_class_ids,
            )
            success_count += 1
        except Exception as e:
            print(f"  错误：{e}")
            error_count += 1
    
    print(f"\n处理完成！")
    print(f"  成功: {success_count}")
    print(f"  失败: {error_count}")
    print(f"  输出目录: {output_dir}")


def parse_args():
    parser = argparse.ArgumentParser(description="批量可视化真值标签")
    parser.add_argument("--image-dir", type=str, required=True, help="图像目录")
    parser.add_argument("--label-dir", type=str, required=True, help="标签目录")
    parser.add_argument("--output", type=str, default="./gt_visualization_batch", help="输出目录")
    parser.add_argument("--calib-dir", type=str, default=None,
                       help="标定文件目录（包含L2_calib/camera4.json）")
    parser.add_argument("--roi", type=int, nargs=2, default=None, metavar=("WIDTH", "HEIGHT"),
                       help="ROI尺寸（宽 高），如: --roi 704 352")
    parser.add_argument("--virtual-fx", type=float, default=None,
                       help="虚拟焦距，用于深度归一化")
    parser.add_argument("--ori-img-size", type=int, nargs=2, default=None, metavar=("WIDTH", "HEIGHT"),
                       help="原始图像尺寸（宽 高）")
    parser.add_argument("--max-samples", type=int, default=None,
                       help="最大处理样本数")
    parser.add_argument("--sample-interval", type=int, default=1,
                       help="采样间隔，每隔 N 帧处理一帧（默认1，即处理全部）")
    parser.add_argument("--file-extension", type=str, default=".jpg",
                       help="图像文件扩展名")
    parser.add_argument("--label-extension", type=str, default=None,
                       help="标签文件扩展名，默认自动在 .txt/.json 中查找")
    parser.add_argument("--classes", type=str, nargs="+", default=None,
                       help="仅可视化指定类别，支持类别名/类别ID，可传多个值或逗号分隔")
    
    return parser.parse_args()


def main():
    args = parse_args()
    target_class_ids = resolve_target_class_ids(args.classes, names=DEFAULT_CLASS_NAMES)
    
    visualize_batch(
        image_dir=args.image_dir,
        label_dir=args.label_dir,
        output_dir=args.output,
        calib_dir=args.calib_dir,
        roi_size=tuple(args.roi) if args.roi else None,
        virtual_fx=args.virtual_fx,
        ori_img_size=tuple(args.ori_img_size) if args.ori_img_size else None,
        max_samples=args.max_samples,
        sample_interval=args.sample_interval,
        file_extension=args.file_extension,
        label_extension=args.label_extension,
        target_class_ids=target_class_ids,
    )


if __name__ == "__main__":
    main()
