"""
单帧图像真值可视化脚本

用途：读取一帧图像和对应的真值标签文件，实现2D和3D的可视化

使用方法：
    python scripts_for_gt/visualize_single_frame.py \
        --image path/to/image.jpg \
        --label path/to/label.txt \
        --output output_dir \
        --roi 704 352 \
        --virtual-fx 500
        
示例：
    python scripts_for_gt/visualize_single_frame.py \
        --image /path/to/images/frame_001.jpg \
        --label /path/to/labels/frame_001.txt \
        --output ./gt_visualization
"""

import argparse
import json
import os
import sys
from pathlib import Path

import cv2
import numpy as np
import torch

# Add the visualization directory to path so plots_3d can be imported directly
FILE = Path(__file__).resolve()
VIS_DIR = FILE.parent
if str(VIS_DIR) not in sys.path:
    sys.path.insert(0, str(VIS_DIR))

from plots_3d import (
    decode_and_reconstruct_3d_box_from_target,
    plot_3d_boxes_from_decoded_targets,
    draw_bev_blank,
    drawbev,
    Annotator,
    colors,
)

JSON_TYPE_NAME_TO_CLASS = {
    # Vehicles — face_3d_classes [0-8]
    "vehicle": 0,            # generic vehicle fallback -> car
    "car": 0,
    "suv": 1,
    "pickup": 2,
    "medium_car": 3,
    "van": 4,
    "bus": 5,
    "truck": 6,
    "tanker": 6,
    "large_truck": 6,
    "construction_vehicle": 6,
    "special_vehicle": 7,
    "unknown": 8,
    # Pedestrian / cyclists — complete_3d_classes [9-12]
    "pedestrian": 9,
    "bicyclist": 10,
    "motorcyclist": 10,
    "bicycle": 11,
    "motorcycle": 11,
    "rider": 10,             # generic rider fallback -> bicyclist
    "tricycle": 12,
    "tricyclist": 12,
    # Other objects
    "traffic_sign": 13,
    "wheel": 14,
    "plate": 15,
    "face": 16,
}

DEFAULT_CLASS_NAMES = {
    0: "car", 1: "suv", 2: "pickup", 3: "medium_car", 4: "van",
    5: "bus", 6: "truck", 7: "special_vehicle", 8: "unknown",
    9: "pedestrian", 10: "bicyclist", 11: "bicycle", 12: "tricycle",
    13: "traffic_sign", 14: "wheel", 15: "plate", 16: "face",
}

CLASS_FILTER_GROUPS = {
    "vehicle": set(range(0, 9)),
    "vehicles": set(range(0, 9)),
    "vru": {9, 10, 11, 12},
    "cyclist": {10, 11},
    "cyclists": {10, 11},
    "two_wheeler": {10, 11},
    "two_wheelers": {10, 11},
}


def _safe_float(value):
    try:
        return float(value)
    except (TypeError, ValueError):
        return np.nan


def _map_json_label_class(obj):
    type_name = str(obj.get("type_name", "")).strip().lower()
    if type_name in JSON_TYPE_NAME_TO_CLASS:
        return JSON_TYPE_NAME_TO_CLASS[type_name]

    type_value = obj.get("type", None)
    try:
        type_id = int(type_value)
    except (TypeError, ValueError):
        return None

    # Legacy coarse ids written directly into JSON (old 5-class scheme).
    # 0=vehicle->car(0), 1=pedestrian->9, 2=bicycle->11, 3=rider->10, 13=tricycle->12
    legacy_type_map = {0: 0, 1: 9, 2: 11, 3: 10, 13: 12}
    if type_name in legacy_type_map:
        return JSON_TYPE_NAME_TO_CLASS.get(type_name, legacy_type_map.get(type_id))
    if type_id in legacy_type_map:
        return legacy_type_map[type_id]

    # Fine-grained ids from convert_json_to_json.py / convert_txt_to_json.py.
    # 0=car,1=suv,2=pickup,3=medium_car,4=van,5=bus,6=truck,7=special_vehicle,8=unknown
    if 0 <= type_id <= 8:
        return type_id
    if type_id == 9:   return 9   # pedestrian
    if type_id == 10:  return 10  # bicyclist / motorcyclist
    if type_id == 11:  return 11  # bicycle / motorcycle
    if type_id == 12:  return 12  # tricycle
    return None


def _xyxy_to_xywhn(box2d, img_width, img_height):
    x1, y1, x2, y2 = [_safe_float(v) for v in box2d[:4]]
    return np.array([
        (x1 + x2) * 0.5 / img_width,
        (y1 + y2) * 0.5 / img_height,
        (x2 - x1) / img_width,
        (y2 - y1) / img_height,
    ], dtype=np.float32)


def _normalize_uv(x_px, y_px, img_width, img_height):
    return x_px / img_width, y_px / img_height


def _iter_json_objects(payload):
    if isinstance(payload, dict):
        if all(isinstance(value, dict) for value in payload.values()):
            return payload.values()
        if isinstance(payload.get("objects"), list):
            return payload["objects"]
        return []
    if isinstance(payload, list):
        return payload
    return []


def _load_json_label_file(label_path, image_size=None):
    """读取 labels_json 文件并转换为可视化脚本使用的47维标签格式。"""
    if image_size is None:
        raise ValueError("JSON 标签需要提供 image_size=(width, height) 以便将 box2d 像素坐标转换为归一化坐标")

    img_width, img_height = image_size
    if img_width <= 0 or img_height <= 0:
        raise ValueError(f"无效的 image_size: {image_size}")

    with open(label_path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    labels = []
    for obj in _iter_json_objects(payload):
        cls = _map_json_label_class(obj)
        if cls is None:
            continue

        box2d = obj.get("box2d") or []
        if len(box2d) < 4:
            continue

        temp = np.full(47, np.nan, dtype=np.float32)
        temp[0] = cls
        temp[1:5] = _xyxy_to_xywhn(box2d, img_width, img_height)

        ori_3d = obj.get("3d_ori") or []
        ori_values = [_safe_float(v) for v in ori_3d]
        if len(ori_values) >= 12 and ori_values[0] != -1.0:
            temp[5:12] = np.array(ori_values[:7], dtype=np.float32)
            temp[12], temp[13] = _normalize_uv(ori_values[7], ori_values[8], img_width, img_height)
            temp[14] = ori_values[11]

        for face_name, offset in (("3d_front", 15), ("3d_back", 23), ("3d_left", 31), ("3d_right", 39)):
            face_values = [_safe_float(v) for v in (obj.get(face_name) or [])]
            if len(face_values) < 8:
                continue
            temp[offset:offset + 4] = np.array(face_values[:4], dtype=np.float32)
            if face_values[0] == -1.0:
                temp[offset + 4] = -1.0
                temp[offset + 5] = -1.0
            else:
                temp[offset + 4], temp[offset + 5] = _normalize_uv(face_values[4], face_values[5], img_width, img_height)
            temp[offset + 6] = face_values[6]
            temp[offset + 7] = face_values[7]

        labels.append(temp)

    if len(labels) > 0:
        return np.stack(labels, axis=0)
    return np.empty((0, 47), dtype=np.float32)


def load_label_file(label_path, image_size=None):
    """读取标签文件，返回47维的标签数组。
    
    标签格式参考：
    - txt: 6/18/50列格式
    - json: labels_json 目录下的逐帧 JSON 标签

    Returns:
        np.ndarray: (N, 47) 标签数组
    """
    label_path = Path(label_path)
    if label_path.suffix.lower() == ".json":
        return _load_json_label_file(label_path, image_size=image_size)

    labels = []
    with open(label_path, 'r', encoding='utf-8') as f:
        lines = f.read().strip().splitlines()
        for line in lines:
            line = line.strip().split(' ')
            line = line[:-1]  # 移除最后一列（通常是换行符）
            
            if len(line) == 6:  # 仅2D bbox
                temp = np.full(47, np.nan, dtype=np.float32)
                temp[:6] = np.array([float(x) for x in line[:6]], np.float32)
                labels.append(temp)
            elif len(line) == 18 and int(float(line[0])) in [1, 2, 3]:  # 行人、自行车、骑手
                temp = np.full(47, np.nan, dtype=np.float32)
                useful_gt_info = [float(line[x]) for x in (list(range(14)) + [16])]
                temp[0:15] = np.array(useful_gt_info, np.float32)
                labels.append(temp)
            elif len(line) == 50 and int(float(line[0])) in [0, 13]:  # 车辆、三轮车（带面信息）
                used_gt_ = [float(line[x]) for x in (list(range(14)) + [16])]
                temp = used_gt_
                temp.extend([float(x) for x in line[18:]])
                labels.append(np.array(temp, np.float32))
            else:
                print(f"警告：跳过无效标签格式，长度={len(line)}")
    
    if len(labels) > 0:
        labels = np.stack(labels, axis=0)
        assert labels.shape[1] == 47, f"标签应为47列，实际为{labels.shape[1]}列"
    else:
        labels = np.empty((0, 47), dtype=np.float32)
    
    return labels


def resolve_target_class_ids(target_classes, names=None):
    """解析目标类别，支持类别名、类别ID、逗号分隔和常见类别组。"""
    if not target_classes:
        return None

    if isinstance(target_classes, str):
        target_classes = [target_classes]

    names = DEFAULT_CLASS_NAMES if names is None else names
    name_to_id = {str(name).strip().lower(): cls for cls, name in names.items()}
    alias_to_ids = {name: {cls} for name, cls in JSON_TYPE_NAME_TO_CLASS.items()}
    alias_to_ids.update(CLASS_FILTER_GROUPS)

    resolved_ids = set()
    unknown_classes = []

    for raw_value in target_classes:
        for class_token in str(raw_value).split(","):
            class_name = class_token.strip().lower()
            if not class_name:
                continue

            if class_name in name_to_id:
                resolved_ids.add(name_to_id[class_name])
                continue

            if class_name in alias_to_ids:
                resolved_ids.update(alias_to_ids[class_name])
                continue

            try:
                resolved_ids.add(int(class_name))
            except ValueError:
                unknown_classes.append(class_name)

    if unknown_classes:
        available_classes = ", ".join(str(name) for name in names.values())
        raise ValueError(
            f"未知类别: {', '.join(sorted(set(unknown_classes)))}；"
            f"可用类别: {available_classes}"
        )

    return resolved_ids or None


def format_target_classes(target_class_ids, names=None):
    """将类别ID集合格式化为便于打印的名称列表。"""
    if not target_class_ids:
        return "all"

    names = DEFAULT_CLASS_NAMES if names is None else names
    return ", ".join(names.get(cls_id, str(cls_id)) for cls_id in sorted(target_class_ids))


def filter_labels_by_class(labels, target_class_ids):
    """按类别过滤标签。"""
    if target_class_ids is None or len(labels) == 0:
        return labels

    valid_mask = ~np.isnan(labels[:, 0])
    class_mask = np.zeros(len(labels), dtype=bool)
    class_mask[valid_mask] = np.isin(labels[valid_mask, 0].astype(np.int32), sorted(target_class_ids))
    return labels[class_mask]


def load_calibration(calib_path):
    """读取相机标定文件
    
    Args:
        calib_path: camera4.json 路径
        
    Returns:
        dict: 标定参数字典
    """
    with open(calib_path, 'r') as f:
        calib_params = json.load(f)
    return calib_params


def apply_roi_transform(img, labels, calib_params, roi_size, virtual_fx, ori_img_size):
    """应用ROI变换：裁剪图像和调整标签
    
    Args:
        img: 原始图像 (H, W, 3) BGR
        labels: 标签数组 (N, 47)
        calib_params: 标定参数字典
        roi_size: ROI 尺寸 (width, height)
        virtual_fx: 虚拟焦距
        ori_img_size: 原始图像尺寸 (width, height)
        
    Returns:
        img_roi: 裁剪后的图像
        labels_roi: 调整后的标签
        calib_roi: ROI 标定参数
    """
    oriW, oriH = ori_img_size
    roi_w, roi_h = roi_size
    
    # 获取相机参数
    fx = calib_params['focal_u']
    fy = calib_params['focal_v']
    cx = calib_params['cu']
    cy = calib_params['cv']
    c_pitch = calib_params['pitch']
    
    # 计算消失点和裁剪中心
    vanish_y = cy - fy * np.tan(c_pitch * np.pi / 180)
    crop_center_x = oriW // 2
    crop_center_y = vanish_y
    
    # 计算ROI边界
    roi_x1 = int(crop_center_x - roi_w / 2.0)
    roi_y1 = int(crop_center_y - roi_h / 2.0)
    roi_x2 = roi_x1 + roi_w
    roi_y2 = roi_y1 + roi_h
    
    # 裁剪图像
    img_roi = img[roi_y1:roi_y2, roi_x1:roi_x2, :]
    
    # 调整标签
    labels_roi = post_process_labels_to_roi(
        labels.copy(), oriW, oriH, roi_x1, roi_y1, roi_x2, roi_y2
    )
    
    # 计算缩放比例（如果需要resize到特定大小）
    # 这里假设不需要额外resize，ROI大小即为目标大小
    scale_x = 1.0
    scale_y = 1.0
    
    # 计算ROI标定参数
    calib_roi = {
        'fx': fx * scale_x,
        'fy': fy * scale_y,
        'cx': (cx - roi_x1) * scale_x,
        'cy': (cy - roi_y1) * scale_y,
        'distort_coeffs': calib_params.get('distort_coeffs', []),
        'target_fx': virtual_fx,
        'depth_scale': fx / virtual_fx,  # 用于将归一化深度转换回原始尺度
    }
    
    # 调整z3d深度值（深度归一化）
    if len(labels_roi) > 0:
        scale = virtual_fx / calib_roi['fx']
        labels_roi = scale_z3d(labels_roi, scale)
    
    return img_roi, labels_roi, calib_roi


def post_process_labels_to_roi(labels, oriW, oriH, roi_x1, roi_y1, roi_x2, roi_y2):
    """将标签从原始图像坐标转换到ROI坐标"""
    if len(labels) == 0:
        return labels
    
    roi_width = roi_x2 - roi_x1
    roi_height = roi_y2 - roi_y1
    
    # 将归一化的xywh转换为绝对坐标
    boxes = xywhn2xyxy(labels[:, 1:5], oriW, oriH, padw=0, padh=0)
    
    # 转换到ROI相对坐标
    new_x1 = boxes[:, 0] - roi_x1
    new_y1 = boxes[:, 1] - roi_y1
    new_x2 = boxes[:, 2] - roi_x1
    new_y2 = boxes[:, 3] - roi_y1
    
    # 检查是否完全在ROI外
    curr_outside = ((new_x1 < 0) & (new_x2 < 0)) | \
                   ((new_x1 >= roi_width) & (new_x2 >= roi_width)) | \
                   ((new_y1 < 0) & (new_y2 < 0)) | \
                   ((new_y1 >= roi_height) & (new_y2 >= roi_height))
    
    # 裁剪到ROI边界
    new_x1 = np.clip(new_x1, 0, roi_width - 1)
    new_y1 = np.clip(new_y1, 0, roi_height - 1)
    new_x2 = np.clip(new_x2, 0, roi_width - 1)
    new_y2 = np.clip(new_y2, 0, roi_height - 1)
    
    # 更新归一化坐标
    labels[:, 1] = (new_x1 + new_x2) * 0.5 / roi_width
    labels[:, 2] = (new_y1 + new_y2) * 0.5 / roi_height
    labels[:, 3] = (new_x2 - new_x1) / roi_width
    labels[:, 4] = (new_y2 - new_y1) / roi_height
    
    # 调整UV坐标（中心点投影）
    # 列索引: (12,13), (19,20), (27,28), (35,36), (43,44)
    for xi, yi in [(12, 13), (19, 20), (27, 28), (35, 36), (43, 44)]:
        mask = ~np.isnan(labels[:, xi]) & ~np.isnan(labels[:, yi])
        if mask.any():
            labels[mask, xi] = (labels[mask, xi] * oriW - roi_x1) / roi_width
            labels[mask, yi] = (labels[mask, yi] * oriH - roi_y1) / roi_height
    
    # 移除完全在ROI外的框
    labels = labels[~curr_outside]
    
    return labels


def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
    """将归一化的xywh转换为xyxy"""
    y = np.copy(x)
    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # x1
    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # y1
    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # x2
    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # y2
    return y


def scale_z3d(labels, scale):
    """缩放z3d深度坐标
    
    只缩放Z坐标，因为：
    - 投影公式：u = fx * X / Z + cx
    - 同时缩放fx和Z可以保持投影不变
    - X和Y不需要缩放，因为它们是从UV和Z恢复的
    """
    if len(labels) == 0 or scale == 1.0:
        return labels
    
    labels_scaled = labels.copy()
    
    # 整体中心z3d（列7）
    labels_scaled[:, 7] *= scale
    
    # 前面中心z3d（列17）
    mask = ~np.isnan(labels_scaled[:, 17])
    labels_scaled[mask, 17] *= scale
    
    # 后面中心z3d（列25）
    mask = ~np.isnan(labels_scaled[:, 25])
    labels_scaled[mask, 25] *= scale
    
    # 左面中心z3d（列33）
    mask = ~np.isnan(labels_scaled[:, 33])
    labels_scaled[mask, 33] *= scale
    
    # 右面中心z3d（列41）
    mask = ~np.isnan(labels_scaled[:, 41])
    labels_scaled[mask, 41] *= scale
    
    return labels_scaled


def plot_2d_boxes(img, labels, names=None, label_text=None, scale_factor=2):
    """绘制2D边界框
    
    Args:
        img: 图像 (H, W, 3) BGR
        labels: 标签数组 (N, 47)
        names: 类别名称字典
        label_text: 标签文字（如"2D GT"）
        scale_factor: 放大倍数
        
    Returns:
        img_annotated: 带标注的图像 BGR
    """
    h, w = img.shape[:2]
    
    # 放大图像以获得更清晰的可视化
    h_new = h * scale_factor
    w_new = w * scale_factor
    img_scaled = cv2.resize(img, (w_new, h_new), interpolation=cv2.INTER_LINEAR)
    
    # 转换为RGB用于Annotator
    img_rgb = cv2.cvtColor(img_scaled, cv2.COLOR_BGR2RGB)
    
    # 创建标注器
    fs = int((h_new + w_new) * 0.01)
    annotator = Annotator(img_rgb, line_width=max(2, round(fs / 8)), font_size=fs, pil=True)
    
    # 绘制边界
    annotator.rectangle([0, 0, w_new, h_new], None, (255, 255, 255), width=2)
    
    # 绘制每个框
    for label in labels:
        if np.isnan(label[0]):
            continue
        
        cls = int(label[0])
        x_center, y_center, box_w, box_h = label[1:5]
        
        # 转换为像素坐标
        x1 = (x_center - box_w / 2) * w_new
        y1 = (y_center - box_h / 2) * h_new
        x2 = (x_center + box_w / 2) * w_new
        y2 = (y_center + box_h / 2) * h_new
        
        box = [x1, y1, x2, y2]
        color = colors(cls)
        
        # 类别名称
        if names:
            cls_name = names.get(cls, str(cls))
        else:
            cls_name = str(cls)
        
        annotator.box_label(box, cls_name, color=color)
    
    # 转换回BGR并添加标签文字
    result = np.array(annotator.im)
    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
    
    if label_text:
        cv2.putText(result, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 
                   1.5, (0, 255, 255), 3, cv2.LINE_AA)
    
    return result


def visualize_single_frame(
    image_path,
    label_path,
    output_dir,
    calib_path=None,
    roi_size=None,
    virtual_fx=None,
    ori_img_size=None,
    names=None,
    save_types=None,
    target_class_ids=None,
):
    """可视化单帧图像的2D和3D真值
    
    Args:
        image_path: 图像路径
        label_path: 标签路径
        output_dir: 输出目录
        calib_path: 相机标定文件路径
        roi_size: ROI尺寸 (width, height)，如果为None则不应用ROI
        virtual_fx: 虚拟焦距
        ori_img_size: 原始图像尺寸 (width, height)
        names: 类别名称字典
        save_types: 保存类型列表 ['2d', '3d', 'bev', 'combined']，None表示全部保存
        target_class_ids: 目标类别ID集合，None表示不过滤
    """
    # 创建输出目录
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 读取图像
    img = cv2.imread(str(image_path))
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {image_path}")
    print(f"已读取图像: {image_path}, 尺寸: {img.shape[:2]}")
    
    # 默认类别名称
    if names is None:
        names = DEFAULT_CLASS_NAMES.copy()

    # 读取标签
    labels = load_label_file(label_path, image_size=(img.shape[1], img.shape[0]))
    print(f"已读取 {len(labels)} 个标签")

    if target_class_ids is not None:
        raw_label_count = len(labels)
        labels = filter_labels_by_class(labels, target_class_ids)
        print(
            f"类别过滤后保留 {len(labels)}/{raw_label_count} 个标签: "
            f"{format_target_classes(target_class_ids, names=names)}"
        )

    # 如果没有指定ori_img_size，使用图像实际尺寸
    if ori_img_size is None:
        ori_img_size = (img.shape[1], img.shape[0])
    
    # 检查是否需要应用ROI变换
    apply_roi = roi_size is not None and calib_path is not None and virtual_fx is not None
    
    if apply_roi:
        # 读取标定文件
        calib_params = load_calibration(calib_path)
        print(f"已读取标定文件: {calib_path}")
        
        # 应用ROI变换
        img_processed, labels_processed, calib = apply_roi_transform(
            img, labels, calib_params, roi_size, virtual_fx, ori_img_size
        )
        print(f"应用ROI变换，ROI尺寸: {roi_size}")
    else:
        # 不应用ROI，使用原始图像和标签
        img_processed = img
        labels_processed = labels
        
        # 如果提供了calib_path，读取标定参数
        if calib_path is not None:
            calib_params = load_calibration(calib_path)
            calib = {
                'fx': calib_params['focal_u'],
                'fy': calib_params['focal_v'],
                'cx': calib_params['cu'],
                'cy': calib_params['cv'],
                'distort_coeffs': calib_params.get('distort_coeffs', []),
                'depth_scale': 1.0,  # 不缩放深度
            }
        else:
            calib = None
            print("警告：未提供标定文件，无法进行3D可视化")
    
    # 默认保存所有类型
    if save_types is None:
        save_types = ['2d', '3d', 'bev', 'combined']
    
    output_filename = Path(image_path).stem
    
    # ========== 2D可视化 ==========
    img_2d = None
    if '2d' in save_types or 'combined' in save_types:
        print("\n生成2D可视化...")
        img_2d = plot_2d_boxes(img_processed, labels_processed, names=names, label_text="2D GT", scale_factor=1)
        
        # 保存2D可视化
        if '2d' in save_types:
            save_path_2d = output_dir / f"{output_filename}_2d_gt.jpg"
            cv2.imwrite(str(save_path_2d), img_2d)
            print(f"已保存2D可视化: {save_path_2d}")
    
    # ========== 3D可视化 ==========
    img_3d_bgr = None
    bev_image = None
    decoded_results = []
    
    need_3d = '3d' in save_types or 'combined' in save_types
    need_bev = 'bev' in save_types or 'combined' in save_types
    
    if calib is not None and len(labels_processed) > 0 and (need_3d or need_bev):
        print("\n生成3D可视化...")
        
        # 准备图像张量（BGR格式，归一化到[0,1]）
        img_tensor = torch.from_numpy(img_processed.transpose(2, 0, 1)).float() / 255.0
        img_tensor = img_tensor.unsqueeze(0)  # (1, 3, H, W)
        
        h_img, w_img = img_processed.shape[:2]
        
        # 解码3D框
        decoded_results = []
        for label in labels_processed:
            # 在decode函数中需要添加batch_idx列（列0）
            target = np.concatenate([[0], label])  # (48,)
            
            decoded = decode_and_reconstruct_3d_box_from_target(
                target, calib, w_img, h_img
            )
            if decoded is not None:
                decoded_results.append(decoded)
        
        print(f"成功解码 {len(decoded_results)} 个3D框")
        
        if len(decoded_results) > 0:
            # 3D投影可视化
            if need_3d:
                img_3d = plot_3d_boxes_from_decoded_targets(
                    img_tensor,
                    [decoded_results],
                    [str(image_path)],
                    calib=[calib],
                    names=names,
                    label_text="3D GT",
                    scale_factor=2
                )
                
                if img_3d is not None:
                    # 转换为BGR
                    img_3d_bgr = cv2.cvtColor(img_3d, cv2.COLOR_RGB2BGR)
                    
                    # 保存3D可视化
                    if '3d' in save_types:
                        save_path_3d = output_dir / f"{output_filename}_3d_gt.jpg"
                        cv2.imwrite(str(save_path_3d), img_3d_bgr)
                        print(f"已保存3D可视化: {save_path_3d}")
            
            # BEV可视化
            if need_bev:
                print("\n生成BEV可视化...")
                bev_image = draw_bev_blank()
                for decoded in decoded_results:
                    if decoded and decoded.get('object_3d') is not None:
                        bev_image = drawbev(bev_image, decoded['object_3d'], is_pred=False)
                
                # 保存BEV可视化
                if 'bev' in save_types:
                    save_path_bev = output_dir / f"{output_filename}_bev_gt.jpg"
                    cv2.imwrite(str(save_path_bev), bev_image)
                    print(f"已保存BEV可视化: {save_path_bev}")
            
            # 生成组合图像（2D + 3D 上下组合，然后与 BEV 左右组合）
            if 'combined' in save_types and img_2d is not None and img_3d_bgr is not None and bev_image is not None:
                print("\n生成组合可视化...")
                
                # 步骤1: 确保2D和3D宽度一致
                h_2d, w_2d = img_2d.shape[:2]
                h_3d, w_3d = img_3d_bgr.shape[:2]
                
                # 以2D的宽度为准，调整3D图像
                if w_3d != w_2d:
                    h_3d_resized = int(h_3d * w_2d / w_3d)
                    img_3d_bgr_resized = cv2.resize(img_3d_bgr, (w_2d, h_3d_resized))
                else:
                    img_3d_bgr_resized = img_3d_bgr
                    h_3d_resized = h_3d
                
                # 步骤2: 添加垂直间隔并上下拼接2D和3D
                padding_height = 20
                padding_vertical = np.zeros((padding_height, w_2d, 3), dtype=np.uint8)
                left_side = np.vstack([img_2d, padding_vertical, img_3d_bgr_resized])
                
                # 步骤3: 调整BEV高度以匹配左侧总高度
                h_left, w_left = left_side.shape[:2]
                h_bev, w_bev = bev_image.shape[:2]
                
                # 保持BEV的宽高比，调整高度匹配左侧
                w_bev_resized = int(w_bev * h_left / h_bev)
                bev_image_resized = cv2.resize(bev_image, (w_bev_resized, h_left))
                
                # 步骤4: 添加水平间隔并左右拼接
                padding_width = 20
                padding_horizontal = np.zeros((h_left, padding_width, 3), dtype=np.uint8)
                combined = np.hstack([left_side, padding_horizontal, bev_image_resized])
                
                save_path_combined = output_dir / f"{output_filename}_combined.jpg"
                cv2.imwrite(str(save_path_combined), combined)
                print(f"已保存组合可视化: {save_path_combined}")
        else:
            print("警告：没有成功解码的3D框")
    else:
        if calib is None:
            print("跳过3D可视化：未提供标定文件")
        else:
            print("跳过3D可视化：没有标签")
    
    print(f"\n可视化完成！输出目录: {output_dir}")


def parse_args():
    parser = argparse.ArgumentParser(description="单帧图像真值可视化脚本")
    parser.add_argument("--image", type=str, required=True, help="图像路径")
    parser.add_argument("--label", type=str, required=True, help="标签文件路径")
    parser.add_argument("--output", type=str, default="./gt_visualization", help="输出目录")
    parser.add_argument("--calib", type=str, default=None, 
                       help="相机标定文件路径（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("--classes", type=str, nargs="+", default=None,
                       help="仅可视化指定类别，支持类别名/类别ID，可传多个值或逗号分隔")
    
    args = parser.parse_args()
    
    # 如果没有提供calib路径，尝试从图像路径推断
    if args.calib is None:
        image_path = Path(args.image)
        base = image_path.parent.parent  # dataset root
        # 兼容 calib/L2_calib/camera4.json 和 calib/camera4.json 两种目录结构
        for calib_path in [
            base / 'calib' / 'L2_calib' / 'camera4.json',
            base / 'calib' / 'camera4.json',
        ]:
            if calib_path.exists():
                args.calib = str(calib_path)
                print(f"自动找到标定文件: {args.calib}")
                break
        else:
            print(f"警告：未找到标定文件（已检查 calib/L2_calib/camera4.json 和 calib/camera4.json），将仅进行2D可视化")
    
    return args


def main():
    args = parse_args()
    
    names = DEFAULT_CLASS_NAMES.copy()
    target_class_ids = resolve_target_class_ids(args.classes, names=names)
    
    visualize_single_frame(
        image_path=args.image,
        label_path=args.label,
        output_dir=args.output,
        calib_path=args.calib,
        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,
        names=names,
        target_class_ids=target_class_ids,
    )


if __name__ == "__main__":
    main()
