# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""Plotting utils."""

import contextlib
import math
import os
from copy import copy
from pathlib import Path

import cv2
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sn
import torch
from PIL import Image, ImageDraw
from scipy.ndimage.filters import gaussian_filter1d
from ultralytics.utils.plotting import Annotator

from utils import TryExcept, threaded
from utils.general import LOGGER, clip_boxes, increment_path, xywh2xyxy, xyxy2xywh, denormalize_3d_predictions
from utils.metrics import fitness

# Settings
RANK = int(os.getenv("RANK", -1))
matplotlib.rc("font", **{"size": 11})
matplotlib.use("Agg")  # for writing to files only

# Global color map for 2D bounding boxes (RGB format)
CLASS_COLOR_MAP = {
    0: (255, 100, 100),   # vehicle - light red
    1: (100, 255, 100),   # pedestrian - light green
    2: (100, 100, 255),   # bicycle - light blue
    3: (255, 255, 100),   # rider - yellow
    4: (255, 100, 255),   # motorcycle - magenta
    5: (100, 255, 255),   # tricycle - cyan
    6: (255, 150, 50),    # bus - orange
    7: (150, 50, 255),    # truck - purple
    8: (50, 255, 150),    # van - teal
    9: (255, 200, 100),   # car - gold
    10: (200, 100, 255),  # barrier - lavender
    11: (100, 200, 255),  # cone - sky blue
    12: (255, 100, 200),  # sign - pink
    13: (200, 255, 100),  # tricycle - lime
}

def get_class_color(cls_id):
    """Get consistent color for a given class ID.

    Args:
        cls_id: Class ID (integer)

    Returns:
        RGB color tuple
    """
    return CLASS_COLOR_MAP.get(cls_id, (200, 200, 200))  # Default gray for unknown classes


class Colors:
    """Provides an RGB color palette derived from Ultralytics color scheme for visualization tasks."""

    def __init__(self):
        """Initializes the Colors class with a palette derived from Ultralytics color scheme, converting hex codes to
        RGB.

        Colors derived from `hex = matplotlib.colors.TABLEAU_COLORS.values()`.
        """
        hexs = (
            "FF3838",
            "FF9D97",
            "FF701F",
            "FFB21D",
            "CFD231",
            "48F90A",
            "92CC17",
            "3DDB86",
            "1A9334",
            "00D4BB",
            "2C99A8",
            "00C2FF",
            "344593",
            "6473FF",
            "0018EC",
            "8438FF",
            "520085",
            "CB38FF",
            "FF95C8",
            "FF37C7",
        )
        self.palette = [self.hex2rgb(f"#{c}") for c in hexs]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        """Returns color from palette by index `i`, in BGR format if `bgr=True`, else RGB; `i` is an integer index."""
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):
        """Converts hexadecimal color `h` to an RGB tuple (PIL-compatible) with order (R, G, B)."""
        return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4))


colors = Colors()  # create instance for 'from utils.plots import colors'


def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")):
    """
    Args:
        x: Features to be visualized
        module_type: Module type
        stage: Module stage within model
        n: Maximum number of feature maps to plot
        save_dir: Directory to save results.
    """
    if ("Detect" not in module_type) and (
        "Segment" not in module_type
    ):  # 'Detect' for Object Detect task,'Segment' for Segment task
        _batch, channels, height, width = x.shape  # batch, channels, height, width
        if height > 1 and width > 1:
            f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename

            blocks = torch.chunk(x[0].cpu(), channels, dim=0)  # select batch index 0, block by channels
            n = min(n, channels)  # number of plots
            _fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True)
            ax = ax.ravel()
            plt.subplots_adjust(wspace=0.05, hspace=0.05)
            for i in range(n):
                ax[i].imshow(blocks[i].squeeze())  # cmap='gray'
                ax[i].axis("off")

            LOGGER.info(f"Saving {f}... ({n}/{channels})")
            plt.savefig(f, dpi=300, bbox_inches="tight")
            plt.close()
            np.save(str(f.with_suffix(".npy")), x[0].cpu().numpy())  # npy save


def hist2d(x, y, n=100):
    """Generates a logarithmic 2D histogram, useful for visualizing label or evolution distributions.

    Used in used in labels.png and evolve.png.
    """
    xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
    hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
    xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
    yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
    return np.log(hist[xidx, yidx])


def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
    """Applies a low-pass Butterworth filter to `data` with specified `cutoff`, `fs`, and `order`."""
    from scipy.signal import butter, filtfilt

    # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
    def butter_lowpass(cutoff, fs, order):
        """Applies a low-pass Butterworth filter to a signal with specified cutoff frequency, sample rate, and filter
        order.
        """
        nyq = 0.5 * fs
        normal_cutoff = cutoff / nyq
        return butter(order, normal_cutoff, btype="low", analog=False)

    b, a = butter_lowpass(cutoff, fs, order=order)
    return filtfilt(b, a, data)  # forward-backward filter


def output_to_target(output, max_det=300):
    """Converts YOLOv5 model output to [batch_id, class_id, x, y, w, h, conf] format for plotting, limiting detections
    to `max_det`.
    """
    targets = []
    for i, o in enumerate(output):
        box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
        j = torch.full((conf.shape[0], 1), i)
        targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
    return torch.cat(targets, 0).numpy()


@threaded
def plot_images(images, targets, paths=None, fname="images.jpg", names=None, is_labels=True):
    """Plots an image grid with labels from YOLOv5 predictions or targets, saving to `fname`.

    Args:
        images: Batch of images
        targets: Target labels or predictions
        paths: Image paths
        fname: Output filename
        names: Class names
        is_labels: If True, treats targets as ground truth labels (no conf filtering).
                   If False, treats as predictions (applies conf threshold > 0.25).
    """
    if isinstance(images, torch.Tensor):
        images = images.cpu().float().numpy().copy()  # Make a copy to avoid modifying original
    if isinstance(targets, torch.Tensor):
        targets = targets.cpu().numpy()

    max_size = 1920  # max image size
    max_subplots = 16  # max image subplots, i.e. 4x4
    bs, _, h, w = images.shape  # batch size, _, height, width
    bs = min(bs, max_subplots)  # limit plot images
    ns = np.ceil(bs**0.5)  # number of subplots (square)
    if np.max(images[0]) <= 1:
        images = images * 255  # de-normalise (create new array, don't modify in place)

    # Build Image
    mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)  # init
    for i, im in enumerate(images):
        if i == max_subplots:  # if last batch has fewer images than we expect
            break
        x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
        im = im.transpose(1, 2, 0)
        mosaic[y : y + h, x : x + w, :] = im

    # Resize (optional)
    scale = max_size / ns / max(h, w)
    if scale < 1:
        h = math.ceil(scale * h)
        w = math.ceil(scale * w)
        # Use INTER_AREA for downscaling to avoid artifacts
        mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)), interpolation=cv2.INTER_AREA)

    # Annotate
    fs = int((h + w) * ns * 0.01)  # font size
    # Convert BGR to RGB for Annotator which expects RGB input (pil=True)
    mosaic = mosaic[..., ::-1]  # BGR to RGB
    annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
    for i in range(bs):
        x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
        annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2)  # borders
        if paths:
            annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220))  # filenames
        if len(targets) > 0:
            ti = targets[targets[:, 0] == i]  # image targets
            boxes = xywh2xyxy(ti[:, 2:6]).T
            classes = ti[:, 1].astype("int")
            conf = None if is_labels else ti[:, 6]  # confidence only for predictions

            if boxes.shape[1]:
                if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
                    boxes[[0, 2]] *= w  # scale to pixels
                    boxes[[1, 3]] *= h
                elif scale < 1:  # absolute coords need scale if image scales
                    boxes *= scale
            boxes[[0, 2]] += x
            boxes[[1, 3]] += y
            for j, box in enumerate(boxes.T.tolist()):
                cls = classes[j]
                color = colors(cls)
                from utils.general import get_class_name
                cls = get_class_name(names, cls) if names else cls
                if is_labels or conf[j] > 0.25:  # 0.25 conf thresh for predictions
                    label = f"{cls}" if is_labels else f"{cls} {conf[j]:.1f}"
                    annotator.box_label(box, label, color=color)
    annotator.im.save(fname)  # save


def rotation_3d_in_axis(points, angles, axis=0):
    """Rotate points around a specified axis.

    Args:
        points: (N, 3) array of 3D points
        angles: rotation angle in radians
        axis: 0=X, 1=Y, 2=Z

    Returns:
        rotated points (N, 3)
    """
    rot_sin = np.sin(angles)
    rot_cos = np.cos(angles)
    ones = np.ones_like(rot_cos)
    zeros = np.zeros_like(rot_cos)
    if axis == 1:  # Y axis (right-handed system: X=right, Y=down, Z=forward)
        # Standard right-handed Y-axis rotation matrix (transposed for post-multiplication)
        # R_y(θ)^T = [cos(θ)  0  -sin(θ)]
        #            [  0     1     0   ]
        #            [sin(θ)  0   cos(θ)]
        # Using points @ rot_mat (post-multiplication), so we use the transposed form
        rot_mat = np.stack([
            np.stack([rot_cos, zeros, -rot_sin]),
            np.stack([zeros, ones, zeros]),
            np.stack([rot_sin, zeros, rot_cos])
        ])
    elif axis == 2:  # Z axis
        rot_mat = np.stack([
            np.stack([rot_cos, rot_sin, zeros]),
            np.stack([-rot_sin, rot_cos, zeros]),
            np.stack([zeros, zeros, ones])
        ])
    elif axis == 0:  # X axis
        rot_mat = np.stack([
            np.stack([ones, zeros, zeros]),
            np.stack([zeros, rot_cos, rot_sin]),
            np.stack([zeros, -rot_sin, rot_cos])
        ])
    else:
        raise ValueError(f'axis should in range [0, 1, 2], got {axis}')
    return np.dot(points, rot_mat)


def compute_3d_box_corners_4face(center_3d, dimensions, rotation, face_type=0):
    """Compute 8 corners of a 3D bounding box from a face center point.

    The center_3d is the center of a specific face, not the box center.

    Args:
        center_3d: (x, y, z) center of the face in camera coordinates
        dimensions: (length, height, width) of the box
        rotation: rot_y (rotation around y-axis in radians)
        face_type: 0=front, 1=tail, 2=left, 3=right

    Returns:
        corners: (8, 3) array of corner coordinates
            Corner ordering: [0,1,3,2,4,5,7,6] pattern from unravel_index
    """
    l, h, w = dimensions

    # Create 8 corners using unravel_index pattern (same as reference)
    corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1).astype(np.float64)
    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]

    # Offset based on face type (face center becomes origin)
    # face_type: 0=front, 1=tail, 2=left, 3=right
    # IMPORTANT: In object local frame, X=0 is REAR, X=length is FRONT
    # After rotation, this aligns correctly with camera frame
    if face_type == 1:  # tail (back/rear face) - at X=0 (REAR in object frame)
        corners_norm = corners_norm - [0, 0.5, 0.5]
    elif face_type == 0:  # front face - at X=length (FRONT in object frame)
        corners_norm = corners_norm - [1, 0.5, 0.5]
    elif face_type == 3:  # right face
        corners_norm = corners_norm - [0.5, 0.5, 0]
    elif face_type == 2:  # left face
        corners_norm = corners_norm - [0.5, 0.5, 1]
    else:  # whole object center
        corners_norm = corners_norm - [0.5, 0.5, 0.5]

    # Scale by dimensions [l, h, w]
    dims = np.array([l, h, w])
    corners = dims.reshape([1, 3]) * corners_norm.reshape([8, 3])

    # Rotate around Y axis
    corners = rotation_3d_in_axis(corners, rotation, axis=1)

    # Translate to face center position
    corners += np.array(center_3d).reshape(1, 3)

    return corners


def read_calib_from_path(img_path):
    """Read calibration parameters from the dataset directory structure.

    Args:
        img_path: Path to the image file

    Returns:
        dict with calibration parameters or None if not found
    """
    import json
    base_path = os.path.dirname(img_path)
    calib_path = base_path.replace('images', 'calib') + '/L2_calib/camera4.json'

    if not os.path.exists(calib_path):
        return None

    try:
        with open(calib_path, 'r') as f:
            calib_params = json.load(f)
            return {
                'fx': calib_params['focal_u'],
                'fy': calib_params['focal_v'],
                'cx': calib_params['cu'],
                'cy': calib_params['cv'],
                'yaw': calib_params.get('yaw', 0),
                'pitch': calib_params.get('pitch', 0),
            }
    except Exception:
        return None


def apply_fisheye_distortion(x, y, distort_coeffs):
    """Apply fisheye distortion to normalized camera coordinates.

    Args:
        x, y: normalized camera coordinates (x/z, y/z)
        distort_coeffs: [k1, k2, k3, k4] fisheye distortion coefficients

    Returns:
        xd, yd: distorted normalized coordinates
    """
    if not distort_coeffs or len(distort_coeffs) < 4:
        return x, y

    k1, k2, k3, k4 = distort_coeffs[:4]
    r = np.sqrt(x * x + y * y)

    # Avoid division by zero
    if r < 1e-8:
        return x, y

    theta = np.arctan(r)

    # KB (Kannala-Brandt) fisheye distortion model:
    # theta_d = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8)
    # For KB model: r_d = theta_d (NOT tan(theta_d))
    theta_d = theta * (1 + k1 * theta**2 + k2 * theta**4 + k3 * theta**6 + k4 * theta**8)

    scale = theta_d / r
    return x * scale, y * scale


def remove_fisheye_distortion(xd, yd, distort_coeffs, max_iter=20):
    """Remove fisheye distortion from normalized camera coordinates.

    Args:
        xd, yd: distorted normalized camera coordinates
        distort_coeffs: [k1, k2, k3, k4] fisheye distortion coefficients
        max_iter: maximum iterations for Newton-Raphson solver

    Returns:
        xn, yn: undistorted normalized coordinates
    """
    if not distort_coeffs or len(distort_coeffs) < 4:
        return xd, yd

    k1, k2, k3, k4 = distort_coeffs[:4]

    # Compute distorted radius
    r_d = np.sqrt(xd * xd + yd * yd)

    # Handle center point
    if r_d < 1e-8:
        return xd, yd

    # For KB model: r_d = theta_d, so theta_d = r_d directly
    theta_d = r_d

    # Solve for theta using Newton-Raphson iteration
    # We need to solve: theta_d = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8)
    # Better initial guess: use theta_d / (1 + k1*theta_d^2) as starting point
    theta_d2 = theta_d * theta_d
    theta = theta_d / (1 + k1*theta_d2)  # Better initial guess

    for _ in range(max_iter):
        # f(theta) = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8) - theta_d
        theta2 = theta * theta
        theta4 = theta2 * theta2
        theta6 = theta4 * theta2
        theta8 = theta4 * theta4

        f = theta * (1 + k1*theta2 + k2*theta4 + k3*theta6 + k4*theta8) - theta_d

        # f'(theta) = 1 + 3*k1*theta^2 + 5*k2*theta^4 + 7*k3*theta^6 + 9*k4*theta^8
        f_prime = 1 + 3*k1*theta2 + 5*k2*theta4 + 7*k3*theta6 + 9*k4*theta8

        # Newton-Raphson update
        theta_new = theta - f / f_prime

        # Check convergence
        if abs(theta_new - theta) < 1e-8:
            theta = theta_new
            break

        theta = theta_new

    # For KB model: r = tan(theta)
    r = np.tan(theta)

    # Scale back to get undistorted coordinates
    scale = r / r_d
    return xd * scale, yd * scale


def project_3d_to_2d_with_distortion(points_3d, calib):
    """Project 3D points to 2D using pre-adjusted camera calibration with fisheye distortion.

    Args:
        points_3d: (N, 3) array of 3D points in camera coordinates
        calib: dict with fx, fy, cx, cy, distort_coeffs (already adjusted for ROI-resized image)

    Returns:
        points_2d: (N, 2) array of 2D points in image coordinates
    """
    fx, fy = calib['fx'], calib['fy']
    cx, cy = calib['cx'], calib['cy']
    distort_coeffs = calib.get('distort_coeffs', [])

    points_2d = []
    for pt in points_3d:
        x, y, z = pt
        if z > 0.1:  # Only project points in front of camera
            # Normalize coordinates
            xn = x / z
            yn = y / z

            # Apply fisheye distortion
            xd, yd = apply_fisheye_distortion(xn, yn, distort_coeffs)

            # Project to image coordinates (calib already adjusted for ROI-resized image)
            u = fx * xd + cx
            v = fy * yd + cy
            points_2d.append([u, v])
        else:
            points_2d.append([np.nan, np.nan])

    return np.array(points_2d)


def sample_3d_edge(p1, p2, num_samples=10):
    """Sample points along a 3D edge between two points.

    Args:
        p1, p2: 3D points (x, y, z)
        num_samples: number of samples along the edge

    Returns:
        (num_samples, 3) array of sampled 3D points
    """
    t = np.linspace(0, 1, num_samples).reshape(-1, 1)
    return p1 + t * (p2 - p1)


def project_3d_box_edges_with_distortion(corners_3d, calib, samples_per_edge=10):
    """Project 3D box edges to 2D with distortion by sampling along edges.

    Args:
        corners_3d: (8, 3) array of 3D corner coordinates
        calib: dict with fx, fy, cx, cy, distort_coeffs (already adjusted for ROI-resized image)
        samples_per_edge: number of samples per edge for curved projection

    Returns:
        dict with edge_name -> (N, 2) array of 2D points for each edge
    """
    # Define edges: (start_idx, end_idx)
    # Back face edges (indices 4-7)
    # Front face edges (indices 0-3)
    # Connecting edges
    edges = {
        'back_0': (4, 5), 'back_1': (5, 6), 'back_2': (6, 7), 'back_3': (7, 4),
        'connect_0': (0, 4), 'connect_1': (1, 5), 'connect_2': (2, 6), 'connect_3': (3, 7),
        'front_0': (0, 1), 'front_1': (1, 2), 'front_2': (2, 3), 'front_3': (3, 0),
        'front_x1': (0, 2), 'front_x2': (1, 3),  # X marks on front face
    }

    edge_points_2d = {}
    for edge_name, (i, j) in edges.items():
        # Sample 3D points along the edge
        p1, p2 = corners_3d[i], corners_3d[j]
        sampled_3d = sample_3d_edge(p1, p2, samples_per_edge)

        # Project with distortion
        sampled_2d = project_3d_to_2d_with_distortion(sampled_3d, calib)
        edge_points_2d[edge_name] = sampled_2d

    return edge_points_2d


def plot_box3d_on_img_with_distortion(img, edge_points_2d, color_front=(255, 0, 0), color_back=(0, 0, 255), color_side=(0, 255, 255), thickness=1):
    """Draw 3D bounding box on image using edge points (handles distortion).

    Args:
        img: image to draw on
        edge_points_2d: dict with edge_name -> (N, 2) array of 2D points
        color_front: color for front face edges (default: red)
        color_back: color for back face edges (default: blue)
        color_side: color for connecting side edges (default: cyan)
        thickness: line thickness

    Returns:
        img with 3D box drawn
    """
    front_edges = ['front_0', 'front_1', 'front_2', 'front_3', 'front_x1', 'front_x2']
    back_edges = ['back_0', 'back_1', 'back_2', 'back_3', 'back_x1', 'back_x2']

    for edge_name, points in edge_points_2d.items():
        # Skip if any points are invalid
        if np.any(np.isnan(points)):
            continue

        pts = points.astype(np.int32)

        # Determine color based on edge type
        if edge_name in front_edges:
            color = color_front
            line_thickness = thickness
        elif edge_name in back_edges:
            color = color_back
            line_thickness = thickness
        else:
            # Connecting edges (side edges)
            color = color_side
            line_thickness = thickness

        # Draw polyline through sampled points
        cv2.polylines(img, [pts], isClosed=False, color=color, thickness=line_thickness, lineType=cv2.LINE_AA)

    return img


def project_3d_to_2d_with_calib(points_3d, calib):
    """Project 3D points to 2D using pre-adjusted camera calibration.

    Args:
        points_3d: (N, 3) array of 3D points in camera coordinates
        calib: dict with fx, fy, cx, cy (already adjusted for ROI-resized image)

    Returns:
        points_2d: (N, 2) array of 2D points in image coordinates
    """
    fx, fy = calib['fx'], calib['fy']
    cx, cy = calib['cx'], calib['cy']

    points_2d = []
    for pt in points_3d:
        x, y, z = pt
        if z > 0.1:  # Only project points in front of camera
            u = fx * x / z + cx
            v = fy * y / z + cy
            points_2d.append([u, v])
        else:
            points_2d.append([np.nan, np.nan])

    return np.array(points_2d)


def project_3d_to_2d_simple(points_3d, img_size):
    """Project 3D points to 2D using simple pinhole camera model.

    Args:
        points_3d: (N, 3) array of 3D points in camera coordinates
        img_size: (width, height) of image

    Returns:
        points_2d: (N, 2) array of 2D points in pixel coordinates
    """
    w, h = img_size

    # Estimate focal length from image size (typical for automotive cameras)
    # Using a reasonable approximation based on common camera setups
    fx = fy = w * 1.2
    cx, cy = w / 2, h / 2

    points_2d = []
    for pt in points_3d:
        x, y, z = pt
        if z > 0.1:  # Only project points in front of camera
            u = fx * x / z + cx
            v = fy * y / z + cy
            points_2d.append([u, v])
        else:
            points_2d.append([np.nan, np.nan])

    return np.array(points_2d)


def plot_box3d_on_img(img, corners_2d, color_front=(255, 0, 0), color_back=(0, 0, 255), color_side=(0, 255, 255), thickness=1):
    """Draw 3D bounding box on image.

    Args:
        img: image to draw on
        corners_2d: (8, 2) array of 2D corner coordinates
            Corner order matches compute_3d_box_corners_4face output
        color_front: color for front face edges (default: red)
        color_back: color for back face edges (default: blue)
        color_side: color for connecting side edges (default: cyan)
        thickness: line thickness

    Returns:
        img with 3D box drawn
    """
    # Line indices for drawing the box
    # Back face (indices 4-7), connecting edges, front face (indices 0-3)
    line_indices = (
        (4, 5), (5, 6), (6, 7), (7, 4),  # back face
        (0, 4), (1, 5), (2, 6), (3, 7),  # connecting edges
        (0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3)  # front face with X
    )

    front_edges = [(0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3)]
    back_edges = [(4, 5), (5, 6), (6, 7), (7, 4)]

    corners = corners_2d.astype(np.int32)

    for start, end in line_indices:
        try:
            pt1 = (corners[start, 0], corners[start, 1])
            pt2 = (corners[end, 0], corners[end, 1])

            # Determine color based on edge type
            if (start, end) in front_edges:
                cv2.line(img, pt1, pt2, color_front, thickness, cv2.LINE_AA)
            elif (start, end) in back_edges:
                cv2.line(img, pt1, pt2, color_back, thickness, cv2.LINE_AA)
            else:
                # Connecting side edges
                cv2.line(img, pt1, pt2, color_side, thickness, cv2.LINE_AA)
        except Exception:
            pass

    return img


def convert_predictions_to_target_format(pred_3d, calib, image_size):
    """Convert 41-channel post-processed predictions to 47-dimension target format.

    This utility function transforms post-processed model predictions (with UV in absolute pixels
    and rot_y decoded) to match the target label format, enabling consistent plotting.

    Args:
        pred_3d: Post-processed 3D predictions array with shape (N, 41) containing:
            - 0-5: Front face (z, u_abs, v_abs, l, h, w)
            - 6-11: Rear face (z, u_abs, v_abs, l, h, w)
            - 12-17: Left face (z, u_abs, v_abs, l, h, w)
            - 18-23: Right face (z, u_abs, v_abs, l, h, w)
            - 24-40: Whole object (z, u_abs, v_abs, l, h, w, yaw_cls×4, rot_y, ..., cut_cls×3)
            Note: UV coordinates should already be in absolute pixels (after post-processing)
            Note: rot_y should already be decoded (channel 34 contains rot_y, not yaw_reg)
        calib: Calibration dict with 'fx', 'fy', 'cx', 'cy' keys
        image_size: Tuple of (width, height) in pixels for UV normalization

    Returns:
        Array with shape (N, 47) matching target label format (without batch index):
            - 0: class (filled with 0, should be set externally from 2D predictions)
            - 1-4: [x, y, w, h] 2D bbox (filled with 0, should be set externally)
            - 5-7: [x3d, y3d, z3d] whole object 3D center
            - 8-10: [l, h, w] whole object size
            - 11: rot_y
            - 12-13: [xc, yc] whole object UV (normalized [0,1])
            - 14: alpha (observation angle)
            - 15-22: front face [x3d, y3d, z3d, alpha, xc, yc, score, visible]
            - 23-30: rear face
            - 31-38: left face
            - 39-46: right face

    Note: This function expects post-processed predictions (UV in absolute pixels, rot_y decoded).
    It is used for plotting inference results where predictions have already been post-processed.
    """
    if isinstance(pred_3d, torch.Tensor):
        pred_3d = pred_3d.cpu().numpy()

    n = pred_3d.shape[0]
    out = np.zeros((n, 47), dtype=pred_3d.dtype)

    # Get calibration parameters
    fx = calib.get('fx', image_size[0] / 2)
    fy = calib.get('fy', image_size[1] / 2)
    cx = calib.get('cx', image_size[0] / 2)
    cy = calib.get('cy', image_size[1] / 2)

    # Get image dimensions for UV normalization
    img_w, img_h = image_size

    # Helper function: inverse projection (u, v, z) -> (x3d, y3d, z3d)
    def inverse_project(u, v, z):
        x3d = (u - cx) * z / fx
        y3d = (v - cy) * z / fy
        z3d = z
        return x3d, y3d, z3d

    # Helper function: compute alpha from rot_y and position
    def compute_alpha(rot_y, x3d, z3d):
        # alpha = rot_y - atan2(x3d, z3d)
        return rot_y - np.arctan2(x3d, z3d)

    # --- Whole object (channels 24-40 -> dims 5-14) ---
    whole_z = pred_3d[:, 24]
    whole_u = pred_3d[:, 25]  # Already in absolute pixels
    whole_v = pred_3d[:, 26]  # Already in absolute pixels
    whole_l = pred_3d[:, 27]
    whole_h = pred_3d[:, 28]
    whole_w = pred_3d[:, 29]
    rot_y = pred_3d[:, 34]  # Already decoded rot_y

    # Inverse projection for whole object
    whole_x3d, whole_y3d, whole_z3d = inverse_project(whole_u, whole_v, whole_z)
    whole_alpha = compute_alpha(rot_y, whole_x3d, whole_z3d)

    out[:, 5] = whole_x3d
    out[:, 6] = whole_y3d
    out[:, 7] = whole_z3d
    out[:, 8] = whole_l
    out[:, 9] = whole_h
    out[:, 10] = whole_w
    out[:, 11] = rot_y
    out[:, 12] = whole_u / img_w  # Normalize to [0, 1]
    out[:, 13] = whole_v / img_h  # Normalize to [0, 1]
    out[:, 14] = whole_alpha

    # --- Front face (channels 0-5 -> dims 15-22) ---
    front_z = pred_3d[:, 0]
    front_u = pred_3d[:, 1]  # Already in absolute pixels
    front_v = pred_3d[:, 2]  # Already in absolute pixels

    front_x3d, front_y3d, front_z3d = inverse_project(front_u, front_v, front_z)
    front_alpha = compute_alpha(rot_y, front_x3d, front_z3d)

    out[:, 15] = front_x3d
    out[:, 16] = front_y3d
    out[:, 17] = front_z3d
    out[:, 18] = front_alpha
    out[:, 19] = front_u / img_w  # Normalize to [0, 1]
    out[:, 20] = front_v / img_h  # Normalize to [0, 1]
    out[:, 21] = 1.0  # score (placeholder)
    out[:, 22] = 1.0  # visible (placeholder)

    # --- Rear face (channels 6-11 -> dims 23-30) ---
    rear_z = pred_3d[:, 6]
    rear_u = pred_3d[:, 7]  # Already in absolute pixels
    rear_v = pred_3d[:, 8]  # Already in absolute pixels

    rear_x3d, rear_y3d, rear_z3d = inverse_project(rear_u, rear_v, rear_z)
    rear_alpha = compute_alpha(rot_y, rear_x3d, rear_z3d)

    out[:, 23] = rear_x3d
    out[:, 24] = rear_y3d
    out[:, 25] = rear_z3d
    out[:, 26] = rear_alpha
    out[:, 27] = rear_u / img_w  # Normalize to [0, 1]
    out[:, 28] = rear_v / img_h  # Normalize to [0, 1]
    out[:, 29] = 1.0  # score
    out[:, 30] = 1.0  # visible

    # --- Left face (channels 12-17 -> dims 31-38) ---
    left_z = pred_3d[:, 12]
    left_u = pred_3d[:, 13]  # Already in absolute pixels
    left_v = pred_3d[:, 14]  # Already in absolute pixels

    left_x3d, left_y3d, left_z3d = inverse_project(left_u, left_v, left_z)
    left_alpha = compute_alpha(rot_y, left_x3d, left_z3d)

    out[:, 31] = left_x3d
    out[:, 32] = left_y3d
    out[:, 33] = left_z3d
    out[:, 34] = left_alpha
    out[:, 35] = left_u / img_w  # Normalize to [0, 1]
    out[:, 36] = left_v / img_h  # Normalize to [0, 1]
    out[:, 37] = 1.0  # score
    out[:, 38] = 1.0  # visible

    # --- Right face (channels 18-23 -> dims 39-46) ---
    right_z = pred_3d[:, 18]
    right_u = pred_3d[:, 19]  # Already in absolute pixels
    right_v = pred_3d[:, 20]  # Already in absolute pixels

    right_x3d, right_y3d, right_z3d = inverse_project(right_u, right_v, right_z)
    right_alpha = compute_alpha(rot_y, right_x3d, right_z3d)

    out[:, 39] = right_x3d
    out[:, 40] = right_y3d
    out[:, 41] = right_z3d
    out[:, 42] = right_alpha
    out[:, 43] = right_u / img_w  # Normalize to [0, 1]
    out[:, 44] = right_v / img_h  # Normalize to [0, 1]
    out[:, 45] = 1.0  # score
    out[:, 46] = 1.0  # visible

    return out


def _reconstruct_3d_box_from_face(face_uv, face_z, dims, rot_y, face_type, calib):
    """Reconstruct 3D box corners from visible face center.

    Args:
        face_uv: (u, v) 2D pixel coordinates of face center
        face_z: Depth of face center in meters
        dims: (l, h, w) object dimensions in meters
        rot_y: Rotation yaw angle in radians
        face_type: 0=front, 1=rear, 2=left, 3=right
        calib: Calibration dict with fx, fy, cx, cy, distort_coeffs

    Returns:
        corners_3d: (8, 3) array of 3D corner coordinates in camera frame
        or None if invalid
    """
    if calib is None:
        return None

    fx, fy = calib['fx'], calib['fy']
    cx, cy = calib['cx'], calib['cy']
    distort_coeffs = calib.get('distort_coeffs', [])

    u_face, v_face = face_uv

    # Back-project face center to 3D
    # Normalize pixel coordinates
    xd = (u_face - cx) / fx
    yd = (v_face - cy) / fy

    # Remove fisheye distortion if present
    if distort_coeffs and len(distort_coeffs) >= 4:
        xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs)
    else:
        xn, yn = xd, yd

    # Back-project to 3D
    x3d_face = xn * face_z
    y3d_face = yn * face_z
    z3d_face = face_z

    # Extract dimensions
    l, h, w = dims

    if np.isnan(l) or np.isnan(h) or np.isnan(w) or np.isnan(rot_y):
        return None

    # Use existing helper function that takes face center as input
    face_center_3d = np.array([x3d_face, y3d_face, z3d_face])
    corners_3d = compute_3d_box_corners_4face(
        face_center_3d, dims, rot_y, face_type=face_type
    )

    object_3d = [x3d_face, y3d_face, z3d_face, l, h, w, rot_y, face_type]

    # Return corners directly without reordering
    return corners_3d, object_3d


def decode_and_reconstruct_3d_box_from_target(target, calib, img_width, img_height,
                                                face_3d_classes=None, complete_3d_classes=None):
    """Decode and reconstruct 3D box from ground truth target format.

    Args:
        target: (48,) array in target format:
            - col 0: image index
            - col 1: class
            - cols 2-5: [x, y, w, h] normalized 2D bbox
            - cols 6-8: [x3d, y3d, z3d] 3D center in camera coords
            - cols 9-11: [length, height, width] 3D dimensions
            - col 12: rot_y (yaw rotation)
            - cols 13-14: [xc, yc] 2D center of 3D box projection (normalized)
            - col 15: alpha
            - cols 16-23: front face [x3d, y3d, z3d, alpha, xc, yc, score, visible]
            - cols 24-31: rear face
            - cols 32-39: left face
            - cols 40-47: right face
        calib: Calibration dict
        img_width: Image width in pixels
        img_height: Image height in pixels
        face_3d_classes: List/set of class IDs that use face-based reconstruction (default: [0, 1, 2, 3, 4, 5])
        complete_3d_classes: List/set of class IDs that use complete 3D box (default: [6, 7, 8, 9])

    Returns:
        dict or None: {
            'should_draw': bool,
            'cls': int,
            'corners_3d': np.ndarray (8, 3) or None,
            'face_center_2d': tuple or None,  # For face-based classes
            'face_color': tuple or None,  # For face-based classes
            'center_2d': tuple or None,  # For complete box classes
        }
    """
    # Default class lists: face-based reconstruction for vehicles (0) and tricycles (13),
    # complete 3D box for pedestrian (1), bicycle (2), rider (3)
    if face_3d_classes is None:
        face_3d_classes = [0, 13]
    if complete_3d_classes is None:
        complete_3d_classes = [1, 2, 3]

    if len(target) < 15 or np.isnan(target[1]):
        return None

    cls = int(target[1])

    # Skip unsupported classes (not in either list)
    if cls not in face_3d_classes and cls not in complete_3d_classes:
        return None

    # Get depth_scale for converting normalized depth back to original calibration space
    # In ROI mode with target_fx, xyz3d are scaled by target_fx/fx_orig during training
    # depth_scale = fx_orig/target_fx converts them back
    depth_scale = calib['depth_scale']  # Required field, no default

    result = {
        'should_draw': True,
        'cls': cls,
        'corners_3d': None,
        'face_center_2d': None,
        'face_color': None,
        'center_2d': None,
    }

    if cls in face_3d_classes:  # Classes with face annotations - use best visible face
        if len(target) < 48:
            return None
        
        # Find best face
        face_offsets = [16, 24, 32, 40]
        # Face colors in BGR format for OpenCV: front=red, rear=blue, left=green, right=yellow
        face_colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255)]
        
        best_face_type = -1
        best_score = -1
        best_face_data = None
        
        for face_type, face_offset in enumerate(face_offsets):
            if target.shape[0] < face_offset + 8:
                continue
            face_data = target[face_offset:face_offset + 8]
            x3d_face, y3d_face, z3d_face, alpha, xc_face, yc_face, score, is_visible = face_data
            
            # Skip invalid faces
            if is_visible == -1 or np.isnan(is_visible) or is_visible != 1:
                continue
            if np.isnan(score) or score < 0.3:
                continue
            if np.isnan(xc_face) or np.isnan(yc_face) or np.isnan(z3d_face) or z3d_face <= 0:
                continue
            
            # Track best face by score
            if score > best_score:
                best_score = score
                best_face_type = face_type
                best_face_data = face_data
        
        # Reconstruct from best face
        if best_face_type != -1 and best_face_data is not None:
            xc_face = best_face_data[4]
            yc_face = best_face_data[5]
            z3d_face = best_face_data[2]

            # Scale z back to original calibration space (from normalized depth)
            z3d_face = z3d_face * depth_scale

            # Convert normalized UV to absolute pixels
            u_face = xc_face * img_width
            v_face = yc_face * img_height
            
            # Get dimensions and rotation
            dims = target[9:12]  # [l, h, w]
            rot_y = target[12]
            
            corners_3d, object_3d = _reconstruct_3d_box_from_face(
                (u_face, v_face), z3d_face, dims, rot_y, best_face_type, calib
            )
            
            if corners_3d is not None:
                # No reordering here - will be done in visualization
                result['corners_3d'] = corners_3d
                result['face_center_2d'] = (u_face, v_face)
                result['face_color'] = face_colors[best_face_type]
                result['object_3d'] = object_3d

    elif cls in complete_3d_classes:  # Classes with complete 3D box but no face annotations
        # Extract 3D box parameters
        x3d, y3d, z3d = target[6:9]
        dimensions = target[9:12]  # [l, h, w]
        rot_y = target[12]
        xc_norm, yc_norm = target[13:15]

        # Scale only z3d back to original calibration space (from normalized depth)
        # X and Y don't need scaling - only Z is normalized for training
        z3d = z3d * depth_scale

        # Validate
        if np.isnan(z3d) or z3d <= 0 or np.any(np.isnan(dimensions)):
            return None
        if np.isnan(x3d) or np.isnan(y3d):
            return None

        # Compute 3D box corners
        center_3d = np.array([x3d, y3d, z3d])
        corners_3d = compute_3d_box_corners_4face(
            center_3d, dimensions, rot_y, face_type=-1
        )

        # NOTE: GT does NOT need corner reordering (unlike predictions)
        # This suggests GT rot_y convention is different from predicted rot_y
        result['corners_3d'] = corners_3d
        result['center_2d'] = (xc_norm * img_width, yc_norm * img_height)
        # Create object_3d for BEV visualization: [x, y, z, l, h, w, rot_y, face_type]
        result['object_3d'] = [x3d, y3d, z3d, dimensions[0], dimensions[1], dimensions[2], rot_y, -1]

    return result


def decode_and_reconstruct_3d_box(pred_2d_box, pred_3d_box, calib, conf_thres,
                                   face_3d_classes=None, complete_3d_classes=None):
    """Decode and reconstruct 3D box from compact 34-channel predictions.

    Args:
        pred_2d_box: (6,) array [x1, y1, x2, y2, conf, cls]
        pred_3d_box: (34,) array in compact format
        calib: Calibration dict
        conf_thres: Confidence threshold
        face_3d_classes: List/set of class IDs that use face-based reconstruction (default: [0, 1, 2, 3, 4, 5])
        complete_3d_classes: List/set of class IDs that use complete 3D box (default: [6, 7, 8, 9])

    Returns:
        dict or None: {
            'should_draw': bool,
            'cls': int,
            'corners_3d': np.ndarray (8, 3) or None,
            'face_center_2d': tuple or None,
            'face_color': tuple or None,
        }
    """
    # Default class lists: face-based reconstruction for vehicles (0) and tricycles (13),
    # complete 3D box for pedestrian (1), bicycle (2), rider (3)
    if face_3d_classes is None:
        face_3d_classes = [0, 13]
    if complete_3d_classes is None:
        complete_3d_classes = [1, 2, 3]

    conf = pred_2d_box[4]
    if conf < conf_thres:
        return None

    cls = int(pred_2d_box[5])

    # Skip unsupported classes (not in either list)
    if cls not in face_3d_classes and cls not in complete_3d_classes:
        return None

    # NOTE: depth_scale is NOT applied here for predictions because it's already
    # applied in the model forward pass (models/yolo.py:436)
    # Only targets need depth_scale applied (in decode_and_reconstruct_3d_box_from_target)

    result = {
        'should_draw': True,
        'cls': cls,
        'corners_3d': None,
        'face_center_2d': None,
        'face_color': None,
        'cut_cls': None,
    }

    if cls in face_3d_classes:  # Classes with face annotations - face-based reconstruction
        face_offsets = [0, 6, 12, 18]
        # Face colors in BGR format for OpenCV: front=red, rear=blue, left=green, right=yellow
        face_colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255)]
        
        # Parse cut_cls
        # NOTE: Decoding method MUST match training loss type:
        # - BCE loss training -> use sigmoid for independent class probabilities
        # - CrossEntropy loss training -> use softmax for normalized probabilities
        # Current default: softmax (matches CrossEntropy)
        cut_cls_logits = pred_3d_box[31:34]
        
        # Method 1: Softmax (for CrossEntropy trained models) - CURRENT DEFAULT
        # cut_cls_probs = np.exp(cut_cls_logits) / np.sum(np.exp(cut_cls_logits))
        
        # Method 2: Sigmoid (for BCE trained models) - UNCOMMENT if using BCE loss
        cut_cls_probs = 1 / (1 + np.exp(-cut_cls_logits))
        
        cut_cls_pred = np.argmax(cut_cls_probs)
        
        # Determine candidate faces
        if cut_cls_pred == 1:  # cut_in
            candidate_faces = [0]
        elif cut_cls_pred == 2:  # cut_out
            candidate_faces = [1]
        else:  # normal
            candidate_faces = [0, 1, 2, 3]
        
        # Select best face
        best_face_type = -1
        best_visible_score = -1.0
        for face_type in candidate_faces:
            face_offset = face_offsets[face_type]
            z_face = pred_3d_box[face_offset]
            u_face = pred_3d_box[face_offset + 1]
            v_face = pred_3d_box[face_offset + 2]
            visible_score = pred_3d_box[face_offset + 5]
            
            if np.isnan(z_face) or z_face <= 0 or np.isnan(u_face) or np.isnan(v_face):
                continue
            if np.isnan(visible_score):
                continue
            
            if visible_score > best_visible_score:
                best_visible_score = visible_score
                best_face_type = face_type
        
        if best_face_type == -1:
            return None
        
        # Reconstruct 3D box
        face_offset = face_offsets[best_face_type]
        z_face = pred_3d_box[face_offset]
        u_face = pred_3d_box[face_offset + 1]
        v_face = pred_3d_box[face_offset + 2]
        dims = pred_3d_box[27:30]  # [l, h, w]
        rot_y = pred_3d_box[30]

        # z is already scaled in model forward pass, use directly
        corners_3d, object_3d = _reconstruct_3d_box_from_face(
            (u_face, v_face), z_face, dims, rot_y, best_face_type, calib
        )
        
        if corners_3d is not None:
            # No reordering here - will be done in visualization
            result['corners_3d'] = corners_3d
            result['face_center_2d'] = (u_face, v_face)
            result['face_color'] = face_colors[best_face_type]
            result['object_3d'] = object_3d
            result['cut_cls'] = cut_cls_pred
        else:
            return None

    elif cls in complete_3d_classes:  # Classes with complete 3D box but no face annotations
        # Extract whole box info
        z_whole = pred_3d_box[24]
        u_whole = pred_3d_box[25]
        v_whole = pred_3d_box[26]
        l, h, w = pred_3d_box[27:30]
        rot_y = pred_3d_box[30]

        # z is already scaled in model forward pass, use directly

        if calib is None:
            return None

        fx, fy = calib['fx'], calib['fy']
        cx, cy = calib['cx'], calib['cy']
        distort_coeffs = calib.get('distort_coeffs', [])
        
        # Back-project to 3D
        xd = (u_whole - cx) / fx
        yd = (v_whole - cy) / fy
        
        if distort_coeffs and len(distort_coeffs) >= 4:
            xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs)
        else:
            xn, yn = xd, yd
        
        x_center = xn * z_whole
        y_center = yn * z_whole
        
        # Reconstruct 8 corners
        corners_3d = compute_3d_box_corners_4face(
            (x_center, y_center, z_whole), (l, h, w), rot_y, face_type=-1
        )

        result['corners_3d'] = corners_3d
        result['object_3d'] = [x_center, y_center, z_whole, l, h, w, rot_y, -1]
    
    return result


def draw_3d_box_from_corners(im, corners_3d, calib, img_shape,
                             face_center_2d=None, face_color=None, thickness=1):
    """Project and draw 3D box from corner coordinates.

    Args:
        im: Image array (H, W, 3) BGR
        corners_3d: (8, 3) 3D corner coordinates in camera frame
        calib: Calibration dict
        img_shape: (w, h)
        face_center_2d: (u, v) or None, face center marker for face-based
        face_color: tuple or None, face center marker color
        thickness: line thickness

    Returns:
        im: Modified image with 3D box drawn
    """
    w, h = img_shape

    # Reorder corners: compute_3d_box_corners_4face outputs indices 0-3 at rear, 4-7 at front
    # plot_box3d_on_img expects indices 0-3 = front (red), indices 4-7 = rear (blue)
    # So we swap: [4,5,6,7,0,1,2,3] -> indices 0-3 = front, indices 4-7 = rear
    corners_3d = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]]

    # Colors in BGR format for OpenCV:
    # plot_box3d_on_img draws indices 0-3 with color_front and 4-7 with color_back
    color_front = (0, 0, 255)   # Red in BGR (B=0, G=0, R=255) for front/head (indices 0-3)
    color_back = (255, 0, 0)    # Blue in BGR (B=255, G=0, R=0) for rear/back (indices 4-7)
    color_side = (255, 255, 0)  # Cyan in BGR (B=255, G=255, R=0) for side edges

    # Project to 2D
    if calib and calib.get('distort_coeffs'):
        edge_points_2d = project_3d_box_edges_with_distortion(
            corners_3d, calib, samples_per_edge=15
        )
        im = plot_box3d_on_img_with_distortion(
            im, edge_points_2d, color_front, color_back, color_side, thickness=thickness
        )
    else:
        if calib is not None:
            corners_2d = project_3d_to_2d_with_calib(corners_3d, calib)
        else:
            corners_2d = project_3d_to_2d_simple(corners_3d, (w, h))

        if not np.any(np.isnan(corners_2d)):
            im = plot_box3d_on_img(
                im, corners_2d, color_front, color_back, color_side, thickness=thickness
            )

    # Draw face center marker if provided
    if face_center_2d is not None and face_color is not None:
        u, v = face_center_2d
        cv2.circle(im, (int(u), int(v)), 2, face_color, -1, cv2.LINE_AA)

    return im


def plot_3d_boxes_from_decoded(im, decoded_results, paths, calib=None, names=None, label_text=None, scale_factor=2):
    """Plot 3D boxes from pre-decoded results (inference script layer).

    This function only handles projection and drawing, not decoding/reconstruction.
    The 3D decoding should be done at the inference script level before calling this.

    Args:
        im: Batch tensor (N, 3, H, W) normalized [0, 1] in BGR format (from dataloader3d)
        decoded_results: List of lists, where decoded_results[i] contains decoded results
                        for image i. Each decoded result is a dict from decode_and_reconstruct_3d_box:
                        {
                            'should_draw': bool,
                            'cls': int,
                            'corners_3d': np.ndarray (8, 3) or None,
                            'face_center_2d': tuple or None,
                            'face_color': tuple or None,
                        }
        paths: List of image paths
        calib: Calibration dict or list of calib dicts
        names: Class names dict
        label_text: Optional text to display on top-left (e.g., "3D GT" or "3D Pred")
        scale_factor: Scale factor for output image resolution (default 2 for clearer images)

    Returns:
        np.ndarray: (H*scale_factor, W*scale_factor, 3) RGB image with 3D boxes drawn, or None
    """
    import cv2

    if im.ndim == 3:
        im = im.unsqueeze(0)

    im_np = im[0].cpu().numpy().transpose(1, 2, 0)
    im_np = np.ascontiguousarray(im_np * 256, dtype=np.uint8)
    # Input is BGR from dataloader3d, keep as BGR for cv2 drawing
    h_orig, w_orig = im_np.shape[:2]

    # Upscale image for clearer visualization
    h_new = h_orig * scale_factor
    w_new = w_orig * scale_factor
    im_bgr = cv2.resize(im_np, (w_new, h_new), interpolation=cv2.INTER_LINEAR)

    # Handle calibration - scale calibration parameters for upscaled image
    if isinstance(calib, list):
        calib = calib[0] if len(calib) > 0 else None

    # Create scaled calibration for drawing
    if calib is not None:
        calib_scaled = {
            'fx': calib['fx'] * scale_factor,
            'fy': calib['fy'] * scale_factor,
            'cx': calib['cx'] * scale_factor,
            'cy': calib['cy'] * scale_factor,
            'distort_coeffs': calib.get('distort_coeffs', []),
        }
    else:
        calib_scaled = None

    # Draw each decoded 3D box
    for decoded in decoded_results[0]:  # First image in batch
        if not decoded or not decoded['should_draw']:
            continue

        if decoded['corners_3d'] is not None:
            corners_3d = decoded['corners_3d']

            # Scale face_center_2d if present
            face_center_2d = decoded.get('face_center_2d')
            if face_center_2d is not None:
                face_center_2d = (face_center_2d[0] * scale_factor, face_center_2d[1] * scale_factor)

            # No reordering here - draw_3d_box_from_corners handles it for all classes
            im_bgr = draw_3d_box_from_corners(
                im_bgr,
                corners_3d,
                calib_scaled,
                (w_new, h_new),
                face_center_2d=face_center_2d,
                face_color=decoded.get('face_color')
            )

    # Draw label text on top-left (after drawing boxes so it's on top)
    if label_text:
        cv2.putText(im_bgr, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, cv2.LINE_AA)

    # Convert back to RGB
    im_rgb = cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB)
    return im_rgb


def plot_3d_boxes_from_decoded_targets(im, decoded_results, paths, calib=None, names=None, label_text=None, scale_factor=2):
    """Plot 3D boxes from pre-decoded ground truth targets.

    This function only handles projection and drawing, not decoding/reconstruction.
    The 3D decoding should be done at the script level before calling this.

    Args:
        im: Batch tensor (N, 3, H, W) normalized [0, 1] in BGR format (from dataloader3d)
        decoded_results: List of lists, where decoded_results[i] contains decoded results
                        for image i. Each decoded result is a dict from decode_and_reconstruct_3d_box_from_target:
                        {
                            'should_draw': bool,
                            'cls': int,
                            'corners_3d': np.ndarray (8, 3) or None,
                            'face_center_2d': tuple or None,  # For vehicles/tricycles
                            'face_color': tuple or None,  # For vehicles/tricycles
                            'center_2d': tuple or None,  # For pedestrians/bicycles/riders
                        }
        paths: List of image paths
        calib: Calibration dict or list of calib dicts
        names: Class names dict
        label_text: Optional text to display on top-left (e.g., "3D GT" or "3D Pred")
        scale_factor: Scale factor for output image resolution (default 2 for clearer images)

    Returns:
        np.ndarray: (H*scale_factor, W*scale_factor, 3) RGB image with 3D boxes drawn, or None
    """
    import cv2

    if im.ndim == 3:
        im = im.unsqueeze(0)

    im_np = im[0].cpu().numpy().transpose(1, 2, 0)
    im_np = np.ascontiguousarray(im_np * 255, dtype=np.uint8)
    # Input is BGR from dataloader3d, keep as BGR for cv2 drawing
    h_orig, w_orig = im_np.shape[:2]

    # Upscale image for clearer visualization
    h_new = h_orig * scale_factor
    w_new = w_orig * scale_factor
    im_bgr = cv2.resize(im_np, (w_new, h_new), interpolation=cv2.INTER_LINEAR)

    # Handle calibration - scale calibration parameters for upscaled image
    if isinstance(calib, list):
        calib = calib[0] if len(calib) > 0 else None

    # Create scaled calibration for drawing
    if calib is not None:
        calib_scaled = {
            'fx': calib['fx'] * scale_factor,
            'fy': calib['fy'] * scale_factor,
            'cx': calib['cx'] * scale_factor,
            'cy': calib['cy'] * scale_factor,
            'distort_coeffs': calib.get('distort_coeffs', []),
        }
    else:
        calib_scaled = None

    # Draw each decoded 3D box
    for decoded in decoded_results[0]:  # First image in batch
        if not decoded or not decoded['should_draw']:
            continue

        if decoded['corners_3d'] is not None:
            cls = decoded['cls']

            if cls in [0, 13]:  # vehicles/tricycles - face-based
                # Scale face_center_2d if present
                face_center_2d = decoded.get('face_center_2d')
                if face_center_2d is not None:
                    face_center_2d = (face_center_2d[0] * scale_factor, face_center_2d[1] * scale_factor)

                # Draw 3D box from corners
                im_bgr = draw_3d_box_from_corners(
                    im_bgr,
                    decoded['corners_3d'],
                    calib_scaled,
                    (w_new, h_new),
                    face_center_2d=face_center_2d,
                    face_color=decoded.get('face_color')
                )

            elif cls in [1, 2, 3]:  # pedestrians/bicycles/riders - complete box
                # Use centralized draw function (handles reordering and colors)
                im_bgr = draw_3d_box_from_corners(
                    im_bgr,
                    decoded['corners_3d'],
                    calib_scaled,
                    (w_new, h_new),
                    thickness=2
                )

                # Draw center point (scaled)
                if decoded.get('center_2d'):
                    color = colors(cls)
                    u, v = decoded['center_2d']
                    u_scaled, v_scaled = int(u * scale_factor), int(v * scale_factor)
                    cv2.circle(im_bgr, (u_scaled, v_scaled), 4, color, -1, cv2.LINE_AA)
                    cv2.circle(im_bgr, (u_scaled, v_scaled), 6, (255, 255, 255), 2, cv2.LINE_AA)

    # Draw label text on top-left (after drawing boxes so it's on top)
    if label_text:
        cv2.putText(im_bgr, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, cv2.LINE_AA)

    # Convert back to RGB
    im_rgb = cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB)
    return im_rgb


def _plot_images_3d_core(images, targets, paths=None, fname="images.jpg", names=None,
                         calibs=None, return_array=False, face_3d_classes=None, complete_3d_classes=None):
    """Core implementation for plotting 2x2 grid: target 2D, target 3D, prediction 2D, prediction 3D.

    Layout for TensorBoard logging (single sample):
        (1,1) Target 2D    | (1,2) Target 3D
        (2,1) Prediction 2D | (2,2) Prediction 3D

    Args:
        images: Batch of images (B, C, H, W) in BGR format (from dataloader3d)
        targets: Target tensor with 48 columns per target
        paths: List of image paths
        fname: Output filename
        names: Class names dict
        calibs: Tuple of calibration dicts from dataloader (pre-adjusted for ROI-resized image)
        return_array: If True, return numpy array instead of saving to file
        predictions_2d: List of 2D predictions per image (optional, for combined view)
        predictions_3d: 3D predictions (optional, for combined view)
        conf_thres: Confidence threshold for predictions
        face_3d_classes: List/set of class IDs that use face-based reconstruction
        complete_3d_classes: List/set of class IDs that use complete 3D box
    """
    if isinstance(images, torch.Tensor):
        images = images.cpu().float().numpy()
    if isinstance(targets, torch.Tensor):
        targets = targets.cpu().numpy()

    # Only plot first sample for 2x2 grid
    _, _, h, w = images.shape

    if np.max(images[0]) <= 1:
        images = images * 255  # de-normalise

    # Convert BGR to RGB (dataloader3d returns BGR format)
    images = images[:, ::-1, :, :]  # Reverse channel order: BGR -> RGB

    # Get first image
    im = images[0].transpose(1, 2, 0).astype(np.uint8)
    im_target_2d = im.copy()
    im_target_3d = im.copy()

    # Get targets for first image
    ti = targets[targets[:, 0] == 0] if len(targets) > 0 else np.array([])

    # Get calibration for first image
    calib = None
    if calibs is not None and len(calibs) > 0:
        calib = calibs[0]
    elif paths:
        calib = read_calib_from_path(paths[0])

    # --- Draw Target 2D boxes ---
    for t in ti:
        if len(t) < 6 or np.isnan(t[1]):
            continue
        cls = int(t[1])
        color = get_class_color(cls)

        x_center, y_center, bw, bh = t[2:6]
        if np.any(np.isnan([x_center, y_center, bw, bh])):
            continue
        x1 = int((x_center - bw / 2) * w)
        y1 = int((y_center - bh / 2) * h)
        x2 = int((x_center + bw / 2) * w)
        y2 = int((y_center + bh / 2) * h)

        cv2.rectangle(im_target_2d, (x1, y1), (x2, y2), color, 1)

    # --- Draw Target 3D boxes ---
    # Resolve defaults here (same values used inside decode_and_reconstruct_3d_box_from_target)
    _face_3d_classes = face_3d_classes if face_3d_classes is not None else [0, 13]
    _complete_3d_classes = complete_3d_classes if complete_3d_classes is not None else [1, 2, 3]
    for t in ti:
        if len(t) < 15 or np.isnan(t[1]):
            continue

        cls = int(t[1])
        decoded = decode_and_reconstruct_3d_box_from_target(t, calib, w, h, _face_3d_classes, _complete_3d_classes)

        if decoded and decoded['should_draw'] and decoded['corners_3d'] is not None:
            if cls in _face_3d_classes:  # Face-based classes (vehicles/tricycles)
                im_target_3d = draw_3d_box_from_corners(
                    im_target_3d,
                    decoded['corners_3d'],
                    calib,
                    (w, h),
                    face_center_2d=decoded.get('face_center_2d'),
                    face_color=decoded.get('face_color')
                )
            elif cls in _complete_3d_classes:  # Complete box classes (pedestrians/bicycles/riders)
                corners_3d = decoded['corners_3d']
                color = colors(cls)

                if calib and calib.get('distort_coeffs'):
                    corners_3d_reordered = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]]
                    edge_points_2d = project_3d_box_edges_with_distortion(
                        corners_3d_reordered, calib, samples_per_edge=15
                    )
                    im_target_3d = plot_box3d_on_img_with_distortion(
                        im_target_3d, edge_points_2d, (255, 0, 0), (0, 0, 255), (0, 255, 255), thickness=1
                    )
                else:
                    if calib is not None:
                        corners_2d = project_3d_to_2d_with_calib(corners_3d, calib)
                    else:
                        corners_2d = project_3d_to_2d_simple(corners_3d, (w, h))

                    if not np.any(np.isnan(corners_2d)):
                        im_target_3d = plot_box3d_on_img(
                            im_target_3d, corners_2d, (255, 0, 0), (0, 0, 255), (0, 255, 255), thickness=1
                        )

                if decoded.get('center_2d'):
                    u, v = decoded['center_2d']
                    cv2.circle(im_target_3d, (int(u), int(v)), 3, color, -1, cv2.LINE_AA)

    # Create 2x2 figure: target 2D, target 3D (top row)
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # Add labels
    cv2.putText(im_target_2d, 'Target 2D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2)
    cv2.putText(im_target_3d, 'Target 3D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2)

    axes[0, 0].imshow(im_target_2d)
    axes[0, 0].set_title('Target 2D', fontsize=12)
    axes[0, 0].axis('off')

    axes[0, 1].imshow(im_target_3d)
    axes[0, 1].set_title('Target 3D', fontsize=12)
    axes[0, 1].axis('off')

    # Bottom row: placeholder for predictions (empty images with text)
    im_pred_2d = im.copy()
    im_pred_3d = im.copy()
    cv2.putText(im_pred_2d, 'Prediction 2D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2)
    cv2.putText(im_pred_3d, 'Prediction 3D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2)

    axes[1, 0].imshow(im_pred_2d)
    axes[1, 0].set_title('Prediction 2D', fontsize=12)
    axes[1, 0].axis('off')

    axes[1, 1].imshow(im_pred_3d)
    axes[1, 1].set_title('Prediction 3D', fontsize=12)
    axes[1, 1].axis('off')

    if paths:
        fig.suptitle(Path(paths[0]).name[:50], fontsize=10)

    plt.tight_layout()

    if return_array:
        dpi = 150
        fig.set_dpi(dpi)
        fig.canvas.draw()

        buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
        fig_w, fig_h = fig.canvas.get_width_height()
        img_array = buf.reshape(fig_h, fig_w, 4)
        img_array = img_array[:, :, :3]
        plt.close(fig)
        return img_array
    else:
        # Use fig.savefig (not plt.savefig) to avoid matplotlib global-state race condition.
        # plt.savefig() saves the global "current figure", which can be overwritten by another
        # thread calling plt.subplots() concurrently (ni=0,1,2 all trigger threads in quick
        # succession). fig.savefig() operates on the local figure object and is thread-safe.
        fig.savefig(fname, dpi=150, bbox_inches='tight')
        plt.close(fig)


@threaded
def _plot_images_3d_threaded(images, targets, paths=None, fname="images.jpg", names=None,
                              calibs=None, face_3d_classes=None, complete_3d_classes=None):
    """Threaded wrapper for plot_images_3d when saving to file."""
    _plot_images_3d_core(images, targets, paths, fname, names, calibs, return_array=False,
                         face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)


def plot_images_3d(images, targets, paths=None, fname="images.jpg", names=None,
                   calibs=None, return_array=False, face_3d_classes=None, complete_3d_classes=None):
    """Plots 2D and 3D labels in separate subplots for visualization.

    When return_array=False (default): runs in a separate thread like plot_images, returns Thread.
    When return_array=True: runs synchronously and returns numpy array for TensorBoard.

    Args:
        face_3d_classes: List/set of class IDs that use face-based reconstruction
        complete_3d_classes: List/set of class IDs that use complete 3D box
    """
    if return_array:
        # Synchronous call - return numpy array directly
        return _plot_images_3d_core(images, targets, paths, fname, names, calibs, return_array=True,
                                    face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)
    else:
        # Threaded call - same behavior as plot_images
        return _plot_images_3d_threaded(images, targets, paths, fname, names, calibs,
                                        face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)


def plot_predictions_3d(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg",
                        names=None, calibs=None, conf_thres=0.25, return_array=False,
                        face_3d_classes=None, complete_3d_classes=None):
    """Plots 3D predictions on images.

    Args:
        images: Batch of images (B, C, H, W)
        predictions_2d: 2D detection predictions (after NMS) - list of tensors per image, each (N, 6) [x1, y1, x2, y2, conf, cls]
        predictions_3d: 3D predictions tensor with shape (B, N, 34) or list of (N, 34) arrays per image
        paths: Image paths
        fname: Output filename
        names: Class names
        calibs: Calibration parameters for each image
        conf_thres: Confidence threshold for filtering predictions
        return_array: If True, return numpy array instead of saving to file
        face_3d_classes: List/set of class IDs that use face-based reconstruction
        complete_3d_classes: List/set of class IDs that use complete 3D box

    Returns:
        Thread object (when running in background) or numpy array if return_array=True
    """
    if return_array:
        # Synchronous call - return numpy array directly
        return _plot_predictions_3d_core(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres, return_array=True,
                                         face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)
    else:
        # Threaded call
        return _plot_predictions_3d_threaded(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres,
                                             face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)


def _plot_predictions_3d_core(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg",
                               names=None, calibs=None, conf_thres=0.25, return_array=False,
                               face_3d_classes=None, complete_3d_classes=None):
    """Core implementation for plotting 2x2 grid with predictions: target 2D, target 3D, prediction 2D, prediction 3D.

    Layout for TensorBoard logging (single sample):
        (1,1) Target 2D    | (1,2) Target 3D
        (2,1) Prediction 2D | (2,2) Prediction 3D

    Note: This function only draws predictions. For targets, use plot_val_comparison_3d.

    Args:
        images: Batch of images (B, C, H, W) in BGR format (from dataloader3d)
        predictions_2d: List of 2D predictions per image, each (N, 6) [x1, y1, x2, y2, conf, cls]
        predictions_3d: 3D predictions - list of (N, 34) arrays per image
        paths: List of image paths
        fname: Output filename
        names: Class names dict
        calibs: Tuple of calibration dicts from dataloader
        conf_thres: Confidence threshold
        return_array: If True, return numpy array instead of saving to file
    """
    if isinstance(images, torch.Tensor):
        images = images.cpu().float().numpy()

    # Only plot first sample for 2x2 grid
    _, _, h, w = images.shape

    if np.max(images[0]) <= 1:
        images = images * 255  # de-normalise

    # Convert BGR to RGB (dataloader3d returns BGR format)
    images = images[:, ::-1, :, :]  # Reverse channel order: BGR -> RGB

    # Get first image
    im = images[0].transpose(1, 2, 0).astype(np.uint8)
    im_pred_2d = im.copy()
    im_pred_3d = im.copy()

    # Get calibration for first image
    calib = None
    if calibs is not None and len(calibs) > 0:
        calib = calibs[0]
    elif paths:
        calib = read_calib_from_path(paths[0])

    # Get predictions for first image
    pred_2d = None
    if predictions_2d is not None and len(predictions_2d) > 0:
        pred_2d = predictions_2d[0]
        if isinstance(pred_2d, torch.Tensor):
            pred_2d = pred_2d.cpu().numpy()

    pred_3d = None
    if predictions_3d is not None:
        if isinstance(predictions_3d, list) and len(predictions_3d) > 0:
            pred_3d = predictions_3d[0]
            if isinstance(pred_3d, torch.Tensor):
                pred_3d = pred_3d.cpu().numpy()
        elif isinstance(predictions_3d, torch.Tensor):
            pred_3d = predictions_3d[0].cpu().numpy()
        elif isinstance(predictions_3d, np.ndarray) and len(predictions_3d) > 0:
            pred_3d = predictions_3d[0]

    # --- Draw Prediction 2D boxes ---
    if pred_2d is not None:
        for pred in pred_2d:
            if len(pred) < 6:
                continue
            x1, y1, x2, y2, conf, cls = pred[:6]

            cls = int(cls)
            color = get_class_color(cls)
            cv2.rectangle(im_pred_2d, (int(x1), int(y1)), (int(x2), int(y2)), color, 1)

    # --- Draw Prediction 3D boxes ---
    if pred_2d is not None and pred_3d is not None and len(pred_3d) > 0:
        for j, pred_2d_box in enumerate(pred_2d):
            if len(pred_2d_box) < 6 or j >= len(pred_3d):
                continue

            box_info = decode_and_reconstruct_3d_box(
                pred_2d_box, pred_3d[j], calib, conf_thres, face_3d_classes, complete_3d_classes
            )

            if box_info is None or not box_info['should_draw']:
                continue

            if box_info['corners_3d'] is None:
                continue

            im_pred_3d = draw_3d_box_from_corners(
                im_pred_3d,
                box_info['corners_3d'],
                calib,
                (w, h),
                face_center_2d=box_info['face_center_2d'],
                face_color=box_info['face_color']
            )

    # Create 2x2 figure with placeholders for targets (top row)
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # Top row: placeholder for targets
    im_target_2d = im.copy()
    im_target_3d = im.copy()
    cv2.putText(im_target_2d, 'Target 2D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2)
    cv2.putText(im_target_3d, 'Target 3D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2)

    axes[0, 0].imshow(im_target_2d)
    axes[0, 0].set_title('Target 2D', fontsize=12)
    axes[0, 0].axis('off')

    axes[0, 1].imshow(im_target_3d)
    axes[0, 1].set_title('Target 3D', fontsize=12)
    axes[0, 1].axis('off')

    # Bottom row: predictions
    cv2.putText(im_pred_2d, 'Prediction 2D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2)
    cv2.putText(im_pred_3d, 'Prediction 3D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2)

    axes[1, 0].imshow(im_pred_2d)
    axes[1, 0].set_title('Prediction 2D', fontsize=12)
    axes[1, 0].axis('off')

    axes[1, 1].imshow(im_pred_3d)
    axes[1, 1].set_title('Prediction 3D', fontsize=12)
    axes[1, 1].axis('off')

    if paths:
        fig.suptitle(Path(paths[0]).name[:50], fontsize=10)

    plt.tight_layout()

    if return_array:
        dpi = 150
        fig.set_dpi(dpi)
        fig.canvas.draw()

        buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
        fig_w, fig_h = fig.canvas.get_width_height()
        img_array = buf.reshape(fig_h, fig_w, 4)
        img_array = img_array[:, :, :3]
        plt.close(fig)
        return img_array
    else:
        plt.savefig(fname, dpi=150, bbox_inches='tight')
        plt.close(fig)


@threaded
def _plot_predictions_3d_threaded(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg",
                                  names=None, calibs=None, conf_thres=0.25,
                                  face_3d_classes=None, complete_3d_classes=None):
    """Threaded wrapper for plotting 3D predictions."""
    _plot_predictions_3d_core(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres, return_array=False,
                              face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes)


def plot_val_comparison_3d(images, targets, predictions_2d, predictions_3d=None, paths=None,
                           fname="val_comparison.jpg", names=None, calibs=None,
                           return_array=False, face_3d_classes=None, complete_3d_classes=None):
    """Plot 2x2 comparison grid: target 2D, target 3D, prediction 2D, prediction 3D.

    This is the main function for TensorBoard logging during validation.

    Layout:
        (1,1) Target 2D    | (1,2) Target 3D
        (2,1) Prediction 2D | (2,2) Prediction 3D

    Args:
        images: Batch tensor (B, 3, H, W) normalized [0, 1] in BGR format
        targets: Target tensor with 48 columns per target
        predictions_2d: List of 2D predictions per image, each (N, 6)
        predictions_3d: 3D predictions - list of (N, 34) arrays per image
        paths: List of image paths
        fname: Output filename
        names: Class names dict
        calibs: Tuple of calibration dicts
        return_array: If True, return numpy array instead of saving to file
        face_3d_classes: List/set of class IDs that use face-based reconstruction
        complete_3d_classes: List/set of class IDs that use complete 3D box

    Returns:
        numpy array (H, W, 3) RGB if return_array=True, else None
    """
    if isinstance(images, torch.Tensor):
        images = images.cpu().float().numpy()
    if isinstance(targets, torch.Tensor):
        targets = targets.cpu().numpy()

    _, _, h, w = images.shape

    if np.max(images[0]) <= 1:
        images = images * 255

    # Keep BGR format for cv2 drawing (will convert to RGB at the end for matplotlib)
    # images = images[:, ::-1, :, :]  # Don't convert here

    # Get first image (BGR format)
    im = images[0].transpose(1, 2, 0).astype(np.uint8)
    im_target_2d = im.copy()
    im_target_3d = im.copy()
    im_pred_2d = im.copy()
    im_pred_3d = im.copy()

    # Get targets for first image
    ti = targets[targets[:, 0] == 0] if len(targets) > 0 else np.array([])

    # Get calibration
    calib = None
    if calibs is not None and len(calibs) > 0:
        calib = calibs[0]
    elif paths:
        calib = read_calib_from_path(paths[0])

    # Apply project-correct defaults for class categorization (face-based vs complete 3D box)
    # face_3d_classes: only vehicles (0) and tricycles (13) have GT face annotations
    # complete_3d_classes: pedestrian (1), bicycle (2), rider (3) use full 3D box without face data
    _face_3d_classes = face_3d_classes if face_3d_classes is not None else [0, 13]
    _complete_3d_classes = complete_3d_classes if complete_3d_classes is not None else [1, 2, 3]

    # --- Draw Target 2D boxes ---
    for t in ti:
        if len(t) < 6 or np.isnan(t[1]):
            continue
        cls = int(t[1])
        color = get_class_color(cls)

        x_center, y_center, bw, bh = t[2:6]
        if np.any(np.isnan([x_center, y_center, bw, bh])):
            continue
        x1 = int((x_center - bw / 2) * w)
        y1 = int((y_center - bh / 2) * h)
        x2 = int((x_center + bw / 2) * w)
        y2 = int((y_center + bh / 2) * h)

        cv2.rectangle(im_target_2d, (x1, y1), (x2, y2), color, 1)

    # --- Draw Target 3D boxes ---
    for t in ti:
        if len(t) < 15 or np.isnan(t[1]):
            continue

        cls = int(t[1])
        decoded = decode_and_reconstruct_3d_box_from_target(t, calib, w, h, _face_3d_classes, _complete_3d_classes)

        if decoded and decoded['should_draw'] and decoded['corners_3d'] is not None:
            if cls in _face_3d_classes:
                im_target_3d = draw_3d_box_from_corners(
                    im_target_3d,
                    decoded['corners_3d'],
                    calib,
                    (w, h),
                    face_center_2d=decoded.get('face_center_2d'),
                    face_color=decoded.get('face_color')
                )
            elif cls in _complete_3d_classes:
                # Use draw_3d_box_from_corners which handles corner reordering internally
                # This ensures consistent front/rear face colors (red=front, blue=rear)
                # and fixes the head/rear reversal bug in TensorBoard logs
                im_target_3d = draw_3d_box_from_corners(
                    im_target_3d,
                    decoded['corners_3d'],
                    calib,
                    (w, h),
                    face_center_2d=decoded.get('center_2d'),
                    face_color=colors(cls)
                )

    # --- Draw Prediction 2D boxes ---
    pred_2d = None
    if predictions_2d is not None and len(predictions_2d) > 0:
        pred_2d = predictions_2d[0]
        if isinstance(pred_2d, torch.Tensor):
            pred_2d = pred_2d.cpu().numpy()

        for pred in pred_2d:
            if len(pred) < 6:
                continue
            x1, y1, x2, y2, conf, cls = pred[:6]

            cls = int(cls)
            color = get_class_color(cls)
            cv2.rectangle(im_pred_2d, (int(x1), int(y1)), (int(x2), int(y2)), color, 1)

    # --- Draw Prediction 3D boxes ---
    pred_3d = None
    if predictions_3d is not None:
        if isinstance(predictions_3d, list) and len(predictions_3d) > 0:
            pred_3d = predictions_3d[0]
            if isinstance(pred_3d, torch.Tensor):
                pred_3d = pred_3d.cpu().numpy()
        elif isinstance(predictions_3d, torch.Tensor):
            pred_3d = predictions_3d[0].cpu().numpy()
        elif isinstance(predictions_3d, np.ndarray) and len(predictions_3d) > 0:
            pred_3d = predictions_3d[0]

    if pred_2d is not None and pred_3d is not None and len(pred_3d) > 0:
        for j, pred_2d_box in enumerate(pred_2d):
            if len(pred_2d_box) < 6 or j >= len(pred_3d):
                continue

            box_info = decode_and_reconstruct_3d_box(
                pred_2d_box, pred_3d[j], calib, 0.0, face_3d_classes, complete_3d_classes  # Predictions already filtered by conf_thres_plot
            )

            if box_info is None or not box_info['should_draw']:
                continue

            if box_info['corners_3d'] is None:
                continue

            im_pred_3d = draw_3d_box_from_corners(
                im_pred_3d,
                box_info['corners_3d'],
                calib,
                (w, h),
                face_center_2d=box_info['face_center_2d'],
                face_color=box_info['face_color']
            )

    # Create 2x2 figure
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # Convert BGR to RGB for matplotlib display
    axes[0, 0].imshow(im_target_2d[:, :, ::-1])
    axes[0, 0].set_title('Target 2D', fontsize=14, fontweight='bold', color='darkblue')
    axes[0, 0].axis('off')

    axes[0, 1].imshow(im_target_3d[:, :, ::-1])
    axes[0, 1].set_title('Target 3D', fontsize=14, fontweight='bold', color='darkblue')
    axes[0, 1].axis('off')

    axes[1, 0].imshow(im_pred_2d[:, :, ::-1])
    axes[1, 0].set_title('Prediction 2D', fontsize=14, fontweight='bold', color='darkgreen')
    axes[1, 0].axis('off')

    axes[1, 1].imshow(im_pred_3d[:, :, ::-1])
    axes[1, 1].set_title('Prediction 3D', fontsize=14, fontweight='bold', color='darkgreen')
    axes[1, 1].axis('off')

    if paths:
        fig.suptitle(Path(paths[0]).name[:50], fontsize=12, fontweight='bold')

    plt.tight_layout()

    if return_array:
        dpi = 300
        fig.set_dpi(dpi)
        fig.canvas.draw()

        buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
        fig_w, fig_h = fig.canvas.get_width_height()
        img_array = buf.reshape(fig_h, fig_w, 4)
        img_array = img_array[:, :, :3]
        plt.close(fig)
        return img_array
    else:
        plt.savefig(fname, dpi=300, bbox_inches='tight')
        plt.close(fig)
        return None


def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=""):
    """Plots learning rate schedule for given optimizer and scheduler, saving plot to `save_dir`."""
    optimizer, scheduler = copy(optimizer), copy(scheduler)  # do not modify originals
    y = []
    for _ in range(epochs):
        scheduler.step()
        y.append(optimizer.param_groups[0]["lr"])
    plt.plot(y, ".-", label="LR")
    plt.xlabel("epoch")
    plt.ylabel("LR")
    plt.grid()
    plt.xlim(0, epochs)
    plt.ylim(0)
    plt.savefig(Path(save_dir) / "LR.png", dpi=200)
    plt.close()


def plot_val_txt():
    """Plots 2D and 1D histograms of bounding box centers from 'val.txt' using matplotlib, saving as 'hist2d.png' and
    'hist1d.png'.

    Example: from utils.plots import *; plot_val()
    """
    x = np.loadtxt("val.txt", dtype=np.float32)
    box = xyxy2xywh(x[:, :4])
    cx, cy = box[:, 0], box[:, 1]

    fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
    ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
    ax.set_aspect("equal")
    plt.savefig("hist2d.png", dpi=300)

    _fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
    ax[0].hist(cx, bins=600)
    ax[1].hist(cy, bins=600)
    plt.savefig("hist1d.png", dpi=200)


def plot_targets_txt():
    """Plots histograms of object detection targets from 'targets.txt', saving the figure as 'targets.jpg'.

    Example: from utils.plots import *; plot_targets_txt()
    """
    x = np.loadtxt("targets.txt", dtype=np.float32).T
    s = ["x targets", "y targets", "width targets", "height targets"]
    _fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
    ax = ax.ravel()
    for i in range(4):
        ax[i].hist(x[i], bins=100, label=f"{x[i].mean():.3g} +/- {x[i].std():.3g}")
        ax[i].legend()
        ax[i].set_title(s[i])
    plt.savefig("targets.jpg", dpi=200)


def plot_val_study(file="", dir="", x=None):
    """Plots validation study results from 'study*.txt' files in a directory or a specific file, comparing model
    performance and speed.

    Example: from utils.plots import *; plot_val_study()
    """
    save_dir = Path(file).parent if file else Path(dir)
    plot2 = False  # plot additional results
    if plot2:
        ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)[1].ravel()

    _fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
    # for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]:
    for f in sorted(save_dir.glob("study*.txt")):
        y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
        x = np.arange(y.shape[1]) if x is None else np.array(x)
        if plot2:
            s = ["P", "R", "mAP@.5", "mAP@.5:.95", "t_preprocess (ms/img)", "t_inference (ms/img)", "t_NMS (ms/img)"]
            for i in range(7):
                ax[i].plot(x, y[i], ".-", linewidth=2, markersize=8)
                ax[i].set_title(s[i])

        j = y[3].argmax() + 1
        ax2.plot(
            y[5, 1:j],
            y[3, 1:j] * 1e2,
            ".-",
            linewidth=2,
            markersize=8,
            label=f.stem.replace("study_coco_", "").replace("yolo", "YOLO"),
        )

    ax2.plot(
        1e3 / np.array([209, 140, 97, 58, 35, 18]),
        [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
        "k.-",
        linewidth=2,
        markersize=8,
        alpha=0.25,
        label="EfficientDet",
    )

    ax2.grid(alpha=0.2)
    ax2.set_yticks(np.arange(20, 60, 5))
    ax2.set_xlim(0, 57)
    ax2.set_ylim(25, 55)
    ax2.set_xlabel("GPU Speed (ms/img)")
    ax2.set_ylabel("COCO AP val")
    ax2.legend(loc="lower right")
    f = save_dir / "study.png"
    print(f"Saving {f}...")
    plt.savefig(f, dpi=300)


@TryExcept()  # known issue https://github.com/ultralytics/yolov5/issues/5395
def plot_labels(labels, names=(), save_dir=Path("")):
    """Plots dataset labels, saving correlogram and label images, handles classes, and visualizes bounding boxes."""
    LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
    c, b = labels[:, 0], labels[:, 1:].transpose()  # classes, boxes
    nc = int(c.max() + 1)  # number of classes
    x = pd.DataFrame(b.transpose(), columns=["x", "y", "width", "height"])
    labels_plot = labels

    # seaborn correlogram (can be slow with large datasets)
    sn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
    plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200)
    plt.close()

    # matplotlib labels
    matplotlib.use("svg")  # faster
    ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
    y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
    with contextlib.suppress(Exception):  # color histogram bars by class
        [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)]  # known issue #3195
    ax[0].set_ylabel("instances")
    if 0 < len(names) < 30:
        ax[0].set_xticks(range(len(names)))
        ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
    else:
        ax[0].set_xlabel("classes")
    sn.histplot(x, x="x", y="y", ax=ax[2], bins=50, pmax=0.9)
    sn.histplot(x, x="width", y="height", ax=ax[3], bins=50, pmax=0.9)

    # rectangles - use subsampled labels
    labels_rect = labels_plot.copy()
    labels_rect[:, 1:3] = 0.5  # center
    labels_rect[:, 1:] = xywh2xyxy(labels_rect[:, 1:]) * 2000
    img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255)
    for cls, *box in labels_rect[:1000]:
        ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls))  # plot
    ax[1].imshow(img)
    ax[1].axis("off")

    for a in [0, 1, 2, 3]:
        for s in ["top", "right", "left", "bottom"]:
            ax[a].spines[s].set_visible(False)

    plt.savefig(save_dir / "labels.jpg", dpi=200)
    matplotlib.use("Agg")
    plt.close()


def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path("images.jpg")):
    """Displays a grid of images with optional labels and predictions, saving to a file."""
    from utils.augmentations import denormalize

    names = names or [f"class{i}" for i in range(1000)]
    blocks = torch.chunk(
        denormalize(im.clone()).cpu().float(), len(im), dim=0
    )  # select batch index 0, block by channels
    n = min(len(blocks), nmax)  # number of plots
    m = min(8, round(n**0.5))  # 8 x 8 default
    _fig, ax = plt.subplots(math.ceil(n / m), m)  # 8 rows x n/8 cols
    ax = ax.ravel() if m > 1 else [ax]
    # plt.subplots_adjust(wspace=0.05, hspace=0.05)
    for i in range(n):
        ax[i].imshow(blocks[i].squeeze().permute((1, 2, 0)).numpy().clip(0.0, 1.0))
        ax[i].axis("off")
        if labels is not None:
            from utils.general import get_class_name
            s = get_class_name(names, labels[i]) + (f"—{get_class_name(names, pred[i])}" if pred is not None else "")
            ax[i].set_title(s, fontsize=8, verticalalignment="top")
    plt.savefig(f, dpi=300, bbox_inches="tight")
    plt.close()
    if verbose:
        LOGGER.info(f"Saving {f}")
        if labels is not None:
            LOGGER.info("True:     " + " ".join(f"{get_class_name(names, i):3s}" for i in labels[:nmax]))
        if pred is not None:
            LOGGER.info("Predicted:" + " ".join(f"{get_class_name(names, i):3s}" for i in pred[:nmax]))
    return f


def plot_evolve(evolve_csv="path/to/evolve.csv"):
    """Plots hyperparameter evolution results from a given CSV, saving the plot and displaying best results.

    Example: from utils.plots import *; plot_evolve()
    """
    evolve_csv = Path(evolve_csv)
    data = pd.read_csv(evolve_csv)
    keys = [x.strip() for x in data.columns]
    x = data.values
    f = fitness(x)
    j = np.argmax(f)  # max fitness index
    plt.figure(figsize=(10, 12), tight_layout=True)
    matplotlib.rc("font", **{"size": 8})
    print(f"Best results from row {j} of {evolve_csv}:")
    for i, k in enumerate(keys[7:]):
        v = x[:, 7 + i]
        mu = v[j]  # best single result
        plt.subplot(6, 5, i + 1)
        plt.scatter(v, f, c=hist2d(v, f, 20), cmap="viridis", alpha=0.8, edgecolors="none")
        plt.plot(mu, f.max(), "k+", markersize=15)
        plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9})  # limit to 40 characters
        if i % 5 != 0:
            plt.yticks([])
        print(f"{k:>15}: {mu:.3g}")
    f = evolve_csv.with_suffix(".png")  # filename
    plt.savefig(f, dpi=200)
    plt.close()
    print(f"Saved {f}")


def plot_results(file="path/to/results.csv", dir=""):
    """Plots training results from a 'results.csv' file; accepts file path and directory as arguments.

    Example: from utils.plots import *; plot_results('path/to/results.csv')
    """
    save_dir = Path(file).parent if file else Path(dir)
    fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
    ax = ax.ravel()
    files = list(save_dir.glob("results*.csv"))
    assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
    for f in files:
        try:
            data = pd.read_csv(f)
            s = [x.strip() for x in data.columns]
            x = data.values[:, 0]
            for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]):
                y = data.values[:, j].astype("float")
                # y[y == 0] = np.nan  # don't show zero values
                ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8)  # actual results
                ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2)  # smoothing line
                ax[i].set_title(s[j], fontsize=12)
                # if j in [8, 9, 10]:  # share train and val loss y axes
                #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
        except Exception as e:
            LOGGER.info(f"Warning: Plotting error for {f}: {e}")
    ax[1].legend()
    fig.savefig(save_dir / "results.png", dpi=200)
    plt.close()


def profile_idetection(start=0, stop=0, labels=(), save_dir=""):
    """Plots per-image iDetection logs, comparing metrics like storage and performance over time.

    Example: from utils.plots import *; profile_idetection()
    """
    ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel()
    s = ["Images", "Free Storage (GB)", "RAM Usage (GB)", "Battery", "dt_raw (ms)", "dt_smooth (ms)", "real-world FPS"]
    files = list(Path(save_dir).glob("frames*.txt"))
    for fi, f in enumerate(files):
        try:
            results = np.loadtxt(f, ndmin=2).T[:, 90:-30]  # clip first and last rows
            n = results.shape[1]  # number of rows
            x = np.arange(start, min(stop, n) if stop else n)
            results = results[:, x]
            t = results[0] - results[0].min()  # set t0=0s
            results[0] = x
            for i, a in enumerate(ax):
                if i < len(results):
                    label = labels[fi] if len(labels) else f.stem.replace("frames_", "")
                    a.plot(t, results[i], marker=".", label=label, linewidth=1, markersize=5)
                    a.set_title(s[i])
                    a.set_xlabel("time (s)")
                    # if fi == len(files) - 1:
                    #     a.set_ylim(bottom=0)
                    for side in ["top", "right"]:
                        a.spines[side].set_visible(False)
                else:
                    a.remove()
        except Exception as e:
            print(f"Warning: Plotting error for {f}; {e}")
    ax[1].legend()
    plt.savefig(Path(save_dir) / "idetection_profile.png", dpi=200)


def save_one_box(xyxy, im, file=Path("im.jpg"), gain=1.02, pad=10, square=False, BGR=False, save=True):
    """Crops and saves an image from bounding box `xyxy`, applied with `gain` and `pad`, optionally squares and adjusts
    for BGR.
    """
    xyxy = torch.tensor(xyxy).view(-1, 4)
    b = xyxy2xywh(xyxy)  # boxes
    if square:
        b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # attempt rectangle to square
    b[:, 2:] = b[:, 2:] * gain + pad  # box wh * gain + pad
    xyxy = xywh2xyxy(b).long()
    clip_boxes(xyxy, im.shape)
    crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR else -1)]
    if save:
        file.parent.mkdir(parents=True, exist_ok=True)  # make directory
        f = str(increment_path(file).with_suffix(".jpg"))
        # cv2.imwrite(f, crop)  # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
        Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0)  # save RGB
    return crop


# bev视角画板
def draw_bev_blank(max_range=200):
    """Create a blank BEV image for visualization.

    Args:
        max_range: Maximum forward range in meters (default 200m)

    Returns:
        BEV image with grid lines
    """
    # Scale: 20 pixels per meter for high resolution (2x for clearer saved images)
    pixels_per_meter = 20
    img_height = max_range * pixels_per_meter  # 4000 for 200m
    img_width = 100 * pixels_per_meter  # 2000 for 100m lateral range (-50m to +50m)

    bevimg = np.ones((img_height, img_width, 3), dtype=np.uint8) * 255
    ego_center_x = img_width // 2
    ego_center_y = img_height  # Ego car position at bottom center

    # Dimensions of ego car in pixels (e.g., 4.5m long, 1.8m wide)
    ego_l_px = int(4.5 * pixels_per_meter)
    ego_w_px = int(1.8 * pixels_per_meter)
    ego_half_l, ego_half_w = ego_l_px // 2, ego_w_px // 2

    ego_box = np.array([
        [ego_center_x - ego_half_w, ego_center_y - ego_half_l],
        [ego_center_x + ego_half_w, ego_center_y - ego_half_l],
        [ego_center_x + ego_half_w, ego_center_y + ego_half_l],
        [ego_center_x - ego_half_w, ego_center_y + ego_half_l]
    ], dtype=np.int32)

    # Plot horizontal grid by distance (every 20m for 200m range)
    grid_step_m = 20  # meters per grid line
    grid_step_px = grid_step_m * pixels_per_meter
    num_grids = max_range // grid_step_m

    for i in range(num_grids + 1):
        y_pos = ego_center_y - i * grid_step_px
        if y_pos >= 0:
            cv2.line(bevimg, (0, y_pos), (img_width, y_pos), (180, 180, 180), 3, cv2.LINE_AA)
            cv2.putText(bevimg, f"{i * grid_step_m}m", (ego_center_x + 15, y_pos - 15),
                       cv2.FONT_HERSHEY_SIMPLEX, 1.2, (80, 80, 80), 2, cv2.LINE_AA)

    # Plot vertical grid (every 10m, range from -50m to +50m)
    lateral_range = 50  # meters each side
    lateral_step_m = 10
    lateral_step_px = lateral_step_m * pixels_per_meter

    for i in range(-lateral_range // lateral_step_m, lateral_range // lateral_step_m + 1):
        x_pos = ego_center_x + i * lateral_step_px
        if 0 <= x_pos < img_width:
            cv2.line(bevimg, (x_pos, 0), (x_pos, img_height), (180, 180, 180), 3, cv2.LINE_AA)
            # Only draw label for non-zero positions to avoid clutter
            if i != 0:
                cv2.putText(bevimg, f"{i * lateral_step_m}m", (x_pos - 40, img_height - 20),
                           cv2.FONT_HERSHEY_SIMPLEX, 1.0, (80, 80, 80), 2, cv2.LINE_AA)

    cv2.drawContours(bevimg, [ego_box], 0, (255, 0, 0), -1)  # Blue filled rectangle for ego

    return bevimg

def cam_corners_front_rear(pred3d, facetype):
    dims = pred3d[3:6]
    corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1)

    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
    # use relative origin [0.5, 1, 0.5]
    if facetype in ['tail', 1]:
        # front
        corners_norm = corners_norm - [0, 0.5, 0.5]
    elif facetype in ['front', 0]:
        # tail
        corners_norm = corners_norm - [1, 0.5, 0.5]
    elif facetype in ['right', 3]:
        # left
        corners_norm = corners_norm - [0.5, 0.5, 0]
    elif facetype in ['left', 2]:
        # right
        corners_norm = corners_norm - [0.5, 0.5, 1]
    elif facetype in ['whole', -1]:
        # center
        corners_norm = corners_norm - [0.5, 0.5, 0.5]
    else:
        raise AssertionError('Non valid face type')
    corners = dims.reshape([1, 3]) * corners_norm.reshape([8, 3])
    # rotate around y axis
    corners = rotation_3d_in_axis(corners, pred3d[6], axis=1)
    corners += pred3d[:3].reshape(1, 3)
    return corners

def drawbev(bevimg, vehicle3d, is_pred=True):
    """Draw a vehicle on BEV image.

    Args:
        bevimg: BEV image to draw on
        vehicle3d: [x, y, z, l, h, w, ..., rotation_y, face_type]
        is_pred: True for predictions (red), False for GT (green)
    """
    x, y, z, l, h, w = vehicle3d[0], vehicle3d[1], vehicle3d[2], vehicle3d[3], vehicle3d[4], vehicle3d[5]
    rotation_y = vehicle3d[-2]
    face_type = vehicle3d[-1]  # 'front','rear','left','right','whole'

    # BEV image parameters - must match draw_bev_blank()
    # Scale: 20 pixels per meter for high resolution (2x for clearer saved images)
    pixels_per_meter = 20
    max_range = 200  # meters forward
    lateral_range = 50  # meters each side

    img_height = bevimg.shape[0]
    img_width = bevimg.shape[1]
    ego_center_x = img_width // 2
    ego_center_y = img_height

    # Filter out objects outside range
    if x > lateral_range or x < -lateral_range or z > max_range or z < 0:
        return bevimg

    corners = cam_corners_front_rear(np.array([x, y, z, l, h, w, rotation_y]), face_type)
    xyz3d_0 = np.mean(corners[4:8, :], axis=0)  # Front_point
    xyz3d_center = np.mean(corners[0:8, :], axis=0)  # center

    # Convert 3D coordinates to BEV pixel coordinates
    # x (lateral) -> horizontal in BEV, z (forward) -> vertical in BEV (inverted)
    center = (int(ego_center_x + xyz3d_center[0] * pixels_per_meter),
              int(ego_center_y - xyz3d_center[2] * pixels_per_meter))
    front_point = (int(ego_center_x + xyz3d_0[0] * pixels_per_meter),
                   int(ego_center_y - xyz3d_0[2] * pixels_per_meter))

    H, W = l * pixels_per_meter, w * pixels_per_meter
    size = (H, W)
    angle = np.degrees(rotation_y)
    rect = (center, size, angle)
    box = cv2.boxPoints(rect)
    box = np.intp(box)

    if is_pred:
        color = (0, 0, 255)  # Red for predictions
    else:
        color = (0, 255, 0)  # Green for ground truth

    cv2.drawContours(bevimg, [box], 0, color, 3, cv2.LINE_AA)
    cv2.arrowedLine(bevimg, center, front_point, color, thickness=3, tipLength=0.3, line_type=cv2.LINE_AA)

    return bevimg
