from __future__ import annotations

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

import cv2
import numpy as np
import yaml

try:
    from .two_roi_types import (
        Decoded3DPrediction,
        Prediction3DAttrs,
        RawCameraCalib,
        ROICropCalib,
        ResizedCalib,
    )
    from .edge_yaw_utils import (
        EDGE_YAW_MAX_LATERAL_DIST_M,
        decode_edge_yaw_selection_from_prediction,
        extract_face_regressed_size_priors_from_prediction,
        reconstruct_edge_based_box_from_selection,
    )
except ImportError:
    from two_roi_types import (
        Decoded3DPrediction,
        Prediction3DAttrs,
        RawCameraCalib,
        ROICropCalib,
        ResizedCalib,
    )
    from edge_yaw_utils import (
        EDGE_YAW_MAX_LATERAL_DIST_M,
        decode_edge_yaw_selection_from_prediction,
        extract_face_regressed_size_priors_from_prediction,
        reconstruct_edge_based_box_from_selection,
    )

IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
ROI_NAMES = ("roi0", "roi1")
GRID_BG_COLOR = (24, 24, 24)
DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M = EDGE_YAW_MAX_LATERAL_DIST_M
DEFAULT_INPUT_IMGSZ = (768, 352)
DEFAULT_STRIDES = (8, 16, 32)
LABEL_FONT = cv2.FONT_HERSHEY_SIMPLEX
LABEL_FONT_SCALE = 0.45
LABEL_THICKNESS = 1
PANEL_TITLE_FONT_SCALE = 0.7
PANEL_TITLE_THICKNESS = 2
HEADING_LINE_GAP_PX = 4
BOX_LABEL_LINE_GAP_PX = 2
BOX_LABEL_PAD_X = 2
BOX_LABEL_PAD_Y = 2
DEFAULT_PALETTE = (
    (56, 56, 255),
    (151, 157, 255),
    (31, 112, 255),
    (29, 178, 255),
    (49, 210, 207),
    (10, 249, 72),
    (23, 204, 146),
    (134, 219, 61),
    (52, 147, 26),
    (187, 212, 0),
    (168, 153, 44),
    (255, 194, 0),
    (147, 69, 52),
    (255, 115, 100),
    (236, 24, 0),
    (255, 56, 132),
    (133, 0, 82),
    (255, 56, 203),
    (200, 149, 255),
    (199, 55, 255),
)
DEFAULT_DATASET_CONFIG = {
    "class_map": {
        "car": 0,
        "suv": 1,
        "pickup": 2,
        "medium_car": 3,
        "van": 4,
        "bus": 5,
        "truck": 6,
        "tanker": 6,
        "large_truck": 6,
        "construction_vehicle": 6,
        "special_vehicle": 7,
        "unknown": 8,
        "pedestrian": 9,
        "bicyclist": 10,
        "motorcyclist": 10,
        "bicycle": 11,
        "motorcycle": 11,
        "tricycle": 12,
        "tricyclist": 12,
        "traffic_sign": 13,
        "wheel": 14,
        "plate": 15,
        "face": 16,
        "car_fake": 17,
        "bicyclist_fake": 18,
        "pedestrian_fake": 19,
        "car_carrier": 6,
        "platform_truck": 6,
    },
    "face_3d_classes": [0, 1, 2, 3, 4, 5, 6, 7, 8, 17],
    "complete_3d_classes": [9, 10, 11, 12, 18, 19],
    "fake_3d_classes": [17, 18, 19],
    "roi_configs": {
        "roi0": {"roi": [1920, 880], "virtual_fx": 537.0, "crop_center_mode": "cxvy"},
        "roi1": {"roi": [768, 352], "virtual_fx": 537.0, "crop_center_mode": "vxvy"},
    },
    "norm_scales_3d": {
        "z3d_scale": 24.415,
        "z3d_offset": 39.937,
        "size_scale": 1.945,
        "size_offset": 3.780,
        "yaw_scale": 1.5707963,
    },
}
DEFAULT_ROI_CONFIGS = DEFAULT_DATASET_CONFIG["roi_configs"]


@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 PreparedROI:
    name: str
    image: np.ndarray
    crop_bounds: tuple[int, int, int, int]
    calib: ResizedCalib
    vp_x: float
    vp_y: float
    crop_center_x: float
    crop_center_y: float


@dataclass(frozen=True)
class RuntimeMetadata:
    class_map: dict[str, int]
    names: dict[int, str]
    face_3d_classes: set[int]
    complete_3d_classes: set[int]
    fake_3d_classes: set[int]
    norm_scales_3d: dict[str, float]


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 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 _load_dataset_config(data_config_path: str = "") -> dict[str, Any]:
    loaded = load_yaml_if_present(data_config_path)
    if not loaded:
        return DEFAULT_DATASET_CONFIG
    return _deep_merge_dict(DEFAULT_DATASET_CONFIG, loaded)


def _deep_merge_dict(base: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]:
    merged = dict(base)
    for key, value in update.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge_dict(merged[key], value)
        else:
            merged[key] = value
    return merged


def _invert_class_map(class_map: dict[str, int]) -> dict[int, str]:
    names = {}
    for name, class_id in class_map.items():
        names.setdefault(int(class_id), str(name))
    return dict(sorted(names.items()))


def load_runtime_metadata(data_config_path: str = "") -> RuntimeMetadata:
    config = _load_dataset_config(data_config_path)
    class_map = {str(name): int(class_id) for name, class_id in (config.get("class_map", {}) or {}).items()}
    return RuntimeMetadata(
        class_map=class_map,
        names=_invert_class_map(class_map),
        face_3d_classes=set(int(v) for v in config.get("face_3d_classes", [])),
        complete_3d_classes=set(int(v) for v in config.get("complete_3d_classes", [])),
        fake_3d_classes=set(int(v) for v in config.get("fake_3d_classes", [])),
        norm_scales_3d={k: float(v) for k, v in config.get("norm_scales_3d", {}).items()},
    )


def infer_model_imgsz(_spec_model_path: Any, override: Optional[tuple[int, int]]) -> tuple[int, int]:
    return override or DEFAULT_INPUT_IMGSZ


def compute_centered_roi_bounds(ori_w: int, ori_h: int, roi_w: int, roi_h: int, center_x: float, center_y: float) -> tuple[int, int, int, int]:
    crop_x1 = int(max(0, min(center_x - roi_w / 2, ori_w - roi_w)))
    crop_y1 = int(max(0, min(center_y - roi_h / 2, ori_h - roi_h)))
    return crop_x1, crop_y1, crop_x1 + roi_w, crop_y1 + roi_h


def adjust_calib_for_roi_crop(
    raw_calib: RawCameraCalib,
    ori_w: int,
    ori_h: int,
    crop_bounds: tuple[int, int, int, int],
) -> ROICropCalib:
    crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds
    cu = raw_calib.get("cu", ori_w / 2) if raw_calib else ori_w / 2
    cv = raw_calib.get("cv", ori_h / 2) if raw_calib else ori_h / 2
    focal_u = raw_calib.get("focal_u", ori_w) if raw_calib else ori_w
    focal_v = raw_calib.get("focal_v", ori_h) if raw_calib else ori_h
    distort_coeffs = raw_calib.get("distort_coeffs", []) if raw_calib else []
    return {
        "focal_u": focal_u,
        "focal_v": focal_v,
        "cu": cu - crop_x1,
        "cv": cv - crop_y1,
        "src_w": crop_x2 - crop_x1,
        "src_h": crop_y2 - crop_y1,
        "distort_coeffs": distort_coeffs,
    }


def build_final_resized_calib(
    focal_u: float,
    focal_v: float,
    cu: float,
    cv: float,
    src_w: int,
    src_h: int,
    target_w: int,
    target_h: int,
    virtual_fx: float,
    distort_coeffs: Optional[list[float]] = None,
) -> ResizedCalib:
    scale_x = target_w / src_w
    scale_y = target_h / src_h
    fx_final = focal_u * scale_x
    return {
        "fx": fx_final,
        "fy": focal_v * scale_y,
        "cx": cu * scale_x,
        "cy": cv * scale_y,
        "distort_coeffs": distort_coeffs if distort_coeffs is not None else [],
        "depth_scale": fx_final / virtual_fx,
    }


def load_camera4_calib(calib_path: str | Path) -> RawCameraCalib:
    calib_path = Path(calib_path)
    with calib_path.open("r", encoding="utf-8") as file:
        payload = json.load(file)
    if "intrinsics" in payload:
        intrinsics = payload.get("intrinsics", {}).get("camera4.json", {}) or {}
        extrinsics = payload.get("extrinsics", {}).get("camera4.json", {}) or {}
        rpy = list(extrinsics.get("rpy", [0.0, 0.0, 0.0]) or [0.0, 0.0, 0.0])
        while len(rpy) < 3:
            rpy.append(0.0)
        pos = extrinsics.get("pos", [0.0, 0.0, 0.0]) or [0.0, 0.0, 0.0]
        if isinstance(pos, dict):
            pos = [pos.get("x", 0.0), pos.get("y", 0.0), pos.get("z", 0.0)]
        pos = list(pos)
        while len(pos) < 3:
            pos.append(0.0)
        return {
            "focal_u": float(intrinsics["focal_u"]),
            "focal_v": float(intrinsics["focal_v"]),
            "cu": float(intrinsics["cu"]),
            "cv": float(intrinsics["cv"]),
            "roll": float(rpy[0]),
            "pitch": float(rpy[1]),
            "yaw": float(rpy[2]),
            "pos": [float(value) for value in pos[:3]],
            "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": "degrees",
        }
    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 _to_radians(angle: float, unit: str) -> float:
    return math.radians(float(angle)) if unit == "degrees" else float(angle)


def _compute_vanishing_point_xy(raw_calib: RawCameraCalib, ori_w: int, ori_h: int) -> tuple[float, float]:
    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)
    return float(cx + fx * math.tan(yaw)), float(cy - fy * math.tan(pitch))


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 _build_roi_calib(
    raw_calib: RawCameraCalib,
    ori_w: int,
    ori_h: int,
    crop_bounds: tuple[int, int, int, int],
    target_size: tuple[int, int],
    virtual_fx: float,
) -> ResizedCalib:
    calib_for_resize = adjust_calib_for_roi_crop(raw_calib, ori_w, ori_h, crop_bounds)
    return 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],
        virtual_fx,
        distort_coeffs=calib_for_resize["distort_coeffs"],
    )


def prepare_roi_image(
    image_bgr: np.ndarray,
    raw_calib: RawCameraCalib,
    spec: ROIModelSpec,
    target_size: tuple[int, int],
    crop_center_y_offset_px: float = 0.0,
) -> 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 + float(crop_center_y_offset_px)
    crop_bounds = compute_centered_roi_bounds(ori_w, ori_h, roi_w, roi_h, crop_center_x, crop_center_y)
    crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds
    cropped = source_image[crop_y1:crop_y2, crop_x1:crop_x2].copy()
    resized = _resize_ground3d_image_in_steps(cropped, target_size, interpolation=cv2.INTER_LINEAR)
    final_calib = _build_roi_calib(raw_calib, ori_w, ori_h, crop_bounds, target_size, spec.virtual_fx)
    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 resolve_case_calib_path(case_dir: str | Path) -> Path:
    case_path = Path(case_dir).resolve()
    calib_candidates = [
        case_path / "calib" / "L2_calib" / "camera4.json",
        case_path / "calib" / "camera4.json",
    ]
    calib_path = next((path for path in calib_candidates if path.is_file()), None)
    if calib_path is None:
        joined = ", ".join(str(path) for path in calib_candidates)
        raise FileNotFoundError(f"camera4.json not found under {case_path}. Checked: {joined}")
    return calib_path


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", resolve_case_calib_path(case_path), 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 iter_case_images(images_dir: str | Path, glob_pattern: str, max_images: int = 0) -> Iterable[Path]:
    image_paths = [path for path in sorted(Path(images_dir).glob(glob_pattern)) if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES]
    return image_paths[:max_images] if max_images > 0 else image_paths


def add_two_roi_inference_args(parser: argparse.ArgumentParser, default_roi_model_paths: dict[str, str]) -> None:
    for roi_name in ROI_NAMES:
        config = DEFAULT_ROI_CONFIGS[roi_name]
        parser.add_argument(f"--{roi_name}-model", type=str, default=default_roi_model_paths[roi_name])
        parser.add_argument(f"--{roi_name}-data", type=str, default="")
        parser.add_argument(f"--{roi_name}-roi", nargs=2, type=int, default=tuple(config["roi"]), metavar=("W", "H"))
        parser.add_argument(f"--{roi_name}-crop-center-mode", type=str, default=str(config["crop_center_mode"]), choices=("cxvy", "vxvy"))
        parser.add_argument(
            f"--{roi_name}-crop-center-y-offset-px",
            type=float,
            default=0.0,
            help=(
                "Optional constant vertical offset (in original-image pixels) applied to this ROI crop center. "
                "Positive values move the crop downward."
            ),
        )
        parser.add_argument(
            f"--{roi_name}-crop-center-y-offset-map",
            type=str,
            default="",
            help=(
                "Optional JSON file that defines per-frame vertical crop-center offsets for this ROI. "
                "Supports frame_id, frame_index, and frame_name keyed offsets."
            ),
        )
        parser.add_argument(f"--{roi_name}-virtual-fx", type=float, default=float(config["virtual_fx"]))
        parser.add_argument(f"--{roi_name}-imgsz", nargs=2, type=int, default=None, metavar=("W", "H"))
        parser.add_argument(f"--{roi_name}-conf", type=float, default=0.25)
        parser.add_argument(f"--{roi_name}-max-det", type=int, default=300)
    parser.add_argument("--classes", nargs="*", type=int, default=None)
    parser.add_argument(
        "--vis-classes",
        nargs="*",
        type=int,
        default=None,
        help="Filter only visualization output to selected class ids; JSON predictions remain unfiltered.",
    )
    parser.add_argument(
        "--vis-class-names",
        nargs="*",
        type=str,
        default=None,
        help="Filter only visualization output to selected class names or aliases such as pedestrian or rider.",
    )
    parser.add_argument("--glob", type=str, default="*.png")
    parser.add_argument("--max-images", type=int, default=0)
    parser.add_argument("--edge-yaw-max-lateral-dist", type=float, default=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M)
    parser.add_argument("--data-config", type=str, default="")


def _resolve_cli_or_config_value(cli_value: Any, default_value: Any, config_value: Any) -> Any:
    return config_value if cli_value == default_value else cli_value


def _load_roi_data_config(args: Any, roi_name: str) -> dict[str, Any]:
    roi_specific_path = str(getattr(args, f"{roi_name}_data", ""))
    shared_path = str(getattr(args, "data_config", ""))
    return load_yaml_if_present(roi_specific_path) or load_yaml_if_present(shared_path)


def _resolve_roi_config(args: Any, roi_name: str) -> dict[str, Any]:
    return (_load_roi_data_config(args, roi_name).get("roi_configs", {}) or {}).get(roi_name, {})


def _build_roi_spec_from_args(args: Any, roi_name: str) -> ROIModelSpec:
    default_config = DEFAULT_ROI_CONFIGS[roi_name]
    resolved_config = _resolve_roi_config(args, roi_name)
    roi_arg = tuple(int(value) for value in getattr(args, f"{roi_name}_roi"))
    crop_mode_arg = str(getattr(args, f"{roi_name}_crop_center_mode"))
    virtual_fx_arg = float(getattr(args, f"{roi_name}_virtual_fx"))
    roi_value = tuple(int(value) for value in resolved_config.get("roi", roi_arg))
    crop_mode_value = str(resolved_config.get("crop_center_mode", crop_mode_arg))
    virtual_fx_value = float(resolved_config.get("virtual_fx", virtual_fx_arg))
    imgsz_override = getattr(args, f"{roi_name}_imgsz")
    return ROIModelSpec(
        name=roi_name.upper(),
        model_path=getattr(args, f"{roi_name}_model"),
        roi_size=_resolve_cli_or_config_value(roi_arg, tuple(default_config["roi"]), roi_value),
        crop_center_mode=_resolve_cli_or_config_value(
            crop_mode_arg,
            str(default_config["crop_center_mode"]),
            crop_mode_value,
        ),
        virtual_fx=_resolve_cli_or_config_value(
            virtual_fx_arg,
            float(default_config["virtual_fx"]),
            virtual_fx_value,
        ),
        imgsz=coerce_imgsz(imgsz_override) if imgsz_override else None,
        conf=float(getattr(args, f"{roi_name}_conf")),
        max_det=int(getattr(args, f"{roi_name}_max_det")),
    )


def build_roi_specs_from_args(args: Any) -> list[ROIModelSpec]:
    return [_build_roi_spec_from_args(args, roi_name) for roi_name in ROI_NAMES]


def build_anchor_cache(imgsz_wh: tuple[int, int], strides: tuple[int, ...] = DEFAULT_STRIDES) -> tuple[np.ndarray, np.ndarray]:
    width, height = int(imgsz_wh[0]), int(imgsz_wh[1])
    anchor_points, stride_tensor = [], []
    for stride in strides:
        h, w = height // int(stride), width // int(stride)
        sy = np.arange(h, dtype=np.float32) + 0.5
        sx = np.arange(w, dtype=np.float32) + 0.5
        grid_y, grid_x = np.meshgrid(sy, sx, indexing="ij")
        anchor_points.append(np.stack((grid_x, grid_y), axis=0).reshape(2, -1))
        stride_tensor.append(np.full((1, h * w), float(stride), dtype=np.float32))
    return np.concatenate(anchor_points, axis=1), np.concatenate(stride_tensor, axis=1)


def _softmax(x: np.ndarray, axis: int) -> np.ndarray:
    x = x - np.max(x, axis=axis, keepdims=True)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)


def _sigmoid(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=np.float32)
    positive = x >= 0
    negative = ~positive
    out = np.empty_like(x, dtype=np.float32)
    out[positive] = 1.0 / (1.0 + np.exp(-x[positive]))
    exp_x = np.exp(x[negative])
    out[negative] = exp_x / (1.0 + exp_x)
    return out


def decode_boxes_xyxy(raw_boxes: np.ndarray, anchors: np.ndarray, strides: np.ndarray) -> np.ndarray:
    batch, channels, anchors_count = raw_boxes.shape
    reg_max = max(1, channels // 4)
    if reg_max == 1:
        dist = raw_boxes.reshape(batch, 4, anchors_count).astype(np.float32, copy=False)
    else:
        logits = raw_boxes.reshape(batch, 4, reg_max, anchors_count)
        probs = _softmax(logits, axis=2)
        bins = np.arange(reg_max, dtype=np.float32).reshape(1, 1, reg_max, 1)
        dist = np.sum(probs * bins, axis=2)
    lt = dist[:, 0:2, :]
    rb = dist[:, 2:4, :]
    boxes = np.concatenate((anchors[None, :, :] - lt, anchors[None, :, :] + rb), axis=1)
    return boxes * strides[None, :, :]


def _topk_descending(values: np.ndarray, k: int, axis: int = -1) -> tuple[np.ndarray, np.ndarray]:
    sorted_idx = np.argsort(values, axis=axis)
    topk_idx = np.take(sorted_idx, indices=range(values.shape[axis] - k, values.shape[axis]), axis=axis)
    topk_idx = np.flip(topk_idx, axis=axis)
    topk_values = np.take_along_axis(values, topk_idx, axis=axis)
    return topk_values, topk_idx


def _get_topk_index_numpy(scores: np.ndarray, max_det: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    batch_size, anchors_count, num_classes = scores.shape
    k = min(int(max_det), anchors_count)

    anchor_max_scores = scores.max(axis=-1)
    _, anchor_indices = _topk_descending(anchor_max_scores, k, axis=1)
    gathered_scores = np.take_along_axis(scores, anchor_indices[..., None], axis=1)
    flat_scores = gathered_scores.reshape(batch_size, -1)
    topk_scores, flat_indices = _topk_descending(flat_scores, k, axis=1)

    flat_anchor_index = np.broadcast_to(anchor_indices[..., None], (batch_size, k, num_classes)).reshape(batch_size, -1)
    selected_anchor_indices = np.take_along_axis(flat_anchor_index, flat_indices, axis=1)

    class_lookup = np.broadcast_to(
        np.arange(num_classes, dtype=np.int64)[None, None, :],
        (batch_size, k, num_classes),
    ).reshape(batch_size, -1)
    class_indices = np.take_along_axis(class_lookup, flat_indices, axis=1).astype(np.float32)
    return topk_scores[..., None], class_indices[..., None], selected_anchor_indices[..., None]


def postprocess_detections(decoded_boxes_xyxy: np.ndarray, raw_scores: np.ndarray, max_det: int) -> tuple[np.ndarray, np.ndarray]:
    scores = _sigmoid(raw_scores)
    boxes = np.transpose(decoded_boxes_xyxy, (0, 2, 1))
    scores = np.transpose(scores, (0, 2, 1))
    topk_scores, class_indices, selected_anchor_indices = _get_topk_index_numpy(scores, max_det)
    selected_boxes = np.take_along_axis(boxes, np.repeat(selected_anchor_indices, 4, axis=2), axis=1)
    detections = np.concatenate((selected_boxes, topk_scores, class_indices), axis=2).astype(np.float32)
    return detections, selected_anchor_indices[..., 0].astype(np.int64)


def _select_topk_last_axis(values: np.ndarray, topk_indices: np.ndarray) -> np.ndarray:
    return np.take_along_axis(values, topk_indices[:, None, :], axis=2)


def denorm_preds_3d(raw_preds_3d: np.ndarray, norm_scales_3d: dict[str, float]) -> np.ndarray:
    z_scale = float(norm_scales_3d.get("z3d_scale", 1.0))
    z_offset = float(norm_scales_3d.get("z3d_offset", 0.0))
    s_scale = float(norm_scales_3d.get("size_scale", 1.0))
    s_offset = float(norm_scales_3d.get("size_offset", 0.0))
    p = raw_preds_3d.astype(np.float32, copy=True)
    p[:, (0, 6, 12, 18, 24), :] = p[:, (0, 6, 12, 18, 24), :] * z_scale + z_offset
    for ch_start in (1, 7, 13, 19, 25):
        p[:, ch_start : ch_start + 2, :] = _sigmoid(p[:, ch_start : ch_start + 2, :]) * 16.0 - 8.0
    for ch_start in (3, 9, 15, 21):
        p[:, ch_start : ch_start + 2, :] = p[:, ch_start : ch_start + 2, :] * s_scale + s_offset
    p[:, 27:30, :] = p[:, 27:30, :] * s_scale + s_offset
    p[:, 34:38, :] = np.tanh(p[:, 34:38, :])
    return p


def denorm_preds_edge(raw_preds_edge: np.ndarray, norm_scales_3d: dict[str, float]) -> np.ndarray:
    z_scale = float(norm_scales_3d.get("z3d_scale", 1.0))
    z_offset = float(norm_scales_3d.get("z3d_offset", 0.0))
    p = raw_preds_edge.astype(np.float32, copy=True)
    for point_start in range(0, p.shape[1], 3):
        p[:, point_start : point_start + 2, :] = _sigmoid(p[:, point_start : point_start + 2, :]) * 16.0 - 8.0
        p[:, point_start + 2, :] = p[:, point_start + 2, :] * z_scale + z_offset
    return p


def select_topk_branch(branch: np.ndarray, topk_indices: np.ndarray) -> np.ndarray:
    return _select_topk_last_axis(branch, topk_indices).transpose(0, 2, 1)


def select_topk_3d_metadata(preds_3d: np.ndarray, anchors: np.ndarray, strides: np.ndarray, topk_indices: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    preds_3d_selected = _select_topk_last_axis(preds_3d, topk_indices).transpose(0, 2, 1)
    anchors_selected = _select_topk_last_axis(anchors[None, :, :], topk_indices)
    strides_selected = np.take_along_axis(strides, topk_indices, axis=1)
    return preds_3d_selected, anchors_selected, strides_selected


def restore_depth_scale(preds_3d: np.ndarray, preds_edge: Optional[np.ndarray], depth_scale: float) -> tuple[np.ndarray, Optional[np.ndarray]]:
    preds_3d = preds_3d.copy()
    preds_3d[:, (0, 6, 12, 18, 24)] *= depth_scale
    if preds_edge is None:
        return preds_3d, None
    preds_edge = preds_edge.copy()
    preds_edge[:, 2::3] *= depth_scale
    return preds_3d, preds_edge


def filter_prediction_rows(
    detections: np.ndarray,
    preds_3d: 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_edge is None else preds_edge[:0], anchors[:, :0], strides[:0]
    return (
        detections[keep_idx],
        preds_3d[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 _palette_color(index: int) -> tuple[int, int, int]:
    color = DEFAULT_PALETTE[int(index) % len(DEFAULT_PALETTE)]
    return int(color[0]), int(color[1]), int(color[2])


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


def _normalize_box_label_lines(label: Any) -> list[str]:
    if label is None:
        return []
    if isinstance(label, str):
        values = [label]
    elif isinstance(label, Iterable):
        values = list(label)
    else:
        values = [str(label)]

    lines = []
    for value in values:
        text = str(value).strip()
        if text:
            lines.append(text)
    return lines


def _draw_box_label(img: np.ndarray, label: Any, x1: int, y1: int, color: tuple[int, int, int]) -> None:
    lines = _normalize_box_label_lines(label)
    if not lines:
        return

    x1 = max(0, int(x1))
    y1 = max(0, int(y1))
    text_sizes = [cv2.getTextSize(line, LABEL_FONT, LABEL_FONT_SCALE, LABEL_THICKNESS)[0] for line in lines]
    text_w = max(size[0] for size in text_sizes)
    total_text_h = sum(size[1] for size in text_sizes)
    box_h = total_text_h + BOX_LABEL_LINE_GAP_PX * max(0, len(lines) - 1) + BOX_LABEL_PAD_Y * 2
    box_w = text_w + BOX_LABEL_PAD_X * 2
    top_y = max(0, y1 - box_h)
    bottom_y = min(img.shape[0] - 1, top_y + box_h)
    right_x = min(img.shape[1] - 1, x1 + box_w)
    cv2.rectangle(img, (x1, top_y), (right_x, bottom_y), color, -1)

    y = top_y + BOX_LABEL_PAD_Y
    for line, (_, text_h) in zip(lines, text_sizes):
        baseline_y = y + text_h
        cv2.putText(
            img,
            line,
            (x1 + BOX_LABEL_PAD_X, baseline_y),
            LABEL_FONT,
            LABEL_FONT_SCALE,
            (255, 255, 255),
            LABEL_THICKNESS,
            cv2.LINE_AA,
        )
        y = baseline_y + BOX_LABEL_LINE_GAP_PX


def draw_2d_boxes(
    image: np.ndarray,
    detections: np.ndarray,
    names: dict[int, str],
    extra_label_lines: Optional[Iterable[Any]] = None,
) -> np.ndarray:
    drawn = image.copy()
    extra_labels = list(extra_label_lines) if extra_label_lines is not None else []
    for index, det in enumerate(detections):
        x1, y1, x2, y2 = np.round(det[:4]).astype(np.int32)
        conf = float(det[4])
        cls_id = int(det[5])
        color = _palette_color(cls_id)
        cv2.rectangle(drawn, (x1, y1), (x2, y2), color, 1, cv2.LINE_AA)
        label_lines = [f"{class_name(names, cls_id)} {conf:.2f}"]
        if index < len(extra_labels):
            label_lines.extend(_normalize_box_label_lines(extra_labels[index]))
        _draw_box_label(drawn, label_lines, x1, y1, color)
    return drawn


def annotate_panel_title(image: np.ndarray, title: str) -> np.ndarray:
    drawn = image.copy()
    cv2.putText(
        drawn,
        title,
        (10, 24),
        LABEL_FONT,
        PANEL_TITLE_FONT_SCALE,
        (0, 255, 255),
        PANEL_TITLE_THICKNESS,
        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 edge_yaw_lateral_distance_m(center_3d: Any) -> Optional[float]:
    if center_3d is None:
        return None
    center = np.asarray(center_3d, dtype=np.float32).reshape(-1)
    if center.size < 1 or not np.isfinite(center[0]):
        return None
    return float(abs(center[0]))


def _has_valid_decoded_corners(decoded: Optional[Decoded3DPrediction]) -> bool:
    return decoded is not None and decoded.corners_3d is not None


def build_edge_heading_decoded(
    base_decoded: Optional[Decoded3DPrediction],
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    bbox_xyxy: np.ndarray,
    calib: ResizedCalib,
    img_w: int,
    whole_attrs: Optional[Prediction3DAttrs],
    edge_yaw_max_lateral_dist_m: float,
    img_h: Optional[int] = None,
    edge_selection: Optional[dict[str, Any]] = None,
) -> tuple[Optional[Decoded3DPrediction], 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.dims
    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

    resolved_face_types = tuple(int(face_type) for face_type in edge_box.get("face_types", ()) or selected_face_types)
    edge_points_2d = edge_selection.get("edge_points_2d")
    edge_points_3d = edge_selection.get("edge_points_3d")
    if base_decoded is None:
        decoded = Decoded3DPrediction(
            corners_3d=np.asarray(edge_box["corners_3d"], dtype=np.float32),
            face_center_2d=None,
            face_color=None,
            visible_face_type=None,
            visible_face_types=resolved_face_types,
            edge_points_2d=None if edge_points_2d is None else np.asarray(edge_points_2d, dtype=np.float32),
            edge_points_3d=None if edge_points_3d is None else np.asarray(edge_points_3d, dtype=np.float32),
            cls_id=-1,
        )
    else:
        decoded = replace(
            base_decoded,
            corners_3d=np.asarray(edge_box["corners_3d"], dtype=np.float32),
            face_center_2d=None,
            face_color=None,
            visible_face_types=resolved_face_types,
            edge_points_2d=None if edge_points_2d is None else np.asarray(edge_points_2d, dtype=np.float32),
            edge_points_3d=None if edge_points_3d is None else np.asarray(edge_points_3d, dtype=np.float32),
        )
    return decoded, float(edge_yaw), True


def build_edge_prediction_artifacts(
    base_decoded: Optional[Decoded3DPrediction],
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    bbox_xyxy: np.ndarray,
    calib: ResizedCalib,
    img_w: int,
    whole_attrs: Optional[Prediction3DAttrs],
    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.dims
    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 _coerce_center_uv(center_uv: Any) -> Optional[tuple[int, int]]:
    if center_uv is None:
        return None
    uv = np.asarray(center_uv, dtype=np.float32).reshape(-1)
    if uv.size < 2 or not np.isfinite(uv[:2]).all():
        return None
    return int(round(float(uv[0]))), int(round(float(uv[1])))


def _heading_text_sizes(lines: list[tuple[str, tuple[int, int, int]]]) -> list[tuple[int, int]]:
    return [cv2.getTextSize(text, LABEL_FONT, LABEL_FONT_SCALE, LABEL_THICKNESS)[0] for text, _ in lines]


def draw_heading_lines(img: np.ndarray, center_uv: Any, lines: list[tuple[str, tuple[int, int, int]]]) -> np.ndarray:
    center = _coerce_center_uv(center_uv)
    if center is None:
        return img

    cx, cy = center
    text_sizes = _heading_text_sizes(lines)
    y = cy - (sum(size[1] for size in text_sizes) + HEADING_LINE_GAP_PX * max(0, len(lines) - 1)) // 2
    for (text, color), (tw, th) in zip(lines, text_sizes):
        baseline_y = y + th
        cv2.putText(img, text, (cx - tw // 2, baseline_y), LABEL_FONT, LABEL_FONT_SCALE, color, LABEL_THICKNESS, cv2.LINE_AA)
        y = baseline_y + HEADING_LINE_GAP_PX
    return img
