from __future__ import annotations

import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Optional

import cv2
import numpy as np
import torch
import yaml

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from tools.model_inference.adapters.video_dir_inference_utils import (
    iter_video_case_frames,
    read_video_frame_index,
    resolve_video_case_paths,
)
from ultralytics.data.ground3d_augment import (
    adjust_calib_for_roi_crop,
    build_final_resized_calib,
    compute_centered_roi_bounds,
)
from ultralytics.nn.tasks import load_checkpoint
from ultralytics.utils.plotting import colors
from ultralytics.utils.plotting_3d import (
    EDGE_YAW_MAX_LATERAL_DIST_M,
    decode_edge_yaw_selection_from_prediction,
    decode_3d_prediction,
    draw_3d_box,
    extract_face_regressed_size_priors_from_prediction,
    extract_3d_attrs_from_prediction,
    project_face_bottom_edge,
    project_partial_face_bottom_edge,
    reconstruct_edge_based_box_from_selection,
)
from ultralytics.utils.torch_utils import select_device


DEFAULT_VISUALIZATION_ROOT = FILE.parent / "visualization"
IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M = EDGE_YAW_MAX_LATERAL_DIST_M
GRID_BG_COLOR = (24, 24, 24)


@dataclass(frozen=True)
class ROIModelSpec:
    name: str
    model_path: str
    roi_size: tuple[int, int]
    crop_center_mode: str
    virtual_fx: float
    imgsz: Optional[tuple[int, int]] = None
    conf: float = 0.25
    max_det: int = 300


@dataclass
class LoadedROIModel:
    spec: ROIModelSpec
    model: torch.nn.Module
    names: dict[int, str]
    face_3d_classes: set[int]
    complete_3d_classes: set[int]
    fake_3d_classes: set[int]
    imgsz: tuple[int, int]


@dataclass
class PreparedROI:
    name: str
    image: np.ndarray
    crop_bounds: tuple[int, int, int, int]
    calib: dict[str, Any]
    vp_x: float
    vp_y: float
    crop_center_x: float
    crop_center_y: float


@dataclass
class InferenceContext:
    device: torch.device
    use_half: bool
    classes: Optional[set[int]]
    roi_models: list[LoadedROIModel]
    edge_yaw_max_lateral_dist_m: float
    inference_batch_size: int


@dataclass
class RawROIOutputs:
    detections: np.ndarray
    preds_3d: np.ndarray
    preds_3d_fake: Optional[np.ndarray]
    preds_edge: Optional[np.ndarray]
    anchors: np.ndarray
    strides: np.ndarray


def _coerce_imgsz(imgsz: Any) -> Optional[tuple[int, int]]:
    if imgsz is None:
        return None
    if isinstance(imgsz, int):
        return (int(imgsz), int(imgsz))
    if isinstance(imgsz, str):
        parts = [part.strip() for part in imgsz.split(",") if part.strip()]
        if len(parts) == 1:
            value = int(parts[0])
            return (value, value)
        if len(parts) == 2:
            return (int(parts[0]), int(parts[1]))
        raise ValueError(f"Unable to parse imgsz={imgsz!r}")
    if isinstance(imgsz, (list, tuple)):
        if len(imgsz) == 1:
            value = int(imgsz[0])
            return (value, value)
        if len(imgsz) == 2:
            return (int(imgsz[0]), int(imgsz[1]))
    raise ValueError(f"Unsupported imgsz value: {imgsz!r}")


def _infer_model_imgsz(model: torch.nn.Module, override: Optional[tuple[int, int]]) -> tuple[int, int]:
    if override is not None:
        return override
    model_args = getattr(model, "args", {}) or {}
    imgsz = _coerce_imgsz(model_args.get("imgsz"))
    return imgsz or (768, 352)


def _load_yaml_if_present(path: str) -> dict[str, Any]:
    if not path:
        return {}
    yaml_path = Path(path)
    if not yaml_path.exists():
        return {}
    with yaml_path.open("r", encoding="utf-8") as file:
        return yaml.safe_load(file) or {}


def camera4_payload_to_raw_calib(payload: dict[str, Any]) -> dict[str, Any]:
    """Normalize a flat or combined `camera4.json` payload to the raw calib format used by ROI inference."""
    if "intrinsics" in payload:
        intrinsics = payload.get("intrinsics", {}).get("camera4.json", {}) or {}
        extrinsics = payload.get("extrinsics", {}).get("camera4.json", {}) or {}
        pitch = 0.0
        rpy = extrinsics.get("rpy", [0.0, 0.0, 0.0])
        if len(rpy) > 1:
            pitch = math.radians(float(rpy[1]))
        return {
            "focal_u": float(intrinsics["focal_u"]),
            "focal_v": float(intrinsics["focal_v"]),
            "cu": float(intrinsics["cu"]),
            "cv": float(intrinsics["cv"]),
            "pitch": pitch,
            "distort_coeffs": list(intrinsics.get("distort_coeffs", [])),
            "image_width": payload.get("image_width", payload.get("img_width")),
            "image_height": payload.get("image_height", payload.get("img_height")),
            "source_format": "combined_calibration",
            "angle_unit": "radians",
        }

    required = ("focal_u", "focal_v", "cu", "cv")
    missing = [key for key in required if key not in payload]
    if missing:
        raise KeyError(f"camera4 payload missing required keys: {missing}")

    # Match ultralytics.data.ground3d_augment.read_calib_from_path() flat-file behavior:
    # keep the payload values as-is instead of applying extra angle conversions here.
    calib = dict(payload)
    calib["focal_u"] = float(payload["focal_u"])
    calib["focal_v"] = float(payload["focal_v"])
    calib["cu"] = float(payload["cu"])
    calib["cv"] = float(payload["cv"])
    calib["distort_coeffs"] = list(payload.get("distort_coeffs", []))
    calib["source_format"] = "flat_camera4"
    calib["angle_unit"] = "degrees"
    return calib


def load_camera4_calib(calib_path: str | Path) -> dict[str, Any]:
    """Load a flat or combined `camera4.json` file into the raw calib format used by training-time ROI code."""
    calib_path = Path(calib_path)
    with calib_path.open("r", encoding="utf-8") as file:
        payload = json.load(file)
    return camera4_payload_to_raw_calib(payload)


def _to_radians(angle: float, unit: str) -> float:
    if unit == "degrees":
        return math.radians(float(angle))
    return float(angle)


def _compute_vanishing_point_xy(raw_calib: dict[str, Any], ori_w: int, ori_h: int) -> tuple[float, float]:
    """Compute vanishing point using the explicit cx/cy/fx/fy/pitch/yaw formula."""
    if raw_calib is None:
        return ori_w / 2.0, ori_h / 2.0

    cx = float(raw_calib.get("cu", ori_w / 2.0))
    cy = float(raw_calib.get("cv", ori_h / 2.0))
    fx = float(raw_calib.get("focal_u", ori_w))
    fy = float(raw_calib.get("focal_v", ori_h))
    angle_unit = str(raw_calib.get("angle_unit", "radians"))
    yaw = _to_radians(float(raw_calib.get("yaw", 0.0)), angle_unit)
    pitch = _to_radians(float(raw_calib.get("pitch", 0.0)), angle_unit)
    vp_x = cx + fx * math.tan(yaw)
    vp_y = cy - fy * math.tan(pitch)
    return float(vp_x), float(vp_y)


def _resize_ground3d_image_in_steps(
    image_bgr: np.ndarray,
    target_size: tuple[int, int],
    interpolation: int = cv2.INTER_LINEAR,
) -> np.ndarray:
    """Match Ground3D training resize with repeated 0.5x downsampling before the final resize."""
    target_w, target_h = int(target_size[0]), int(target_size[1])
    current_h, current_w = image_bgr.shape[:2]

    if (current_w, current_h) == (target_w, target_h):
        return image_bgr

    resized = image_bgr
    if target_w < current_w and target_h < current_h:
        while True:
            next_w = math.ceil(current_w * 0.5)
            next_h = math.ceil(current_h * 0.5)
            if next_w < target_w or next_h < target_h:
                break

            resized = cv2.resize(resized, (next_w, next_h), interpolation=interpolation)
            current_h, current_w = resized.shape[:2]
            if (current_w, current_h) == (target_w, target_h):
                return resized

    return cv2.resize(resized, (target_w, target_h), interpolation=interpolation)


def _prepare_roi_image(
    image_bgr: np.ndarray,
    raw_calib: dict[str, Any],
    spec: ROIModelSpec,
    target_size: tuple[int, int],
) -> PreparedROI:
    source_image = np.ascontiguousarray(image_bgr).copy()
    ori_h, ori_w = source_image.shape[:2]
    roi_w = min(int(spec.roi_size[0]), ori_w)
    roi_h = min(int(spec.roi_size[1]), ori_h)
    vp_x, vp_y = _compute_vanishing_point_xy(raw_calib, ori_w, ori_h)
    crop_center_x = vp_x if spec.crop_center_mode == "vxvy" else ori_w / 2.0
    crop_center_y = vp_y
    crop_bounds = compute_centered_roi_bounds(ori_w, ori_h, roi_w, roi_h, crop_center_x, vp_y)
    crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds

    cropped = source_image[crop_y1:crop_y2, crop_x1:crop_x2].copy()
    crop_size = (crop_x2 - crop_x1, crop_y2 - crop_y1)
    if crop_size == target_size:
        resized = cropped
    else:
        resized = _resize_ground3d_image_in_steps(cropped, target_size, interpolation=cv2.INTER_LINEAR)

    calib_for_resize = adjust_calib_for_roi_crop(raw_calib, ori_w, ori_h, crop_bounds)
    final_calib = build_final_resized_calib(
        calib_for_resize["focal_u"],
        calib_for_resize["focal_v"],
        calib_for_resize["cu"],
        calib_for_resize["cv"],
        calib_for_resize["src_w"],
        calib_for_resize["src_h"],
        target_size[0],
        target_size[1],
        spec.virtual_fx,
        distort_coeffs=calib_for_resize["distort_coeffs"],
    )
    return PreparedROI(
        name=spec.name,
        image=resized,
        crop_bounds=crop_bounds,
        calib=final_calib,
        vp_x=float(vp_x),
        vp_y=float(vp_y),
        crop_center_x=float(crop_center_x),
        crop_center_y=float(crop_center_y),
    )


def prepare_roi_image(
    image_bgr: np.ndarray,
    raw_calib: dict[str, Any],
    spec: ROIModelSpec,
    target_size: tuple[int, int],
) -> PreparedROI:
    return _prepare_roi_image(image_bgr, raw_calib, spec, target_size)


def prepare_roi_batch(
    images_bgr: list[np.ndarray],
    raw_calibs: list[dict[str, Any]],
    spec: ROIModelSpec,
    target_size: tuple[int, int],
) -> list[PreparedROI]:
    if len(images_bgr) != len(raw_calibs):
        raise ValueError(
            f"Expected the same number of images and calibrations, got {len(images_bgr)} images and {len(raw_calibs)} calibrations."
        )
    return [prepare_roi_image(image_bgr, raw_calib, spec, target_size) for image_bgr, raw_calib in zip(images_bgr, raw_calibs)]


def _image_to_tensor(image_bgr: np.ndarray, device: torch.device, use_half: bool) -> torch.Tensor:
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    tensor = torch.from_numpy(np.ascontiguousarray(image_rgb.transpose(2, 0, 1))).to(device)
    tensor = tensor.half() if use_half else tensor.float()
    tensor = tensor / 256.0
    return tensor.unsqueeze(0)


def _images_to_batch_tensor(images_bgr: list[np.ndarray], device: torch.device, use_half: bool) -> torch.Tensor:
    if not images_bgr:
        raise ValueError("Expected at least one ROI image to build a batch tensor.")
    image_batch_rgb = [cv2.cvtColor(np.ascontiguousarray(image_bgr), cv2.COLOR_BGR2RGB).transpose(2, 0, 1) for image_bgr in images_bgr]
    tensor = torch.from_numpy(np.ascontiguousarray(np.stack(image_batch_rgb, axis=0))).to(device)
    tensor = tensor.half() if use_half else tensor.float()
    tensor = tensor / 256.0
    return tensor


def _restore_3d_depth_scale(preds_3d: Optional[np.ndarray], depth_scale: float) -> Optional[np.ndarray]:
    if preds_3d is None:
        return None
    preds_3d = preds_3d.copy()
    for channel in (0, 6, 12, 18, 24):
        preds_3d[:, channel] *= depth_scale
    return preds_3d


def _restore_depth_scale(
    preds_3d: np.ndarray,
    preds_3d_fake: Optional[np.ndarray],
    preds_edge: Optional[np.ndarray],
    depth_scale: float,
) -> tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """Restore canonical ROI-space depths back to metric camera depth for both prediction branches."""
    preds_3d = _restore_3d_depth_scale(preds_3d, depth_scale)
    preds_3d_fake = _restore_3d_depth_scale(preds_3d_fake, depth_scale)

    if preds_edge is None:
        return preds_3d, preds_3d_fake, None

    preds_edge = preds_edge.copy()
    preds_edge[:, 2::3] *= depth_scale
    return preds_3d, preds_3d_fake, preds_edge


def _unpack_model_outputs(outputs: Any) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray], np.ndarray, np.ndarray]:
    batch_outputs = _unpack_model_outputs_batch(outputs)
    if not batch_outputs:
        raise RuntimeError("Expected model forward to return at least one inference sample.")
    first = batch_outputs[0]
    return first.detections, first.preds_3d, first.preds_3d_fake, first.preds_edge, first.anchors, first.strides


def _to_numpy_batch(value: Any) -> Optional[np.ndarray]:
    if value is None:
        return None
    if torch.is_tensor(value):
        return value.detach().float().cpu().numpy()
    return np.asarray(value)


def _split_batch_array(values: Optional[np.ndarray], batch_size: int, value_name: str) -> list[Optional[np.ndarray]]:
    if values is None:
        return [None] * batch_size
    if values.ndim == 0:
        raise RuntimeError(f"Expected batched `{value_name}` values, but got a scalar.")
    if values.ndim == 1:
        if batch_size != 1:
            raise RuntimeError(f"Expected batched `{value_name}` values for {batch_size} samples, but got shape {values.shape}.")
        values = values[None, ...]
    if values.shape[0] != batch_size:
        raise RuntimeError(
            f"Batch size mismatch for `{value_name}`: detections batch={batch_size}, {value_name} batch={values.shape[0]}."
        )
    return [values[index] for index in range(batch_size)]


def _unpack_model_outputs_batch(outputs: Any) -> list[RawROIOutputs]:
    if not isinstance(outputs, (list, tuple)) or len(outputs) < 2:
        raise RuntimeError("Expected model forward to return `(detections, raw_preds)` for Detect3D inference.")

    detections = outputs[0]
    raw_preds = outputs[1]
    if not isinstance(raw_preds, dict):
        raise RuntimeError("Detect3D raw predictions are missing from model output.")

    one2one = raw_preds.get("one2one", {})
    preds_3d = one2one.get("preds_3d_selected")
    preds_3d_fake = one2one.get("preds_3d_fake_selected")
    anchors = one2one.get("anchors_selected")
    strides = one2one.get("strides_selected")
    preds_edge = one2one.get("preds_edge_selected")
    if preds_3d is None or anchors is None or strides is None:
        raise RuntimeError("Detect3D metadata (`preds_3d_selected` / `anchors_selected` / `strides_selected`) is missing.")

    if isinstance(detections, (list, tuple)):
        det_batches = [_to_numpy_batch(det) for det in detections]
    else:
        det_np = _to_numpy_batch(detections)
        if det_np is None:
            raise RuntimeError("Detect3D detections are missing from model output.")
        if det_np.ndim == 2:
            det_np = det_np[None, ...]
        det_batches = [det_np[index] for index in range(det_np.shape[0])]

    batch_size = len(det_batches)
    preds_3d_batches = _split_batch_array(_to_numpy_batch(preds_3d), batch_size, "preds_3d_selected")
    preds_3d_fake_batches = _split_batch_array(_to_numpy_batch(preds_3d_fake), batch_size, "preds_3d_fake_selected")
    preds_edge_batches = _split_batch_array(_to_numpy_batch(preds_edge), batch_size, "preds_edge_selected")
    anchors_batches = _split_batch_array(_to_numpy_batch(anchors), batch_size, "anchors_selected")
    strides_batches = _split_batch_array(_to_numpy_batch(strides), batch_size, "strides_selected")

    return [
        RawROIOutputs(
            detections=np.asarray(det_batches[index], dtype=np.float32),
            preds_3d=np.asarray(preds_3d_batches[index], dtype=np.float32),
            preds_3d_fake=None
            if preds_3d_fake_batches[index] is None
            else np.asarray(preds_3d_fake_batches[index], dtype=np.float32),
            preds_edge=None if preds_edge_batches[index] is None else np.asarray(preds_edge_batches[index], dtype=np.float32),
            anchors=np.asarray(anchors_batches[index], dtype=np.float32),
            strides=np.asarray(strides_batches[index], dtype=np.float32),
        )
        for index in range(batch_size)
    ]


def _filter_prediction_rows(
    detections: np.ndarray,
    preds_3d: np.ndarray,
    preds_3d_fake: Optional[np.ndarray],
    preds_edge: Optional[np.ndarray],
    anchors: np.ndarray,
    strides: np.ndarray,
    conf_thres: float,
    max_det: int,
    classes: Optional[set[int]],
) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray, np.ndarray]:
    keep_idx = _compute_prediction_keep_idx(detections, conf_thres=conf_thres, max_det=max_det, classes=classes)
    if detections.size == 0:
        empty = detections[:0]
        return empty, preds_3d[:0], None if preds_3d_fake is None else preds_3d_fake[:0], None if preds_edge is None else preds_edge[:0], anchors[:, :0], strides[:0]
    return (
        detections[keep_idx],
        preds_3d[keep_idx],
        None if preds_3d_fake is None else preds_3d_fake[keep_idx],
        None if preds_edge is None else preds_edge[keep_idx],
        anchors[:, keep_idx],
        strides[keep_idx],
    )


def _compute_prediction_keep_idx(
    detections: np.ndarray,
    conf_thres: float,
    max_det: int,
    classes: Optional[set[int]],
) -> np.ndarray:
    if detections.size == 0:
        return np.zeros((0,), dtype=np.int64)

    keep = detections[:, 4] >= conf_thres
    if classes is not None:
        keep &= np.isin(detections[:, 5].astype(np.int32), np.asarray(sorted(classes), dtype=np.int32))
    return np.flatnonzero(keep)[:max_det]


def iter_batches(items: Iterable[Any], batch_size: int) -> Iterable[list[Any]]:
    resolved_batch_size = max(1, int(batch_size))
    batch: list[Any] = []
    for item in items:
        batch.append(item)
        if len(batch) >= resolved_batch_size:
            yield batch
            batch = []
    if batch:
        yield batch


def run_model_for_prepared_roi_batch(
    bundle: LoadedROIModel,
    prepared_batch: list[PreparedROI],
    device: Optional[torch.device] = None,
    use_half: Optional[bool] = None,
) -> list[RawROIOutputs]:
    if not prepared_batch:
        return []
    resolved_device = device if device is not None else next(bundle.model.parameters()).device
    resolved_use_half = bool(use_half if use_half is not None else next(bundle.model.parameters()).dtype == torch.float16)
    image_tensor = _images_to_batch_tensor([prepared.image for prepared in prepared_batch], device=resolved_device, use_half=resolved_use_half)
    with torch.inference_mode():
        outputs = bundle.model(image_tensor)

    raw_outputs_batch = _unpack_model_outputs_batch(outputs)
    if len(raw_outputs_batch) != len(prepared_batch):
        raise RuntimeError(
            f"Model output batch size mismatch for {bundle.spec.name}: "
            f"prepared={len(prepared_batch)} raw_outputs={len(raw_outputs_batch)}."
        )

    restored_outputs = []
    for prepared, raw_outputs in zip(prepared_batch, raw_outputs_batch):
        preds_3d, preds_3d_fake, preds_edge = _restore_depth_scale(
            raw_outputs.preds_3d,
            raw_outputs.preds_3d_fake,
            raw_outputs.preds_edge,
            float(prepared.calib.get("depth_scale", 1.0)),
        )
        restored_outputs.append(
            RawROIOutputs(
                detections=raw_outputs.detections,
                preds_3d=preds_3d,
                preds_3d_fake=preds_3d_fake,
                preds_edge=preds_edge,
                anchors=raw_outputs.anchors,
                strides=raw_outputs.strides,
            )
        )
    return restored_outputs


def run_model_for_prepared_roi(
    bundle: LoadedROIModel,
    prepared: PreparedROI,
    device: Optional[torch.device] = None,
    use_half: Optional[bool] = None,
) -> RawROIOutputs:
    outputs_batch = run_model_for_prepared_roi_batch(bundle, [prepared], device=device, use_half=use_half)
    if not outputs_batch:
        raise RuntimeError(f"No inference outputs were produced for {bundle.spec.name}.")
    return outputs_batch[0]


def filter_prediction_outputs(
    raw_outputs: RawROIOutputs,
    conf_thres: float,
    max_det: int,
    classes: Optional[set[int]],
) -> RawROIOutputs:
    detections, preds_3d, preds_3d_fake, preds_edge, anchors, strides = _filter_prediction_rows(
        raw_outputs.detections,
        raw_outputs.preds_3d,
        raw_outputs.preds_3d_fake,
        raw_outputs.preds_edge,
        raw_outputs.anchors,
        raw_outputs.strides,
        conf_thres=conf_thres,
        max_det=max_det,
        classes=classes,
    )
    return RawROIOutputs(
        detections=detections,
        preds_3d=preds_3d,
        preds_3d_fake=preds_3d_fake,
        preds_edge=preds_edge,
        anchors=anchors,
        strides=strides,
    )


def filter_prediction_outputs_batch(
    raw_outputs_batch: list[RawROIOutputs],
    conf_thres: float,
    max_det: int,
    classes: Optional[set[int]],
) -> list[RawROIOutputs]:
    return [
        filter_prediction_outputs(
            raw_outputs=raw_outputs,
            conf_thres=conf_thres,
            max_det=max_det,
            classes=classes,
        )
        for raw_outputs in raw_outputs_batch
    ]


def _class_name(names: dict[int, str], cls_id: int) -> str:
    if isinstance(names, dict):
        return str(names.get(cls_id, cls_id))
    return str(cls_id)


def _infer_fake_3d_classes(model: torch.nn.Module, names: dict[int, str]) -> set[int]:
    explicit = set(getattr(model, "fake_3d_classes", set()) or set())
    if explicit:
        return explicit
    if not isinstance(names, dict):
        return set()
    return {int(cls_id) for cls_id, name in names.items() if str(name).endswith("_fake")}


def _draw_2d_boxes(
    image: np.ndarray,
    detections: np.ndarray,
    names: dict[int, str],
) -> np.ndarray:
    drawn = image.copy()
    for det in detections:
        x1, y1, x2, y2 = np.round(det[:4]).astype(np.int32)
        conf = float(det[4])
        cls_id = int(det[5])
        color = colors(cls_id, bgr=True)
        cv2.rectangle(drawn, (x1, y1), (x2, y2), color, 1, cv2.LINE_AA)
        label = f"{_class_name(names, cls_id)} {conf:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
        cv2.rectangle(drawn, (x1, max(0, y1 - th - 6)), (x1 + tw + 2, y1), color, -1)
        cv2.putText(drawn, label, (x1 + 1, max(th + 1, y1 - 3)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
    return drawn


def _annotate_panel_title(image: np.ndarray, title: str) -> np.ndarray:
    drawn = image.copy()
    cv2.putText(drawn, title, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
    return drawn


def _serialize_scalar(value: Any) -> Any:
    if isinstance(value, (np.floating, np.integer)):
        value = value.item()
    if isinstance(value, float) and not math.isfinite(value):
        return None
    return value


def _serialize_array(values: Any) -> Any:
    if values is None:
        return None
    arr = np.asarray(values)
    if arr.ndim == 0:
        return _serialize_scalar(arr.item())
    return [_serialize_array(item) for item in arr.tolist()]


def _wrapped_angle_diff_rad(lhs: Optional[float], rhs: Optional[float]) -> Optional[float]:
    if lhs is None or rhs is None:
        return None
    lhs_val = float(lhs)
    rhs_val = float(rhs)
    if not math.isfinite(lhs_val) or not math.isfinite(rhs_val):
        return None
    return float((lhs_val - rhs_val + math.pi) % (2 * math.pi) - math.pi)


def _decoded_visible_face_types(decoded: Optional[dict[str, Any]]) -> tuple[int, ...]:
    if not decoded:
        return ()
    raw_types = decoded.get("visible_face_types", ())
    if raw_types is None:
        return ()
    face_types = []
    for face_type in raw_types:
        face_type_scalar = _serialize_scalar(face_type)
        if face_type_scalar is None:
            continue
        face_types.append(int(face_type_scalar))
    return tuple(face_types)


def _build_edge_heading_decoded(
    base_decoded: Optional[dict[str, Any]],
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    bbox_xyxy: np.ndarray,
    calib: dict[str, Any],
    img_w: int,
    whole_attrs: Optional[dict[str, Any]],
    edge_yaw_max_lateral_dist_m: float,
    img_h: Optional[int] = None,
    edge_selection: Optional[dict[str, Any]] = None,
) -> tuple[Optional[dict[str, Any]], float, bool]:
    edge_selection = edge_selection or decode_edge_yaw_selection_from_prediction(
        pred_41,
        pred_edge_60,
        anchor_xy,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=img_w,
        img_h=img_h,
        max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
    )
    selected_face_types = tuple(int(face_type) for face_type in edge_selection.get("face_types", ()))
    edge_yaw = float(edge_selection["yaw"])
    if not bool(edge_selection.get("lateral_ok")):
        return None, float(edge_yaw), False

    regressed_dims = pred_41[27:30] if whole_attrs is None else whole_attrs.get("dims", pred_41[27:30])
    edge_box = reconstruct_edge_based_box_from_selection(
        edge_selection,
        box_center_y_m=None,
        regressed_dims=regressed_dims,
        face_regressed_dims_by_type=extract_face_regressed_size_priors_from_prediction(pred_41),
    )
    if edge_box is None:
        return None, float(edge_yaw), False

    decoded = dict(base_decoded or {})
    resolved_face_types = tuple(int(face_type) for face_type in edge_box.get("face_types", ()) or selected_face_types)
    decoded["corners_3d"] = edge_box["corners_3d"]
    decoded["edge_points_2d"] = edge_selection.get("edge_points_2d")
    decoded["edge_points_3d"] = edge_selection.get("edge_points_3d")
    decoded["visible_face_types"] = resolved_face_types
    decoded["face_center_2d"] = None
    decoded["face_color"] = None
    decoded["edge_box_center_3d"] = edge_box["center"]
    decoded["edge_box_dims"] = edge_box["dims"]
    decoded["edge_box_mode"] = edge_box.get("mode")
    return decoded, float(edge_yaw), True


def _build_edge_prediction_artifacts(
    base_decoded: Optional[dict[str, Any]],
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    bbox_xyxy: np.ndarray,
    calib: dict[str, Any],
    img_w: int,
    whole_attrs: Optional[dict[str, Any]],
    edge_yaw_max_lateral_dist_m: float,
    img_h: Optional[int] = None,
) -> dict[str, Any]:
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred_41,
        pred_edge_60,
        anchor_xy,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=img_w,
        img_h=img_h,
        max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
    )
    regressed_dims = pred_41[27:30] if whole_attrs is None else whole_attrs.get("dims", pred_41[27:30])
    edge_box = reconstruct_edge_based_box_from_selection(
        edge_selection,
        box_center_y_m=None,
        regressed_dims=regressed_dims,
        face_regressed_dims_by_type=extract_face_regressed_size_priors_from_prediction(pred_41),
    )
    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=base_decoded,
        pred_41=pred_41,
        pred_edge_60=pred_edge_60,
        anchor_xy=anchor_xy,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=img_w,
        img_h=img_h,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
        edge_selection=edge_selection,
    )
    return {
        "edge_selection": edge_selection,
        "edge_box": edge_box,
        "heading_decoded": heading_decoded,
        "edge_yaw": float(edge_yaw),
        "edge_confident": bool(edge_confident),
    }


def _draw_heading_lines(
    img: np.ndarray,
    center_uv: Any,
    lines: list[tuple[str, tuple[int, int, int]]],
    font_scale: float = 0.45,
    thickness: int = 1,
    line_gap: int = 4,
) -> np.ndarray:
    if center_uv is None:
        return img
    uv = np.asarray(center_uv, dtype=np.float32).reshape(-1)
    if uv.size < 2 or not np.isfinite(uv[:2]).all():
        return img

    cx, cy = int(round(float(uv[0]))), int(round(float(uv[1])))
    font = cv2.FONT_HERSHEY_SIMPLEX
    text_sizes = [cv2.getTextSize(text, font, font_scale, thickness)[0] for text, _ in lines]
    total_h = sum(size[1] for size in text_sizes) + line_gap * max(0, len(lines) - 1)
    y = cy - total_h // 2

    for (text, color), (tw, th) in zip(lines, text_sizes):
        baseline_y = y + th
        cv2.putText(img, text, (cx - tw // 2, baseline_y), font, font_scale, color, thickness, cv2.LINE_AA)
        y = baseline_y + line_gap
    return img


def _extract_center_3d(center_3d: Any) -> Optional[np.ndarray]:
    if center_3d is None:
        return None
    center = np.asarray(center_3d, dtype=np.float32).reshape(-1)
    if center.size < 3 or not np.isfinite(center[:3]).all():
        return None
    return center


def _format_pose_lines(center_3d: Any, yaw_rad: Optional[float], color: tuple[int, int, int]) -> list[tuple[str, tuple[int, int, int]]]:
    if yaw_rad is None:
        return []
    yaw_value = float(yaw_rad)
    if not math.isfinite(yaw_value):
        return []
    center = _extract_center_3d(center_3d)
    if center is None:
        return []
    lateral = abs(float(center[0]))
    depth_z = float(center[2])
    return [
        (f"{lateral:.1f}", color),
        (f"{depth_z:.1f}", color),
        (f"{math.degrees(yaw_value):.1f}", color),
    ]


def _edge_batch_list(edge_points_2d: Any) -> list[np.ndarray]:
    if edge_points_2d is None:
        return []
    arr = np.asarray(edge_points_2d, dtype=np.float32)
    if arr.ndim == 2:
        return [arr]
    if arr.ndim == 3:
        return [arr[index] for index in range(arr.shape[0])]
    return []


def _project_selected_face_edges_from_box(
    corners_3d: Any,
    face_types: tuple[int, ...],
    face_is_partial: tuple[bool, ...],
    calib: dict[str, Any],
    img_w: int,
    img_h: int,
) -> Optional[np.ndarray]:
    if corners_3d is None:
        return None
    projected_edges = []
    for face_type, is_partial in zip(face_types, face_is_partial):
        if bool(is_partial):
            _, points_2d = project_partial_face_bottom_edge(corners_3d, int(face_type), calib, img_w, img_h, num_samples=5)
        else:
            _, points_2d = project_face_bottom_edge(corners_3d, int(face_type), calib, num_samples=5)
        if points_2d is None:
            return None
        projected_edges.append(np.asarray(points_2d, dtype=np.float32))
    if not projected_edges:
        return None
    if len(projected_edges) == 1:
        return projected_edges[0]
    return np.stack(projected_edges, axis=0)


def _selected_edge_residual_stats(reference_edge_points_2d: Any, projected_edge_points_2d: Any) -> dict[str, Any]:
    reference_batches = _edge_batch_list(reference_edge_points_2d)
    projected_batches = _edge_batch_list(projected_edge_points_2d)
    if not reference_batches or len(reference_batches) != len(projected_batches):
        return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}

    per_face_mean = []
    all_residuals = []
    for ref_points, proj_points in zip(reference_batches, projected_batches):
        if ref_points.shape != proj_points.shape or ref_points.ndim != 2 or ref_points.shape[1] != 2:
            return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}
        residuals = np.linalg.norm(np.asarray(ref_points, dtype=np.float32) - np.asarray(proj_points, dtype=np.float32), axis=1)
        per_face_mean.append(float(np.mean(residuals)))
        all_residuals.append(residuals)

    flat = np.concatenate(all_residuals, axis=0) if all_residuals else np.zeros((0,), dtype=np.float32)
    if flat.size == 0:
        return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}
    return {
        "available": True,
        "mean_px": float(np.mean(flat)),
        "max_px": float(np.max(flat)),
        "per_face_mean_px": [float(value) for value in per_face_mean],
    }


def decode_prepared_roi_predictions(
    bundle: LoadedROIModel,
    prepared: PreparedROI,
    filtered_outputs: RawROIOutputs,
    edge_yaw_max_lateral_dist_m: float,
) -> list[dict[str, Any]]:
    detections = filtered_outputs.detections
    preds_3d = filtered_outputs.preds_3d
    preds_3d_fake = filtered_outputs.preds_3d_fake
    preds_edge = filtered_outputs.preds_edge
    anchors = filtered_outputs.anchors
    strides = filtered_outputs.strides

    predictions = []
    for index, det in enumerate(detections):
        pred_edge = None if preds_edge is None else preds_edge[index]
        bbox_xyxy = det[:4].astype(np.float32)
        cls_id = int(det[5])
        is_face_3d_class = cls_id in bundle.face_3d_classes
        pred_3d = preds_3d_fake[index] if preds_3d_fake is not None and cls_id in bundle.fake_3d_classes else preds_3d[index]
        decoded = decode_3d_prediction(
            pred_3d,
            anchors[:, index],
            float(strides[index]),
            prepared.calib,
            prepared.image.shape[1],
            prepared.image.shape[0],
            bundle.face_3d_classes,
            bundle.complete_3d_classes,
            cls_id,
            pred_edge_60=pred_edge,
            bbox_xyxy=bbox_xyxy,
        )
        face_anchor_type = None if decoded is None else decoded.get("visible_face_type")
        attr_face_type = int(face_anchor_type) if is_face_3d_class and face_anchor_type is not None else None
        attrs = extract_3d_attrs_from_prediction(
            pred_3d,
            anchors[:, index],
            float(strides[index]),
            prepared.calib,
            face_type=attr_face_type,
            pred_edge_60=pred_edge,
        )
        if is_face_3d_class:
            edge_artifacts = _build_edge_prediction_artifacts(
                base_decoded=decoded,
                pred_41=pred_3d,
                pred_edge_60=pred_edge,
                anchor_xy=anchors[:, index],
                stride=float(strides[index]),
                bbox_xyxy=bbox_xyxy,
                calib=prepared.calib,
                img_w=prepared.image.shape[1],
                img_h=prepared.image.shape[0],
                whole_attrs=attrs,
                edge_yaw_max_lateral_dist_m=float(edge_yaw_max_lateral_dist_m),
            )
            edge_selection = edge_artifacts["edge_selection"]
            edge_box = edge_artifacts["edge_box"]
            edge_heading_decoded = edge_artifacts["heading_decoded"]
            edge_yaw = edge_artifacts["edge_yaw"]
            edge_confident = edge_artifacts["edge_confident"]
        else:
            edge_selection = None
            edge_box = None
            edge_heading_decoded = None
            edge_yaw = float("nan")
            edge_confident = False
        predictions.append(
            {
                "bbox_xyxy": bbox_xyxy,
                "confidence": float(det[4]),
                "cls_id": cls_id,
                "pred_41": pred_3d,
                "used_fake_3d_head": bool(preds_3d_fake is not None and cls_id in bundle.fake_3d_classes),
                "pred_edge_60": pred_edge,
                "anchor_xy": anchors[:, index],
                "stride": float(strides[index]),
                "attrs": attrs,
                "decoded": decoded,
                "edge_selection": edge_selection,
                "edge_box": edge_box,
                "edge_heading_decoded": edge_heading_decoded,
                "edge_yaw": float(edge_yaw),
                "edge_confident": bool(edge_confident),
            }
        )
    return predictions


def infer_prepared_roi_batch(
    bundle: LoadedROIModel,
    prepared_batch: list[PreparedROI],
    classes: Optional[set[int]],
    edge_yaw_max_lateral_dist_m: float,
    conf_thres: Optional[float] = None,
    raw_outputs_batch: Optional[list[RawROIOutputs]] = None,
    device: Optional[torch.device] = None,
    use_half: Optional[bool] = None,
) -> list[dict[str, Any]]:
    resolved_raw_outputs_batch = (
        raw_outputs_batch
        if raw_outputs_batch is not None
        else run_model_for_prepared_roi_batch(bundle, prepared_batch, device=device, use_half=use_half)
    )
    if len(resolved_raw_outputs_batch) != len(prepared_batch):
        raise RuntimeError(
            f"Prepared ROI batch size mismatch for {bundle.spec.name}: "
            f"prepared={len(prepared_batch)} raw_outputs={len(resolved_raw_outputs_batch)}."
        )

    resolved_conf_thres = float(bundle.spec.conf if conf_thres is None else conf_thres)
    results = []
    for prepared, raw_outputs in zip(prepared_batch, resolved_raw_outputs_batch):
        filtered_outputs = filter_prediction_outputs(
            raw_outputs=raw_outputs,
            conf_thres=resolved_conf_thres,
            max_det=bundle.spec.max_det,
            classes=classes,
        )
        results.append(
            {
                "prepared": prepared,
                "raw_outputs": raw_outputs,
                "filtered_outputs": filtered_outputs,
                "predictions": decode_prepared_roi_predictions(
                    bundle=bundle,
                    prepared=prepared,
                    filtered_outputs=filtered_outputs,
                    edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
                ),
            }
        )
    return results


def _build_visualized_roi_result(
    bundle: LoadedROIModel,
    prepared: PreparedROI,
    filtered_outputs: RawROIOutputs,
    predictions: list[dict[str, Any]],
) -> dict[str, Any]:
    panel_2d = _draw_2d_boxes(prepared.image, filtered_outputs.detections, bundle.names)
    panel_3d = prepared.image.copy()
    panel_3d_edge = prepared.image.copy()
    pred_records = []

    for prediction in predictions:
        bbox_xyxy = np.asarray(prediction["bbox_xyxy"], dtype=np.float32)
        cls_id = int(prediction["cls_id"])
        decoded = prediction.get("decoded")
        edge_selection = prediction.get("edge_selection") or {
            "face_types": (),
            "face_is_partial": (),
            "edge_points_2d": None,
            "lateral_distance_m": None,
            "lateral_ok": False,
            "two_face_eligible": False,
        }
        edge_box = prediction.get("edge_box")
        heading_decoded = prediction.get("edge_heading_decoded")
        pred_attrs = prediction.get("attrs")
        visible_face_types = _decoded_visible_face_types(decoded)

        selected_face_types = tuple(int(face_type) for face_type in edge_selection.get("face_types", ()))
        selected_face_is_partial = tuple(bool(flag) for flag in edge_selection.get("face_is_partial", ()))
        selected_edge_points_2d = edge_selection.get("edge_points_2d")
        direct_box_selected_edges_2d = _project_selected_face_edges_from_box(
            None if decoded is None else decoded.get("corners_3d"),
            selected_face_types,
            selected_face_is_partial,
            prepared.calib,
            prepared.image.shape[1],
            prepared.image.shape[0],
        )
        edge_box_face_types = tuple(int(face_type) for face_type in (edge_box or {}).get("face_types", ()) or selected_face_types)
        edge_box_selected_edges_2d = _project_selected_face_edges_from_box(
            None if heading_decoded is None else heading_decoded.get("corners_3d"),
            edge_box_face_types,
            selected_face_is_partial,
            prepared.calib,
            prepared.image.shape[1],
            prepared.image.shape[0],
        )
        direct_edge_fit = _selected_edge_residual_stats(selected_edge_points_2d, direct_box_selected_edges_2d)
        edge_edge_fit = _selected_edge_residual_stats(selected_edge_points_2d, edge_box_selected_edges_2d)

        reg_yaw = None if pred_attrs is None else _serialize_scalar(pred_attrs["yaw"])
        center_3d = None if pred_attrs is None else pred_attrs.get("center")
        edge_yaw = prediction.get("edge_yaw")
        edge_vs_reg_yaw_rad = _wrapped_angle_diff_rad(edge_yaw, reg_yaw)
        record = {
            "bbox_xyxy": _serialize_array(bbox_xyxy),
            "confidence": float(prediction["confidence"]),
            "cls_id": cls_id,
            "cls_name": _class_name(bundle.names, cls_id),
            "yaw_rad": reg_yaw,
            "edge_yaw_rad": _serialize_scalar(edge_yaw),
            "edge_yaw_confident": bool(prediction.get("edge_confident")),
            "edge_yaw_lateral_distance_m": _serialize_scalar(edge_selection.get("lateral_distance_m")),
            "edge_yaw_lateral_ok": bool(edge_selection.get("lateral_ok")),
            "edge_yaw_two_face_eligible": bool(edge_selection.get("two_face_eligible")),
            "edge_yaw_selected_face_types": _serialize_array(selected_face_types),
            "edge_yaw_selected_face_is_partial": _serialize_array(selected_face_is_partial),
            "edge_vs_reg_yaw_rad": _serialize_scalar(edge_vs_reg_yaw_rad),
            "selected_edge_direct_box_fit_available": bool(direct_edge_fit["available"]),
            "selected_edge_direct_box_fit_mean_px": _serialize_scalar(direct_edge_fit["mean_px"]),
            "selected_edge_direct_box_fit_max_px": _serialize_scalar(direct_edge_fit["max_px"]),
            "selected_edge_direct_box_fit_per_face_mean_px": _serialize_array(direct_edge_fit["per_face_mean_px"]),
            "selected_edge_edgeyaw_box_fit_available": bool(edge_edge_fit["available"]),
            "selected_edge_edgeyaw_box_fit_mean_px": _serialize_scalar(edge_edge_fit["mean_px"]),
            "selected_edge_edgeyaw_box_fit_max_px": _serialize_scalar(edge_edge_fit["max_px"]),
            "selected_edge_edgeyaw_box_fit_per_face_mean_px": _serialize_array(edge_edge_fit["per_face_mean_px"]),
            "selected_edge_fit_gain_px": _serialize_scalar(
                None
                if direct_edge_fit["mean_px"] is None or edge_edge_fit["mean_px"] is None
                else float(direct_edge_fit["mean_px"]) - float(edge_edge_fit["mean_px"])
            ),
            "edge_box_center_3d": None if edge_box is None else _serialize_array(edge_box["center"]),
            "edge_box_dims": None if edge_box is None else _serialize_array(edge_box["dims"]),
            "edge_box_length_m": _serialize_scalar(None if edge_box is None else edge_box.get("side_length_m")),
            "edge_box_width_m": _serialize_scalar(None if edge_box is None else edge_box.get("width_m")),
            "edge_box_mode": None if edge_box is None else str(edge_box.get("mode")),
            "edge_box_length_source": None if edge_box is None else edge_box.get("length_source"),
            "edge_box_width_source": None if edge_box is None else edge_box.get("width_source"),
            "center_uv": None if pred_attrs is None else _serialize_array(pred_attrs["uv"]),
            "center_3d": _serialize_array(center_3d),
            "visible_face_type": None if decoded is None else _serialize_scalar(decoded.get("visible_face_type")),
            "visible_face_count": len(visible_face_types),
            "visible_face_types": None if decoded is None else _serialize_array(visible_face_types),
            "crop_bounds": list(prepared.crop_bounds),
        }
        pred_records.append(record)

        center_uv = record.get("center_uv")
        if decoded is not None and decoded.get("corners_3d") is not None:
            draw_3d_box(
                panel_3d,
                decoded["corners_3d"],
                prepared.calib,
                decoded.get("face_center_2d"),
                decoded.get("face_color"),
                edge_points_2d=decoded.get("edge_points_2d"),
                edge_color=(0, 255, 0),
                thickness=1,
            )
            reg_lines = _format_pose_lines(record.get("center_3d"), record.get("yaw_rad"), (0, 0, 0))
            if reg_lines:
                _draw_heading_lines(panel_3d, center_uv, reg_lines, font_scale=0.35, thickness=1, line_gap=2)

        if bool(record.get("edge_yaw_confident")) and heading_decoded is not None and heading_decoded.get("corners_3d") is not None:
            draw_3d_box(
                panel_3d_edge,
                heading_decoded["corners_3d"],
                prepared.calib,
                heading_decoded.get("face_center_2d"),
                heading_decoded.get("face_color"),
                edge_points_2d=heading_decoded.get("edge_points_2d"),
                edge_color=(0, 255, 0),
                thickness=1,
            )
            edge_lines = _format_pose_lines(record.get("edge_box_center_3d") or record.get("center_3d"), record.get("edge_yaw_rad"), (0, 0, 0))
            if edge_lines:
                _draw_heading_lines(panel_3d_edge, center_uv, edge_lines, font_scale=0.35, thickness=1, line_gap=2)

    return {
        "prepared": prepared,
        "detections": filtered_outputs.detections,
        "predictions": pred_records,
        "panel_2d": _annotate_panel_title(panel_2d, f"{bundle.spec.name} 2D"),
        "panel_3d": _annotate_panel_title(panel_3d, f"{bundle.spec.name} 3D"),
        "panel_3d_edge": _annotate_panel_title(panel_3d_edge, f"{bundle.spec.name} 3D EdgeRecon (1+ face)"),
    }


def _run_single_roi(
    bundle: LoadedROIModel,
    prepared: PreparedROI,
    device: torch.device,
    use_half: bool,
    classes: Optional[set[int]],
    edge_yaw_max_lateral_dist_m: float,
) -> dict[str, Any]:
    inference_result = infer_prepared_roi_batch(
        bundle=bundle,
        prepared_batch=[prepared],
        classes=classes,
        edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
        raw_outputs_batch=None,
        device=device,
        use_half=use_half,
    )[0]
    return _build_visualized_roi_result(
        bundle=bundle,
        prepared=inference_result["prepared"],
        filtered_outputs=inference_result["filtered_outputs"],
        predictions=inference_result["predictions"],
    )


def _assemble_grid(roi_results: list[dict[str, Any]]) -> np.ndarray:
    row_images = []
    for row_idx, roi_result in enumerate(roi_results):
        panels = [roi_result["panel_2d"], roi_result["panel_3d"], roi_result["panel_3d_edge"]]
        row_h = max(panel.shape[0] for panel in panels)
        row_w = sum(panel.shape[1] for panel in panels)
        row_canvas = np.full((row_h, row_w, 3), GRID_BG_COLOR, dtype=np.uint8)
        x0 = 0
        for panel in panels:
            panel_h, panel_w = panel.shape[:2]
            y0 = max(0, (row_h - panel_h) // 2)
            row_canvas[y0 : y0 + panel_h, x0 : x0 + panel_w] = panel
            x0 += panel_w
        row_images.append(row_canvas)

    grid_h = sum(row.shape[0] for row in row_images)
    grid_w = max(row.shape[1] for row in row_images)
    grid = np.full((grid_h, grid_w, 3), GRID_BG_COLOR, dtype=np.uint8)
    y0 = 0
    for row in row_images:
        row_h, row_w = row.shape[:2]
        x0 = max(0, (grid_w - row_w) // 2)
        grid[y0 : y0 + row_h, x0 : x0 + row_w] = row
        y0 += row_h
    return grid


def iter_case_images(images_dir: str | Path, glob_pattern: str, max_images: int = 0) -> Iterable[Path]:
    images_dir = Path(images_dir)
    image_paths = [path for path in sorted(images_dir.glob(glob_pattern)) if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES]
    if max_images > 0:
        image_paths = image_paths[:max_images]
    return image_paths


def iter_loaded_case_images(
    images_dir: str | Path,
    glob_pattern: str,
    max_images: int = 0,
) -> Iterable[tuple[str, np.ndarray]]:
    for image_path in iter_case_images(images_dir, glob_pattern=glob_pattern, max_images=max_images):
        image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
        if image_bgr is None:
            raise FileNotFoundError(f"Failed to read image: {image_path}")
        yield image_path.name, image_bgr


def build_inference_context(
    roi_specs: list[ROIModelSpec],
    device: str = "0",
    half: bool = False,
    classes: Optional[list[int]] = None,
    edge_yaw_max_lateral_dist_m: float = DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    inference_batch_size: int = 1,
) -> InferenceContext:
    requested_device = device or ""
    try:
        selected_device = select_device(requested_device)
    except ValueError as exc:
        # Allow batch jobs to keep running on CPU-only environments even if the
        # default CLI device is set to `0`.
        if requested_device and requested_device.lower() != "cpu" and not torch.cuda.is_available():
            print(
                f"[WARN] Requested device={requested_device!r} but CUDA is unavailable; "
                "falling back to CPU."
            )
            selected_device = select_device("cpu")
        else:
            raise exc
    use_half = bool(half and selected_device.type != "cpu")
    roi_models = []
    for spec in roi_specs:
        model, _ = load_checkpoint(spec.model_path, device=selected_device, fuse=False)
        if use_half:
            model = model.half()
        names = getattr(model, "names", {}) or {}
        fake_3d_classes = _infer_fake_3d_classes(model, names)
        roi_models.append(
            LoadedROIModel(
                spec=spec,
                model=model.eval(),
                names=names,
                face_3d_classes=set(getattr(model, "face_3d_classes", set())),
                complete_3d_classes=set(getattr(model, "complete_3d_classes", set())),
                fake_3d_classes=fake_3d_classes,
                imgsz=_infer_model_imgsz(model, spec.imgsz),
            )
        )
    return InferenceContext(
        device=selected_device,
        use_half=use_half,
        classes=None if classes is None else set(int(cls_id) for cls_id in classes),
        roi_models=roi_models,
        edge_yaw_max_lateral_dist_m=float(edge_yaw_max_lateral_dist_m),
        inference_batch_size=max(1, int(inference_batch_size)),
    )


def _resolve_case_inputs(case_dir: str = "", images_dir: str = "", calib_file: str = "") -> tuple[Path, Path, Optional[Path]]:
    if case_dir:
        case_path = Path(case_dir).resolve()
        return case_path / "images", case_path / "calib" / "L2_calib" / "camera4.json", case_path
    if not images_dir or not calib_file:
        raise ValueError("Either --case-dir or both --images-dir and --calib-file are required.")
    return Path(images_dir).resolve(), Path(calib_file).resolve(), None


def _resolve_visualization_output_dir(output_dir: str | Path, case_name: str) -> Path:
    output_dir = Path(output_dir)
    if output_dir == DEFAULT_VISUALIZATION_ROOT or output_dir == Path(str(DEFAULT_VISUALIZATION_ROOT)):
        output_dir = output_dir / case_name
    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def create_predictions_payload(
    context: InferenceContext,
    case_name: str,
    source_info: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    payload = {
        "case_name": case_name,
        "edge_yaw_max_lateral_dist_m": context.edge_yaw_max_lateral_dist_m,
        "frames": [],
    }
    if source_info:
        payload.update(source_info)
    return payload


def _run_inference_on_frame_batch(
    context: InferenceContext,
    raw_calib: dict[str, Any],
    frame_batch: list[tuple[int, str, np.ndarray]],
    output_dir: Path,
) -> list[dict[str, Any]]:
    if not frame_batch:
        return []

    frame_payloads = [
        {
            "frame_index": int(frame_index),
            "frame_name": frame_name,
            "rois": {},
        }
        for frame_index, frame_name, _image_bgr in frame_batch
    ]
    roi_results_by_frame: list[list[dict[str, Any]]] = [[] for _ in frame_batch]
    batch_images = [image_bgr for _frame_index, _frame_name, image_bgr in frame_batch]
    raw_calibs = [raw_calib for _ in frame_batch]

    for bundle in context.roi_models:
        prepared_batch = prepare_roi_batch(batch_images, raw_calibs, bundle.spec, bundle.imgsz)
        inference_results = infer_prepared_roi_batch(
            bundle=bundle,
            prepared_batch=prepared_batch,
            classes=context.classes,
            edge_yaw_max_lateral_dist_m=context.edge_yaw_max_lateral_dist_m,
            device=context.device,
            use_half=context.use_half,
        )
        for frame_offset, inference_result in enumerate(inference_results):
            prepared = inference_result["prepared"]
            roi_result = _build_visualized_roi_result(
                bundle=bundle,
                prepared=prepared,
                filtered_outputs=inference_result["filtered_outputs"],
                predictions=inference_result["predictions"],
            )
            roi_results_by_frame[frame_offset].append(roi_result)
            frame_payloads[frame_offset]["rois"][bundle.spec.name.lower()] = {
                "crop_bounds": list(prepared.crop_bounds),
                "vp_x": prepared.vp_x,
                "vp_y": prepared.vp_y,
                "crop_center_x": prepared.crop_center_x,
                "crop_center_y": prepared.crop_center_y,
                "edge_yaw_max_lateral_dist_m": context.edge_yaw_max_lateral_dist_m,
                "calib": {key: _serialize_scalar(value) for key, value in prepared.calib.items()},
                "predictions": roi_result["predictions"],
            }

    for frame_offset, ((frame_index, frame_name, _image_bgr), roi_results) in enumerate(zip(frame_batch, roi_results_by_frame)):
        grid = _assemble_grid(roi_results)
        frame_stem = Path(frame_name).stem or f"frame_{frame_index:06d}"
        grid_path = output_dir / f"{frame_stem}.jpg"
        if not cv2.imwrite(str(grid_path), grid):
            raise IOError(f"Failed to write visualization image: {grid_path}")
        frame_payloads[frame_offset]["visualization"] = str(grid_path)

    return frame_payloads


def _run_inference_on_frame(
    context: InferenceContext,
    raw_calib: dict[str, Any],
    image_bgr: np.ndarray,
    frame_name: str,
    frame_index: int,
    output_dir: Path,
) -> dict[str, Any]:
    frame_payloads = _run_inference_on_frame_batch(
        context=context,
        raw_calib=raw_calib,
        frame_batch=[(frame_index, frame_name, image_bgr)],
        output_dir=output_dir,
    )
    if not frame_payloads:
        raise RuntimeError(f"Failed to run inference for frame {frame_name}.")
    return frame_payloads[0]


def append_image_stream_inference(
    context: InferenceContext,
    frames: Iterable[tuple[str, np.ndarray]],
    raw_calib: dict[str, Any],
    output_dir: str | Path,
    predictions_payload: dict[str, Any],
    frame_index_offset: int = 0,
    frame_name_prefix: str = "",
) -> int:
    output_dir = Path(output_dir)
    num_frames = 0
    batch_size = max(1, int(getattr(context, "inference_batch_size", 1)))
    pending_frames: list[tuple[int, str, np.ndarray]] = []
    for local_frame_index, (frame_name, image_bgr) in enumerate(frames):
        effective_frame_name = f"{frame_name_prefix}_{frame_name}" if frame_name_prefix else frame_name
        pending_frames.append((frame_index_offset + local_frame_index, effective_frame_name, image_bgr))
        if len(pending_frames) < batch_size:
            continue
        frame_payloads = _run_inference_on_frame_batch(
            context=context,
            raw_calib=raw_calib,
            frame_batch=pending_frames,
            output_dir=output_dir,
        )
        predictions_payload["frames"].extend(frame_payloads)
        num_frames += len(frame_payloads)
        pending_frames = []

    if pending_frames:
        frame_payloads = _run_inference_on_frame_batch(
            context=context,
            raw_calib=raw_calib,
            frame_batch=pending_frames,
            output_dir=output_dir,
        )
        predictions_payload["frames"].extend(frame_payloads)
        num_frames += len(frame_payloads)

    return num_frames


def run_image_stream_inference(
    context: InferenceContext,
    frames: Iterable[tuple[str, np.ndarray]],
    raw_calib: dict[str, Any],
    case_name: str,
    output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
    source_info: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    output_dir = _resolve_visualization_output_dir(output_dir, case_name)
    predictions_payload = create_predictions_payload(
        context=context,
        case_name=case_name,
        source_info=source_info,
    )
    num_frames = append_image_stream_inference(
        context=context,
        frames=frames,
        raw_calib=raw_calib,
        output_dir=output_dir,
        predictions_payload=predictions_payload,
        frame_index_offset=0,
        frame_name_prefix="",
    )

    if num_frames == 0:
        raise FileNotFoundError(f"No frames available for case {case_name}")

    predictions_path = output_dir / "predictions.json"
    with predictions_path.open("w", encoding="utf-8") as file:
        json.dump(predictions_payload, file, indent=2, ensure_ascii=False)

    return {
        "case_name": case_name,
        "output_dir": str(output_dir),
        "predictions_path": str(predictions_path),
        "num_frames": num_frames,
    }


def run_case_inference(
    context: InferenceContext,
    case_dir: str = "",
    images_dir: str = "",
    calib_file: str = "",
    output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
    glob_pattern: str = "*.png",
    max_images: int = 0,
) -> dict[str, Any]:
    images_dir_path, calib_path, resolved_case_dir = _resolve_case_inputs(case_dir=case_dir, images_dir=images_dir, calib_file=calib_file)
    raw_calib = load_camera4_calib(calib_path)

    case_name = resolved_case_dir.name if resolved_case_dir is not None else images_dir_path.name
    return run_image_stream_inference(
        context=context,
        frames=iter_loaded_case_images(images_dir_path, glob_pattern=glob_pattern, max_images=max_images),
        raw_calib=raw_calib,
        case_name=case_name,
        output_dir=output_dir,
        source_info={
            "images_dir": str(images_dir_path),
            "calib_file": str(calib_path),
        },
    )


def iter_video_case_images(
    video_path: str | Path,
    frame_index_payload: Optional[dict[str, Any]] = None,
    video_stride: int = 1,
    max_images: int = 0,
) -> Iterable[tuple[str, np.ndarray]]:
    for frame_index, image_bgr, frame_name, _frame_info in iter_video_case_frames(
        video_path,
        frame_index_payload=frame_index_payload,
        frame_stride=video_stride,
        max_frames=max_images,
    ):
        yield f"{int(frame_index):06d}_{Path(frame_name).name}", image_bgr


def run_video_case_inference(
    context: InferenceContext,
    video_case_dir: str | Path,
    output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
    max_images: int = 0,
    video_stride: int = 1,
) -> dict[str, Any]:
    case_dir, video_path, calib_path = resolve_video_case_paths(video_case_dir)
    raw_calib = load_camera4_calib(calib_path)
    frame_index_payload = read_video_frame_index(video_path)

    return run_image_stream_inference(
        context=context,
        frames=iter_video_case_images(
            video_path=video_path,
            frame_index_payload=frame_index_payload,
            video_stride=video_stride,
            max_images=max_images,
        ),
        raw_calib=raw_calib,
        case_name=case_dir.name,
        output_dir=output_dir,
        source_info={
            "images_dir": str(video_path.parent),
            "calib_file": str(calib_path),
            "input_mode": "video_case",
            "video_case_dir": str(case_dir),
            "video_path": str(video_path),
        },
    )


def build_roi_specs_from_args(args: Any) -> list[ROIModelSpec]:
    from tools.pdcl_inference.run_batch_two_roi_infer import build_roi_specs_from_args as _build_roi_specs_from_args

    return _build_roi_specs_from_args(args)


def add_two_roi_inference_args(parser: Any, include_output_dir: bool = True) -> None:
    from tools.pdcl_inference.run_batch_two_roi_infer import add_two_roi_inference_args as _add_two_roi_inference_args

    _add_two_roi_inference_args(parser, include_output_dir=include_output_dir)
