from __future__ import annotations

from dataclasses import dataclass, field
import logging
from pathlib import Path
import re
from typing import Any

import cv2
import numpy as np

try:
    import torch
except ImportError:
    torch = None


LOGGER = logging.getLogger("two_roi_exported_onnx_infer")


DEFAULT_VEHICLE_CLASS_IDS = {0, 13}
DEFAULT_PEDESTRIAN_CLASS_IDS = {1, 3}
DEFAULT_TSR_CLASS_IDS = {6}

VEHICLE_CLASS_NAME_ALIASES = {
    "vehicle",
    "car",
    "car_fake",
    "suv",
    "pickup",
    "medium_car",
    "van",
    "bus",
    "truck",
    "tanker",
    "large_truck",
    "construction_vehicle",
    "special_vehicle",
    "unknown",
    "tricycle",
    "tricyclist",
}
PEDESTRIAN_CLASS_NAME_ALIASES = {
    "pedestrian",
    "person",
    "bicyclist",
    "motorcyclist",
    "cyclist",
    "rider",
}
TSR_CLASS_NAME_ALIASES = {
    "traffic_sign",
    "trafficsign",
    "tsr",
    "small_sign",
    "roadsign",
    "road_sign",
    "sign",
}

VEHICLE_TASK = "vehicle"
PEDESTRIAN_TASK = "pedestrian"
TSR_TASK = "tsr"

PED_POSE_INDEX = ["bg", "w", "x", "a", "d"]
PED_TYPE_INDEX = ["bg", "ped", "ped_cut", "ride"]

ATTRIBUTE_TEXT_COLORS = {
    VEHICLE_TASK: (0, 255, 255),
    PEDESTRIAN_TASK: (0, 255, 0),
    TSR_TASK: (255, 0, 255),
}
ATTRIBUTE_LABEL_BG_COLOR = (36, 36, 36)
ATTRIBUTE_LABEL_BG_ALPHA = 0.72
ATTRIBUTE_LABEL_PAD_X = 2
ATTRIBUTE_LABEL_PAD_Y = 2
ATTRIBUTE_LABEL_GAP_PX = 6


@dataclass
class AttributeModelBundle:
    """Container for optional attribute models and their runtime config."""

    device: torch.device
    car_model: torch.jit.ScriptModule | None = None
    ped_model: torch.jit.ScriptModule | None = None
    tsr_model: torch.jit.ScriptModule | None = None
    tsr_labels: list[str] | None = None
    vehicle_min_conf: float | None = None
    vehicle_min_side: float = 40.0
    pedestrian_min_conf: float | None = None
    pedestrian_min_height: float = 40.0
    tsr_min_conf: float | None = None
    tsr_min_side: float = 16.0
    vehicle_class_ids: set[int] = field(default_factory=set)
    pedestrian_class_ids: set[int] = field(default_factory=set)
    tsr_class_ids: set[int] = field(default_factory=set)
    roi_conf_by_id: dict[int, float] = field(default_factory=dict)


def _require_torch() -> None:
    if torch is None:
        raise ImportError("torch is required for attribute-model inference but is not installed.")


def _normalize_device(device: str | None) -> torch.device:
    _require_torch()
    if device is None:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device)


def _load_torchscript_model(weights: str | Path | None, device: torch.device) -> torch.jit.ScriptModule | None:
    if not weights:
        return None

    _require_torch()
    weights_path = Path(weights)
    if not weights_path.exists():
        raise FileNotFoundError(f"Attribute model not found: {weights_path}")

    LOGGER.info(f"Loading attribute TorchScript model: {weights_path}")
    model = torch.jit.load(str(weights_path), map_location=device)
    model.eval()
    return model


def _load_tsr_labels(label_file: str | Path | None) -> list[str] | None:
    if not label_file:
        return None

    label_path = Path(label_file)
    if not label_path.exists():
        LOGGER.warning(f"TSR label file not found: {label_path}; falling back to raw class ids")
        return None

    with label_path.open("r", encoding="utf-8") as file:
        return [line.strip() for line in file if line.strip()]


def _normalize_class_name(name: str) -> str:
    normalized = re.sub(r"[^a-z0-9]+", "_", str(name or "").strip().lower())
    return normalized.strip("_")


def _is_tsr_name(normalized_name: str) -> bool:
    return (
        normalized_name in TSR_CLASS_NAME_ALIASES
        or normalized_name.startswith("tsr")
        or ("traffic" in normalized_name and "sign" in normalized_name)
    )


def _is_pedestrian_name(normalized_name: str) -> bool:
    return (
        normalized_name in PEDESTRIAN_CLASS_NAME_ALIASES
        or "pedestrian" in normalized_name
        or "bicyclist" in normalized_name
        or "motorcyclist" in normalized_name
        or "cyclist" in normalized_name
        or "rider" in normalized_name
    )


def _is_vehicle_name(normalized_name: str) -> bool:
    return (
        normalized_name in VEHICLE_CLASS_NAME_ALIASES
        or normalized_name == "unknown"
        or "vehicle" in normalized_name
        or "truck" in normalized_name
        or "bus" in normalized_name
        or "tanker" in normalized_name
        or "pickup" in normalized_name
        or normalized_name in {"car", "suv", "van", "tricycle"}
    )


def _resolve_task_class_ids(
    class_names: dict[int, str] | None,
) -> tuple[set[int], set[int], set[int]]:
    if not class_names:
        return (
            set(DEFAULT_VEHICLE_CLASS_IDS),
            set(DEFAULT_PEDESTRIAN_CLASS_IDS),
            set(DEFAULT_TSR_CLASS_IDS),
        )

    vehicle_class_ids: set[int] = set()
    pedestrian_class_ids: set[int] = set()
    tsr_class_ids: set[int] = set()
    for cls_id, cls_name in class_names.items():
        normalized_name = _normalize_class_name(cls_name)
        if _is_tsr_name(normalized_name):
            tsr_class_ids.add(int(cls_id))
            continue
        if _is_pedestrian_name(normalized_name):
            pedestrian_class_ids.add(int(cls_id))
            continue
        if _is_vehicle_name(normalized_name):
            vehicle_class_ids.add(int(cls_id))

    return vehicle_class_ids, pedestrian_class_ids, tsr_class_ids


def _format_class_group(class_ids: set[int], class_names: dict[int, str] | None) -> str:
    if not class_ids:
        return "[]"
    if not class_names:
        return str(sorted(int(cls_id) for cls_id in class_ids))
    return str(
        [
            f"{int(cls_id)}:{class_names.get(int(cls_id), str(int(cls_id)))}"
            for cls_id in sorted(class_ids)
        ]
    )


def load_attribute_models(
    *,
    enable_attr: bool = False,
    class_names: dict[int, str] | None = None,
    roi_conf_by_id: dict[int, float] | None = None,
    car_weights: str | Path | None = None,
    ped_weights: str | Path | None = None,
    tsr_weights: str | Path | None = None,
    tsr_label_file: str | Path | None = None,
    device: str | None = None,
    vehicle_min_conf: float | None = None,
    vehicle_min_side: float = 40.0,
    pedestrian_min_conf: float | None = None,
    pedestrian_min_height: float = 40.0,
    tsr_min_conf: float | None = None,
    tsr_min_side: float = 16.0,
) -> AttributeModelBundle | None:
    """Load optional attribute models for merged inference."""
    if not enable_attr:
        return None

    attr_device = _normalize_device(device)
    bundle = AttributeModelBundle(
        device=attr_device,
        car_model=_load_torchscript_model(car_weights, attr_device),
        ped_model=_load_torchscript_model(ped_weights, attr_device),
        tsr_model=_load_torchscript_model(tsr_weights, attr_device),
        tsr_labels=_load_tsr_labels(tsr_label_file),
        vehicle_min_conf=vehicle_min_conf,
        vehicle_min_side=vehicle_min_side,
        pedestrian_min_conf=pedestrian_min_conf,
        pedestrian_min_height=pedestrian_min_height,
        tsr_min_conf=tsr_min_conf,
        tsr_min_side=tsr_min_side,
        roi_conf_by_id={} if roi_conf_by_id is None else {int(k): float(v) for k, v in roi_conf_by_id.items()},
    )
    (
        bundle.vehicle_class_ids,
        bundle.pedestrian_class_ids,
        bundle.tsr_class_ids,
    ) = _resolve_task_class_ids(class_names)
    LOGGER.info(
        "Attribute models enabled on %s (car=%s, ped=%s, tsr=%s)",
        bundle.device,
        bundle.car_model is not None,
        bundle.ped_model is not None,
        bundle.tsr_model is not None,
    )
    LOGGER.info(
        "Resolved attribute class groups from runtime names: vehicle=%s pedestrian=%s tsr=%s",
        _format_class_group(bundle.vehicle_class_ids, class_names),
        _format_class_group(bundle.pedestrian_class_ids, class_names),
        _format_class_group(bundle.tsr_class_ids, class_names),
    )
    if bundle.roi_conf_by_id:
        LOGGER.info("Attribute confidence gating follows ROI detection thresholds: %s", bundle.roi_conf_by_id)
    return bundle


def _box_width_height(det: dict[str, Any]) -> tuple[float, float]:
    bbox = det["bbox"]
    return float(bbox[2] - bbox[0]), float(bbox[3] - bbox[1])


def _parse_roi_id(roi_id: Any) -> int | None:
    try:
        return int(roi_id)
    except (TypeError, ValueError):
        return None


def _effective_min_conf_for_detection(
    det: dict[str, Any],
    task_min_conf: float | None,
    attribute_models: AttributeModelBundle,
) -> float:
    roi_id = _parse_roi_id(det.get("roi_id"))
    roi_min_conf: float | None = None
    if roi_id is not None and roi_id in attribute_models.roi_conf_by_id:
        roi_min_conf = float(attribute_models.roi_conf_by_id[roi_id])
    elif attribute_models.roi_conf_by_id:
        roi_min_conf = min(float(value) for value in attribute_models.roi_conf_by_id.values())

    thresholds = [float(value) for value in (roi_min_conf, task_min_conf) if value is not None]
    return max(thresholds) if thresholds else 0.0


def _filter_vehicle_detections(
    detections: list[dict[str, Any]],
    attribute_models: AttributeModelBundle,
) -> list[dict[str, Any]]:
    filtered = []
    for det in detections:
        width, height = _box_width_height(det)
        side = max(width, height)
        if det["score"] < _effective_min_conf_for_detection(det, attribute_models.vehicle_min_conf, attribute_models):
            continue
        if side < attribute_models.vehicle_min_side:
            continue
        filtered.append(det)
    return filtered


def _filter_pedestrian_detections(
    detections: list[dict[str, Any]],
    attribute_models: AttributeModelBundle,
) -> list[dict[str, Any]]:
    filtered = []
    for det in detections:
        _, height = _box_width_height(det)
        if det["score"] < _effective_min_conf_for_detection(det, attribute_models.pedestrian_min_conf, attribute_models):
            continue
        if height < attribute_models.pedestrian_min_height:
            continue
        filtered.append(det)
    return filtered


def _filter_tsr_detections(
    detections: list[dict[str, Any]],
    attribute_models: AttributeModelBundle,
) -> list[dict[str, Any]]:
    filtered = []
    for det in detections:
        width, height = _box_width_height(det)
        side = max(width, height)
        if det["score"] < _effective_min_conf_for_detection(det, attribute_models.tsr_min_conf, attribute_models):
            continue
        if side < attribute_models.tsr_min_side:
            continue
        filtered.append(det)
    return filtered


def _crop_region(img: np.ndarray, bbox: list[float]) -> np.ndarray:
    return img[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2]), :]


def _get_img_patch_for_car(
    ori_img: np.ndarray,
    box: list[float],
    side_param: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    rect = np.zeros(4)
    orirect = np.zeros(4)
    bbox = [int(round(float(v))) for v in box]
    w = bbox[2] - bbox[0] + 1
    h = bbox[3] - bbox[1] + 1
    center_x = bbox[0] + w / 2
    center_y = bbox[1] + h / 2

    side_w = int(round(side_param * w))
    side_h = int(round(side_param * h))

    rect[0] = max(0, int(center_x - side_w / 2))
    rect[1] = max(0, int(center_y - side_h / 2))
    rect[2] = min(int(center_x + side_w / 2), ori_img.shape[1])
    rect[3] = min(int(center_y + side_h / 2), ori_img.shape[0])
    orirect[0] = int((side_w - w) / 2) + min(0, int(center_x - side_w / 2))
    orirect[1] = int((side_h - h) / 2) + min(0, int(center_y - side_h / 2))
    orirect[2] = min(int(orirect[0] + w), ori_img.shape[1])
    orirect[3] = min(int(orirect[1] + h), ori_img.shape[0])

    crop_img = _crop_region(ori_img, rect)
    return crop_img, rect, orirect


def _get_img_patch_for_ped(
    ori_img: np.ndarray,
    box: list[float],
    side_param: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    rect = np.zeros(4)
    orirect = np.zeros(4)
    bbox = [int(round(float(v))) for v in box]
    w = bbox[2] - bbox[0] + 1
    h = bbox[3] - bbox[1] + 1
    center_x = bbox[0] + w / 2
    center_y = bbox[1] + h / 2

    side_w = int(round(1.1 * w))
    side_h = int(round(1.2 * h))

    rect[0] = max(0, int(center_x - side_w / 2))
    rect[1] = max(0, int(center_y - side_h / 2))
    rect[2] = min(int(center_x + side_w / 2), ori_img.shape[1])
    rect[3] = min(int(center_y + side_w / 2), ori_img.shape[0])
    orirect[0] = int((side_w - w) / 2) + min(0, int(center_x - side_w / 2))
    orirect[1] = int((side_h - h) / 2) + min(0, int(center_y - side_h / 2))
    orirect[2] = min(int(orirect[0] + w), ori_img.shape[1])
    orirect[3] = min(int(orirect[1] + h), ori_img.shape[0])

    crop_img = _crop_region(ori_img, rect)
    return crop_img, rect, orirect


def _to_merge_img(patches_all: list[np.ndarray], param: list[int]) -> np.ndarray:
    c, h, w, size_reshape = param
    w_h = w * h
    len_img = len(patches_all)
    len_img_new = int(np.ceil(len_img / float(w_h)) * w_h)
    img_all = np.zeros((len_img_new, c, size_reshape, size_reshape))

    for count, patch in enumerate(patches_all):
        img = cv2.resize(patch, (size_reshape, size_reshape), interpolation=cv2.INTER_LINEAR)
        img = img.transpose(2, 0, 1)
        img = np.reshape(img, (c, size_reshape, size_reshape))
        img_all[count, :, :, :] = img

    images = (
        img_all.transpose(1, 0, 2, 3)
        .reshape(img_all.shape[1], -1, size_reshape * h, size_reshape)
        .transpose(0, 1, 3, 2)
        .reshape(img_all.shape[1], -1, size_reshape * w, size_reshape * h)
        .transpose(1, 0, 3, 2)
    )
    return images


def _convert_point_to_rect(dots: np.ndarray, cls_type: int, width: int, height: int) -> list[float]:
    len_dot = int(len(dots) / 2)
    x = np.zeros(shape=len_dot, dtype=np.float32)
    y = np.zeros(shape=len_dot, dtype=np.float32)
    for i in range(len_dot):
        x[i] = float(dots[i * 2])
        y[i] = float(dots[i * 2 + 1])

    res_y0 = (float(y[0]) + float(y[1])) / 2.0
    if cls_type == 1:
        res_x0 = (float(x[0]) + float(x[5])) / 2.0
        res_x1 = (float(x[1]) + float(x[2])) / 2.0
        res_y1 = (float(y[3]) + float(y[4])) / 2.0
    elif cls_type in (2, 3, 11):
        res_x0 = float(x[5])
        res_x1 = float(x[2])
        res_y1 = (float(y[3]) + float(y[4])) / 2.0
    elif 4 <= cls_type <= 10:
        res_x0 = (float(x[0]) + float(x[5])) / 2.0
        res_x1 = (float(x[1]) + float(x[2])) / 2.0
        res_y1 = (float(y[3]) + float(y[4])) / 2.0
    elif cls_type >= 12:
        res_x0 = float(x[0])
        res_x1 = float(x[1])
        res_y1 = (float(y[2]) + float(y[3])) / 2.0
    else:
        raise ValueError(f"Unsupported cls_type for rect conversion: {cls_type}")

    res_x0 = min(max(0, res_x0), width)
    res_x1 = min(max(0, res_x1), width)
    res_y0 = min(max(0, res_y0), height)
    res_y1 = min(max(0, res_y1), height)
    return [res_x0, res_y0, res_x1, res_y1]


def _unscale_and_uncenter(feature: np.ndarray, x0: float, y0: float, width: float, height: float) -> np.ndarray:
    feature = np.array(feature, dtype=np.float32).copy()
    for i in range(len(feature)):
        after_scale = (feature[i] * 2.5 + 40) / 80.0
        if i % 2 == 0:
            feature[i] = after_scale * width + x0
        else:
            feature[i] = after_scale * height + y0
    return feature


def _vehicle_output_analysis(pred: list[torch.Tensor], param: list[int]) -> tuple[np.ndarray, ...]:
    _, h, _, _ = param
    feature = pred[0].detach().cpu().numpy()
    out_shape = feature.shape[-1] // h
    feature = (
        feature.transpose(1, 0, 3, 2)
        .reshape(feature.shape[1], -1, out_shape, out_shape * h)
        .transpose(0, 1, 3, 2)
        .reshape(feature.shape[1], -1, out_shape, out_shape)
        .transpose(1, 0, 2, 3)
    )
    feature = feature[:, :, 0, 0]
    car_type_near_cls = feature[:, 0:24]
    car_type_far_cls = feature[:, 24:48]
    tail_near_cls = feature[:, 48:52]
    tail_far_cls = feature[:, 96:100]
    car_head_type_cls = feature[:, 144:152]
    car_head_ye_double_cls = feature[:, 152:160]
    out_pose_type_cls = feature[:, 168:172]
    fake_car_cls = feature[:, 224:226]
    out_shelter_cls = feature[:, 234:240]
    out_reg1 = feature[:, 240:288]
    out_reg2 = feature[:, 288:336]
    out_reg3 = feature[:, 336:368]
    return (
        car_type_near_cls,
        car_type_far_cls,
        tail_near_cls,
        tail_far_cls,
        car_head_type_cls,
        car_head_ye_double_cls,
        out_pose_type_cls,
        fake_car_cls,
        out_shelter_cls,
        out_reg1,
        out_reg2,
        out_reg3,
    )


def _ped_output_analysis(pred: list[torch.Tensor], param: list[int]) -> tuple[np.ndarray, ...]:
    _, h, _, _ = param
    feature = pred[0].detach().cpu().numpy()
    out_shape = feature.shape[-1] // h
    feature = (
        feature.transpose(1, 0, 3, 2)
        .reshape(feature.shape[1], -1, out_shape, out_shape * h)
        .transpose(0, 1, 3, 2)
        .reshape(feature.shape[1], -1, out_shape, out_shape)
        .transpose(1, 0, 2, 3)
    )
    feature = feature[:, :, 0, 0]
    ped_type_near_cls = feature[:, 0:4]
    ped_type_far_cls = feature[:, 4:8]
    ped_pose_near_cls = feature[:, 10:15]
    ped_pose_far_cls = feature[:, 15:20]
    ped_pose_cut_near_cls = feature[:, 20:25]
    ped_pose_cut_far_cls = feature[:, 25:30]
    ride_pose_near_cls = feature[:, 30:35]
    ride_pose_far_cls = feature[:, 35:40]
    ped_reg_near = feature[:, 40:42]
    ped_reg_far = feature[:, 42:44]
    fake_cls = feature[:, 50:52]
    return (
        ped_type_near_cls,
        ped_pose_near_cls,
        ped_pose_cut_near_cls,
        ride_pose_near_cls,
        ped_type_far_cls,
        ped_pose_far_cls,
        ped_pose_cut_far_cls,
        ride_pose_far_cls,
        ped_reg_near,
        ped_reg_far,
        fake_cls,
    )


def _run_attribute_model(
    model: torch.jit.ScriptModule,
    imgs_merge: np.ndarray,
    device: torch.device,
) -> Any:
    if imgs_merge.size == 0:
        return None

    img2net = torch.from_numpy(imgs_merge.astype(np.float32)).to(device) / 256.0
    with torch.no_grad():
        return model(img2net)


def _slice_prediction_batch(prediction: Any, batch_idx: int) -> Any:
    if prediction is None:
        return None

    if isinstance(prediction, (list, tuple)):
        return [tensor[batch_idx:batch_idx + 1] for tensor in prediction]
    return prediction[batch_idx:batch_idx + 1]


def _infer_vehicle_attributes(
    model: torch.jit.ScriptModule,
    frame: np.ndarray,
    detections: list[dict[str, Any]],
    device: torch.device,
) -> dict[int, dict[str, Any]]:
    if not detections:
        return {}

    param = [3, 4, 4, 96]
    w_h = param[1] * param[2]
    patches_all, crop_rects, side_all, det_indices = [], [], [], []

    for det in detections:
        box = det["bbox"]
        side = float(max(box[2] - box[0], box[3] - box[1]))
        patch, crop_rect, _ = _get_img_patch_for_car(frame, box, side_param=1.1)
        patches_all.append(patch)
        crop_rects.append(crop_rect)
        side_all.append(side)
        det_indices.append(det["index"])

    imgs_merge = _to_merge_img(patches_all, param)
    prediction_batch = _run_attribute_model(model, imgs_merge, device)
    results = {}

    for i in range(imgs_merge.shape[0]):
        pred = _slice_prediction_batch(prediction_batch, i)
        (
            car_type_near_cls,
            car_type_far_cls,
            tail_near_cls,
            tail_far_cls,
            car_head_type_cls,
            car_head_ye_double_cls,
            out_pose_type_cls,
            fake_car_cls,
            out_shelter_cls,
            out_reg1,
            out_reg2,
            out_reg3,
        ) = _vehicle_output_analysis(pred, param)

        for j in range(w_h):
            idx = i * w_h + j
            if idx >= len(detections):
                continue

            crop_rect = crop_rects[idx]
            side = side_all[idx]
            reg_box = None
            out_cls_type = "-1"

            shelter_label = int(np.argmax(out_shelter_cls[j][0:2], 0))
            left_label = int(np.argmax(out_shelter_cls[j][2:4], 0))
            right_label = int(np.argmax(out_shelter_cls[j][4:6], 0))
            if shelter_label == 0:
                shelter_type = 0
            elif left_label == 1 and right_label == 1:
                shelter_type = 3
            elif left_label == 1 and right_label == 0:
                shelter_type = 1
            elif left_label == 0 and right_label == 1:
                shelter_type = 2
            else:
                shelter_type = 4

            is_near = side > 60
            if is_near:
                type_car = int(np.where(car_type_near_cls[j] == np.max(car_type_near_cls[j]))[0][0])
                pose_car = int(np.argmax(tail_near_cls[j], 0))
            else:
                type_car = int(np.where(car_type_far_cls[j] == np.max(car_type_far_cls[j]))[0][0])
                pose_car = int(np.argmax(tail_far_cls[j], 0))

            if pose_car == 0:
                car_pose = "bg"
                context = car_pose
            elif pose_car == 1:
                car_pose = "tail"
                dots_1 = out_reg1[j][0:12]
                dots_2_3 = out_reg1[j][12:24]
                dots_5 = out_reg1[j][24:36]
                dots_11 = out_reg1[j][36:48]
                dots_4_9_10 = out_reg2[j][0:12]
                dots_6 = out_reg2[j][12:24]
                dots_7 = out_reg2[j][24:36]
                dots_8 = out_reg2[j][36:48]
                dots_12_13 = out_reg3[j][16:24]

                if type_car == 1:
                    res_reg = dots_1
                elif type_car in (2, 3):
                    res_reg = dots_2_3
                elif type_car in (4, 9, 10):
                    res_reg = dots_4_9_10
                elif type_car == 5:
                    res_reg = dots_5
                elif type_car == 6:
                    res_reg = dots_6
                elif type_car == 7:
                    res_reg = dots_7
                elif type_car == 8:
                    res_reg = dots_8
                elif type_car == 11:
                    res_reg = dots_11
                elif type_car >= 12:
                    res_reg = dots_12_13
                else:
                    res_reg = None

                if type_car != 0 and res_reg is not None:
                    res_reg = _unscale_and_uncenter(
                        res_reg,
                        crop_rect[0],
                        crop_rect[1],
                        crop_rect[2] - crop_rect[0],
                        crop_rect[3] - crop_rect[1],
                    )
                    reg_box = _convert_point_to_rect(
                        res_reg, type_car, frame.shape[1], frame.shape[0]
                    )
                context = f"{car_pose} {type_car} {shelter_type}"
                out_cls_type = str(type_car)
            elif pose_car == 2:
                car_pose = "head"
                type_head = int(np.argmax(car_head_type_cls[j], 0))
                type_head_ye = int(np.argmax(car_head_ye_double_cls[j], 0))
                dots_car_head = out_reg3[j][0:12]
                dots_car_head_ye = out_reg3[j][12:16]

                if type_head > 0 and type_head < 6:
                    reg_head = dots_car_head
                elif type_head == 6:
                    reg_head = dots_car_head_ye
                else:
                    reg_head = None

                if type_head not in (0, 7) and reg_head is not None:
                    reg_head = _unscale_and_uncenter(
                        reg_head,
                        crop_rect[0],
                        crop_rect[1],
                        crop_rect[2] - crop_rect[0],
                        crop_rect[3] - crop_rect[1],
                    )
                    if type_head < 6:
                        reg_box = _convert_point_to_rect(
                            reg_head, 2, frame.shape[1], frame.shape[0]
                        )
                    else:
                        dots_car_head = _unscale_and_uncenter(
                            dots_car_head,
                            crop_rect[0],
                            crop_rect[1],
                            crop_rect[2] - crop_rect[0],
                            crop_rect[3] - crop_rect[1],
                        )
                        reg_box = _convert_point_to_rect(
                            dots_car_head, 2, frame.shape[1], frame.shape[0]
                        )
                if type_head == 6:
                    type_head = type_head_ye
                out_cls_type = str(type_head)
                context = f"{car_pose} {type_head} {shelter_type}"
            else:
                car_pose = "others"
                double_pose_cls = int(np.argmax(out_pose_type_cls[j][0:4], 0))
                context = f"{car_pose} {double_pose_cls}"
                out_cls_type = str(double_pose_cls)

            is_fake = int(np.argmax(fake_car_cls[j], 0))
            if is_fake == 1:
                context = f"{context}_fake_{is_fake}"

            result = {
                "task": VEHICLE_TASK,
                "label": context,
                "attr_cls": out_cls_type,
                "pose": car_pose,
                "shelter_type": shelter_type,
                "is_fake": is_fake,
            }
            if reg_box is not None:
                result["reg_box"] = [float(v) for v in reg_box]
            results[det_indices[idx]] = result

    return results


def _infer_pedestrian_attributes(
    model: torch.jit.ScriptModule,
    frame: np.ndarray,
    detections: list[dict[str, Any]],
    device: torch.device,
) -> dict[int, dict[str, Any]]:
    if not detections:
        return {}

    param = [3, 4, 4, 96]
    w_h = param[1] * param[2]
    patches_all, crop_rects, side_all, box_all, det_indices = [], [], [], [], []

    for det in detections:
        box = det["bbox"]
        side = float(max(box[2] - box[0], box[3] - box[1]))
        patch, crop_rect, _ = _get_img_patch_for_ped(frame, box, side_param=1.1)
        patches_all.append(patch)
        crop_rects.append(crop_rect)
        side_all.append(side)
        box_all.append(box)
        det_indices.append(det["index"])

    imgs_merge = _to_merge_img(patches_all, param)
    prediction_batch = _run_attribute_model(model, imgs_merge, device)
    results = {}

    for i in range(imgs_merge.shape[0]):
        pred = _slice_prediction_batch(prediction_batch, i)
        (
            ped_type_near_cls,
            ped_pose_near_cls,
            ped_pose_cut_near_cls,
            ride_pose_near_cls,
            ped_type_far_cls,
            ped_pose_far_cls,
            ped_pose_cut_far_cls,
            ride_pose_far_cls,
            ped_reg_near,
            ped_reg_far,
            fake_cls,
        ) = _ped_output_analysis(pred, param)

        for j in range(w_h):
            idx = i * w_h + j
            if idx >= len(detections):
                continue

            crop_rect = crop_rects[idx]
            box = box_all[idx]
            side = side_all[idx]
            is_near = side > 60
            pose_ped = 0

            if is_near:
                ped_reg = ped_reg_near[j].copy()
                type_ped = int(np.argmax(ped_type_near_cls[j], 0))
                if type_ped == 1:
                    pose_ped = int(np.argmax(ped_pose_near_cls[j], 0))
                elif type_ped == 2:
                    pose_ped = int(np.argmax(ped_pose_cut_near_cls[j], 0))
                elif type_ped == 3:
                    pose_ped = int(np.argmax(ride_pose_near_cls[j], 0))
            else:
                ped_reg = ped_reg_far[j].copy()
                type_ped = int(np.argmax(ped_type_far_cls[j], 0))
                if type_ped == 1:
                    pose_ped = int(np.argmax(ped_pose_far_cls[j], 0))
                elif type_ped == 2:
                    pose_ped = int(np.argmax(ped_pose_cut_far_cls[j], 0))
                elif type_ped == 3:
                    pose_ped = int(np.argmax(ride_pose_far_cls[j], 0))

            ped_reg[0] = (ped_reg[0] * 2.5 + 40) / 80.0 * (crop_rect[3] - crop_rect[1]) + crop_rect[1]
            ped_reg[1] = (ped_reg[1] * 2.5 + 40) / 80.0 * (crop_rect[3] - crop_rect[1]) + crop_rect[1]

            context = f"{PED_TYPE_INDEX[type_ped]}_{PED_POSE_INDEX[pose_ped]}"
            if type_ped == 0:
                context = PED_TYPE_INDEX[type_ped]

            is_fake = int(np.argmax(fake_cls[j], 0))
            if is_fake:
                context = f"{context}_fake_{is_fake}"

            result = {
                "task": PEDESTRIAN_TASK,
                "label": context,
                "attr_cls": type_ped,
                "pose": PED_POSE_INDEX[pose_ped],
                "is_fake": is_fake,
                "y_range": [float(ped_reg[0]), float(ped_reg[1])],
            }
            if type_ped == 1:
                result["reg_box"] = [
                    float(box[0]),
                    float(ped_reg[0]),
                    float(box[2]),
                    float(ped_reg[1]),
                ]
            results[det_indices[idx]] = result

    return results


def _infer_tsr_attributes(
    model: torch.jit.ScriptModule,
    frame: np.ndarray,
    detections: list[dict[str, Any]],
    device: torch.device,
    tsr_labels: list[str] | None = None,
) -> dict[int, dict[str, Any]]:
    if not detections:
        return {}

    param = [3, 4, 4, 96]
    w_h = param[1] * param[2]
    patches_all, det_indices = [], []

    for det in detections:
        box = det["bbox"]
        patch, _, _ = _get_img_patch_for_car(frame, box, side_param=1.0)
        patches_all.append(patch)
        det_indices.append(det["index"])

    imgs_merge = _to_merge_img(patches_all, param)
    prediction_batch = _run_attribute_model(model, imgs_merge, device)
    results = {}

    for i in range(imgs_merge.shape[0]):
        pred = _slice_prediction_batch(prediction_batch, i)
        feature = pred[0].detach().cpu().numpy()
        feature = np.reshape(feature, (1, 1, param[1], param[2]))
        out_shape = feature.shape[-1] // param[1]
        feature = (
            feature.transpose(1, 0, 3, 2)
            .reshape(feature.shape[1], -1, out_shape, out_shape * param[1])
            .transpose(0, 1, 3, 2)
            .reshape(feature.shape[1], -1, out_shape, out_shape)
            .transpose(1, 0, 2, 3)
        )
        feature = feature[:, :, 0, 0]

        for j in range(w_h):
            idx = i * w_h + j
            if idx >= len(detections):
                continue

            type_idx = int(np.asarray(feature[j]).item())
            if tsr_labels and 0 <= type_idx < len(tsr_labels):
                label = tsr_labels[type_idx]
            elif type_idx > 3:
                label = str(type_idx)
            else:
                label = "0"

            results[det_indices[idx]] = {
                "task": TSR_TASK,
                "label": label,
                "attr_cls": type_idx,
            }

    return results


def predict_attributes_for_detections(
    attribute_models: AttributeModelBundle | None,
    frame_bgr: np.ndarray,
    detections: list[dict[str, Any]],
) -> list[dict[str, Any] | None]:
    """Run attribute models on merged detections from one frame."""
    if attribute_models is None:
        return [None] * len(detections)

    indexed = []
    for idx, det in enumerate(detections):
        indexed.append(
            {
                "index": idx,
                "cls_id": int(det["cls_id"]),
                "score": float(det["score"]),
                "bbox": [float(v) for v in det["bbox"]],
                "roi_id": det.get("roi_id"),
            }
        )

    results = {}
    vehicle_dets = _filter_vehicle_detections(
        [det for det in indexed if det["cls_id"] in attribute_models.vehicle_class_ids],
        attribute_models,
    )
    pedestrian_dets = _filter_pedestrian_detections(
        [det for det in indexed if det["cls_id"] in attribute_models.pedestrian_class_ids],
        attribute_models,
    )
    tsr_dets = _filter_tsr_detections(
        [det for det in indexed if det["cls_id"] in attribute_models.tsr_class_ids],
        attribute_models,
    )

    if attribute_models.car_model is not None and vehicle_dets:
        results.update(
            _infer_vehicle_attributes(
                attribute_models.car_model, frame_bgr, vehicle_dets, attribute_models.device
            )
        )

    if attribute_models.ped_model is not None and pedestrian_dets:
        results.update(
            _infer_pedestrian_attributes(
                attribute_models.ped_model, frame_bgr, pedestrian_dets, attribute_models.device
            )
        )

    if attribute_models.tsr_model is not None and tsr_dets:
        results.update(
            _infer_tsr_attributes(
                attribute_models.tsr_model,
                frame_bgr,
                tsr_dets,
                attribute_models.device,
                attribute_models.tsr_labels,
            )
        )

    return [results.get(idx) for idx in range(len(detections))]


def _to_jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {k: _to_jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_jsonable(v) for v in value]
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        return float(value)
    return value


def attribute_result_to_json(attribute_result: dict[str, Any] | None) -> dict[str, Any] | None:
    if attribute_result is None:
        return None
    return _to_jsonable(attribute_result)


def draw_attribute_overlays(
    img: np.ndarray,
    detections: list[np.ndarray] | np.ndarray,
    attribute_results: list[dict[str, Any] | None],
) -> np.ndarray:
    """Draw a second attribute label line and optional auxiliary boxes."""
    if not attribute_results:
        return img

    out = img.copy()
    for det, attr in zip(detections, attribute_results):
        if attr is None:
            continue

        x1, y1, x2, y2 = [int(float(v)) for v in det[:4]]
        label = attr.get("label")
        task = attr.get("task")
        color = ATTRIBUTE_TEXT_COLORS.get(task, (255, 255, 255))

        if label:
            (tw, th), bl = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            label_w = tw + ATTRIBUTE_LABEL_PAD_X * 2
            label_h = th + bl + ATTRIBUTE_LABEL_PAD_Y * 2
            max_x0 = max(0, out.shape[1] - label_w - 1)
            max_y0 = max(0, out.shape[0] - label_h - 1)
            box_x0 = min(max(x1, 0), max_x0)
            box_y0 = y2 + ATTRIBUTE_LABEL_GAP_PX
            if box_y0 + label_h >= out.shape[0]:
                box_y0 = y1 - ATTRIBUTE_LABEL_GAP_PX - label_h
            box_y0 = min(max(box_y0, 0), max_y0)
            box_x1 = min(out.shape[1] - 1, box_x0 + label_w)
            box_y1 = min(out.shape[0] - 1, box_y0 + label_h)
            text_y = box_y0 + ATTRIBUTE_LABEL_PAD_Y + th

            overlay = out.copy()
            cv2.rectangle(overlay, (box_x0, box_y0), (box_x1, box_y1), ATTRIBUTE_LABEL_BG_COLOR, -1)
            cv2.addWeighted(overlay, ATTRIBUTE_LABEL_BG_ALPHA, out, 1.0 - ATTRIBUTE_LABEL_BG_ALPHA, 0, out)
            cv2.putText(
                out,
                label,
                (box_x0 + ATTRIBUTE_LABEL_PAD_X, text_y),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                1,
            )

        reg_box = attr.get("reg_box")
        if reg_box is not None and len(reg_box) == 4:
            rx1, ry1, rx2, ry2 = [int(float(v)) for v in reg_box]
            cv2.rectangle(out, (rx1, ry1), (rx2, ry2), color, 2)

    return out
