# SPDX-License-Identifier: Apache-2.0
"""The Autoware BEVDet configuration this port reproduces, with the source of every value.

``AW/`` = autoware_universe ``perception/autoware_tensorrt_bevdet`` @ 9ceaccf (package 0.53.0); ``V/`` = the
``bevdet_vendor`` 0.2.1 inference library (the C++/CUDA code the node links, clean-room: only the behaviour described
in ``research/bevdet/SPEC.md`` is reproduced, no code was copied). ``S:<line>`` = ``research/bevdet/SPEC.md``.

The model yaml (``AW/config/bevdet_r50_4dlongterm_depth.yaml``) and the node parameters
(``AW/config/bevdet.param.yaml``) are embedded as :data:`DEFAULT_MODEL_YAML` / :data:`DEFAULT_NODE_PARAMS` so the
package needs no Autoware checkout; :meth:`BEVDetConfig.from_yaml` reads the real files (a host test checks that
the embedded copy equals them). Values are derived exactly like the vendor (``V/src/bevdet.cpp:136-231``): float32
arithmetic and C float -> int truncation.

Parameter classes (PLAN.md section 0.4): every field below is COMPILE (changes a device shape or the weights' meaning)
unless marked RT-host; the RT-host values are only the defaults of ``api.BEVDet.RUNTIME_PARAMS``.

Pure Python + numpy: importing this module never imports torch / ttnn.
"""
from __future__ import annotations

import copy
from dataclasses import dataclass, field
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple

import numpy as np

__all__ = ["CAMERA_ORDER", "CLASS_NAMES", "AUTOWARE_LABELS", "HEAD_NAMES", "HEAD_CHANNELS", "DEFAULT_MODEL_YAML",
           "DEFAULT_NODE_PARAMS", "BEVDetConfig", "default_config"]

F32 = np.float32

# Network camera order = the node's data_params.cams (AW/config/bevdet.param.yaml:11). The network is permutation
# invariant over cameras (S:580); only the pairing of each image with its own calibration matters.
CAMERA_ORDER: Tuple[str, ...] = ("CAM_FRONT_LEFT", "CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_BACK_LEFT", "CAM_BACK",
                                 "CAM_BACK_RIGHT")
# Network class order (model yaml tasks[0].class_names = node post_process_params.class_names).
CLASS_NAMES: Tuple[str, ...] = ("car", "truck", "construction_vehicle", "bus", "trailer", "barrier", "motorcycle",
                                "bicycle", "pedestrian", "traffic_cone")
# The label the node publishes per network class (AW/src/ros_utils.cpp:32-51 getSemanticType: everything not
# listed there is UNKNOWN).
_AUTOWARE_LABEL = {"car": "CAR", "truck": "TRUCK", "bus": "BUS", "trailer": "TRAILER", "bicycle": "BICYCLE",
                   "motorcycle": "MOTORCYCLE", "pedestrian": "PEDESTRIAN"}
AUTOWARE_LABELS: Tuple[str, ...] = tuple(_AUTOWARE_LABEL.get(c, "UNKNOWN") for c in CLASS_NAMES)
# CenterHead outputs (ONNX graph outputs reg_0 ... heatmap_0; model yaml common_head + num_class).
HEAD_NAMES: Tuple[str, ...] = ("reg_0", "height_0", "dim_0", "rot_0", "vel_0", "heatmap_0")
HEAD_CHANNELS: Dict[str, int] = {"reg_0": 2, "height_0": 1, "dim_0": 3, "rot_0": 2, "vel_0": 2, "heatmap_0": 10}

# AW/config/bevdet_r50_4dlongterm_depth.yaml (keys the vendor reads; the unused ones are listed in S:136).
DEFAULT_MODEL_YAML: Dict[str, Any] = {
    "bev_range": [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0],
    "mean": [123.675, 116.28, 103.53],
    "std": [58.395, 57.12, 57.375],
    "use_depth": True,
    "use_adj": True,
    "adj_num": 8,
    "transform_size": 6,
    "cam_params_size": 27,
    "sampling": "nearest",
    "data_config": {"Ncams": 6, "cams": list(CAMERA_ORDER), "resize_radio": 0.44, "crop": [140, 0], "flip": True,
                    "input_size": [256, 704], "resize": [-0.06, 0.11], "resize_test": 0.0, "rot": [-5.4, 5.4],
                    "src_size": [900, 1600]},
    "grid_config": {"depth": [1.0, 60.0, 0.5], "x": [-51.2, 51.2, 0.8], "y": [-51.2, 51.2, 0.8], "z": [-5, 3, 8]},
    "model": {"bevpool_channels": 80,
              "coder": {"code_size": 9, "max_num": 500, "post_center_range": [-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
                        "score_threshold": 0.1},
              "common_head": {"channels": [2, 1, 3, 2, 2], "names": ["reg", "height", "dim", "rot", "vel"]},
              "down_sample": 16,
              "tasks": [{"class_names": list(CLASS_NAMES), "num_class": 10}]},
    "test_cfg": {"max_per_img": 500, "max_pool_nms": False, "min_radius": [4, 12, 10, 1, 0.85, 0.175],
                 "nms_rescale_factor": [[1.0, 0.7, 0.7, 0.4, 0.55, 1.1, 1.0, 1.0, 1.5, 3.5]], "nms_thr": [0.2],
                 "nms_type": ["rotate"], "post_center_limit_range": [-61.2, -61.2, -10.0, 61.2, 61.2, 10.0],
                 "post_max_size": 500, "pre_max_size": 1000, "score_threshold": 0.1},
}
# AW/config/bevdet.param.yaml (ros__parameters the node reads that change outputs).
DEFAULT_NODE_PARAMS: Dict[str, Any] = {
    "precision": "fp16",  # the deployed default; this port implements the fp32 path (S:166-180, decision D8)
    "data_params": {"camera_count": 6, "cams": list(CAMERA_ORDER)},
    "post_process_params": {"score_threshold": 0.2, "class_names": list(CLASS_NAMES)},
}


def _f32(v: Any) -> np.float32:
    return F32(v)


@dataclass(frozen=True)
class BEVDetConfig:
    """All constants of the deployed BEVDet (vendor fp32 semantics). Build it with :func:`default_config` or
    :meth:`from_yaml`; every derived value follows ``V/src/bevdet.cpp:136-231``."""

    model_yaml: Mapping[str, Any] = field(default_factory=lambda: copy.deepcopy(DEFAULT_MODEL_YAML), repr=False)
    node_params: Mapping[str, Any] = field(default_factory=lambda: copy.deepcopy(DEFAULT_NODE_PARAMS), repr=False)
    # Capacities of the BEVPool tables (TRT max profile, V/src/bevdet.cpp:644-657; S:590). COMPILE: device buffer
    # sizes in the TT port. Fixed constants, never derived from the core grid (PLAN.md 0.2).
    m_max: int = 370_000
    n_max: int = 14_000

    # ---- cameras and images ----------------------------------------------------------------------------
    @property
    def camera_order(self) -> Tuple[str, ...]:
        return tuple(self.model_yaml["data_config"]["cams"])

    @property
    def n_cams(self) -> int:
        return int(self.model_yaml["data_config"]["Ncams"])

    @property
    def src_hw(self) -> Tuple[int, int]:
        """Image size the node feeds the network (it resizes every image to it): (900, 1600)."""
        h, w = self.model_yaml["data_config"]["src_size"]
        return int(h), int(w)

    @property
    def input_hw(self) -> Tuple[int, int]:
        """Network input size after the Preprocess plugin's nearest resize + crop: (256, 704)."""
        h, w = self.model_yaml["data_config"]["input_size"]
        return int(h), int(w)

    @property
    def crop_hw(self) -> Tuple[int, int]:
        """(crop_h, crop_w) = (140, 0): the crop offset in RESIZED pixels (yaml ``crop`` is (h, w))."""
        h, w = self.model_yaml["data_config"]["crop"]
        return int(h), int(w)

    @property
    def resize(self) -> np.float32:
        """``resize_radio`` = input_w / src_w in float32 (V/src/bevdet.cpp:204; the yaml key ``resize_radio`` is not
        read): 0.44f. It is also the ONNX Preprocess attribute (0.4399999976158142)."""
        return F32(self.input_hw[1]) / F32(self.src_hw[1])

    @property
    def mean(self) -> np.ndarray:
        """Per-plane mean, applied to the B, G, R planes in this order (BEVDet trained BGR planes with the "RGB"
        ImageNet statistics, S:188-191)."""
        return np.asarray(self.model_yaml["mean"], F32)

    @property
    def std(self) -> np.ndarray:
        return np.asarray(self.model_yaml["std"], F32)

    @property
    def down_sample(self) -> int:
        return int(self.model_yaml["model"]["down_sample"])

    @property
    def feat_hw(self) -> Tuple[int, int]:
        """Image-feature map size (16, 44) = input / down_sample (integer division as in the vendor)."""
        return self.input_hw[0] // self.down_sample, self.input_hw[1] // self.down_sample

    # ---- depth bins and BEV grid -----------------------------------------------------------------------
    @property
    def depth_bins(self) -> Tuple[np.float32, np.float32, np.float32]:
        return tuple(_f32(v) for v in self.model_yaml["grid_config"]["depth"])  # type: ignore[return-value]

    @property
    def n_depth(self) -> int:
        """D = (int)((60 - 1) / 0.5) = 118 (float32 arithmetic, C truncation)."""
        d0, d1, ds = self.depth_bins
        return int(F32(d1 - d0) / ds)

    def _axis(self, key: str) -> Tuple[np.float32, np.float32, np.float32]:
        return tuple(_f32(v) for v in self.model_yaml["grid_config"][key])  # type: ignore[return-value]

    @property
    def grid_x(self) -> Tuple[np.float32, np.float32, np.float32]:
        return self._axis("x")

    @property
    def grid_y(self) -> Tuple[np.float32, np.float32, np.float32]:
        return self._axis("y")

    @property
    def grid_z(self) -> Tuple[np.float32, np.float32, np.float32]:
        return self._axis("z")

    @property
    def bev_xyz(self) -> Tuple[int, int, int]:
        """(nx, ny, nz) = (128, 128, 1)."""
        out = []
        for lo, hi, step in (self.grid_x, self.grid_y, self.grid_z):
            out.append(int(F32(hi - lo) / step))
        return tuple(out)  # type: ignore[return-value]

    @property
    def bev_hw(self) -> Tuple[int, int]:
        """BEV map (H, W) = (ny, nx) = (128, 128); row = y index, column = x index."""
        nx, ny, _ = self.bev_xyz
        return ny, nx

    @property
    def bev_channels(self) -> int:
        """C = 80 (model.bevpool_channels)."""
        return int(self.model_yaml["model"]["bevpool_channels"])

    # ---- temporal fusion -------------------------------------------------------------------------------
    @property
    def adj_num(self) -> int:
        """History length 8 (COMPILE: the BEV encoder takes (1 + 8) * 80 = 720 channels)."""
        return int(self.model_yaml["adj_num"]) if self.model_yaml.get("use_adj") else 0

    @property
    def bev_in_channels(self) -> int:
        return (1 + self.adj_num) * self.bev_channels

    # ---- head and post-processing ----------------------------------------------------------------------
    @property
    def class_names(self) -> Tuple[str, ...]:
        return tuple(c for task in self.model_yaml["model"]["tasks"] for c in task["class_names"])

    @property
    def num_classes(self) -> int:
        return len(self.class_names)

    @property
    def autoware_labels(self) -> Tuple[str, ...]:
        return tuple(_AUTOWARE_LABEL.get(c, "UNKNOWN") for c in self.class_names)

    @property
    def decode_score_threshold(self) -> float:
        """RT-host. Vendor decode threshold: a cell emits a box when sigmoid(max logit) > 0.1 (test_cfg)."""
        return float(self.model_yaml["test_cfg"]["score_threshold"])

    @property
    def max_pre_nms(self) -> int:
        """RT-host. ``test_cfg.max_per_img`` = 500 is the vendor's PRE-NMS cap (V/src/bevdet.cpp:168,
        V/src/postprocess.cu:183); ``pre_max_size`` (1000) is never read."""
        return int(self.model_yaml["test_cfg"]["max_per_img"])

    @property
    def max_post_nms(self) -> int:
        """RT-host. ``test_cfg.post_max_size`` = 500."""
        return int(self.model_yaml["test_cfg"]["post_max_size"])

    @property
    def nms_iou_threshold(self) -> float:
        """RT-host. ``test_cfg.nms_thr[0]`` = 0.2 (rotated BEV IoU, class-agnostic within the task)."""
        return float(self.model_yaml["test_cfg"]["nms_thr"][0])

    @property
    def nms_rescale_factor(self) -> np.ndarray:
        """RT-host. Scale-NMS factors per class (l, w, h are multiplied before the NMS and divided after)."""
        return np.asarray([f for task in self.model_yaml["test_cfg"]["nms_rescale_factor"] for f in task], F32)

    @property
    def score_threshold(self) -> float:
        """RT-host. The node's output filter: boxes with score < 0.2 are dropped (AW/src/ros_utils.cpp:59)."""
        return float(self.node_params["post_process_params"]["score_threshold"])

    @property
    def cam_params_size(self) -> int:
        return int(self.model_yaml["cam_params_size"])

    @property
    def transform_size(self) -> int:
        return int(self.model_yaml["transform_size"])

    # ---- construction ----------------------------------------------------------------------------------
    @classmethod
    def from_yaml(cls, model_yaml: Any, node_params: Optional[Any] = None, **overrides: Any) -> "BEVDetConfig":
        """Read ``bevdet_r50_4dlongterm_depth.yaml`` (path or mapping) and optionally ``bevdet.param.yaml`` (path
        or mapping; the ``/**: ros__parameters`` wrapper is accepted)."""
        my = _load_yaml(model_yaml)
        npar = copy.deepcopy(DEFAULT_NODE_PARAMS) if node_params is None else _load_yaml(node_params)
        if "/**" in npar:
            npar = npar["/**"]["ros__parameters"]
        return cls(model_yaml=my, node_params=npar, **overrides)

    def check(self) -> None:
        """Assert the structural constants the weights and the TT graph are built for (COMPILE values)."""
        problems = []
        expect = {"n_cams": 6, "input_hw": (256, 704), "feat_hw": (16, 44), "n_depth": 118, "bev_xyz": (128, 128, 1),
                  "bev_channels": 80, "adj_num": 8, "num_classes": 10}
        for name, value in expect.items():
            if getattr(self, name) != value:
                problems.append(f"{name}={getattr(self, name)!r}, the shipped weights need {value!r}")
        if problems:
            raise ValueError("unsupported BEVDet configuration: " + "; ".join(problems))

    def summary(self) -> Dict[str, Any]:
        """JSON-able view (for ``model.info`` and golden metadata)."""
        return {"camera_order": list(self.camera_order), "src_hw": list(self.src_hw), "input_hw": list(self.input_hw),
                "crop_hw": list(self.crop_hw), "resize": float(self.resize), "mean": self.mean.tolist(),
                "std": self.std.tolist(), "feat_hw": list(self.feat_hw), "n_depth": self.n_depth,
                "depth_bins": [float(v) for v in self.depth_bins], "grid_x": [float(v) for v in self.grid_x],
                "grid_y": [float(v) for v in self.grid_y], "grid_z": [float(v) for v in self.grid_z],
                "bev_xyz": list(self.bev_xyz), "bev_channels": self.bev_channels, "adj_num": self.adj_num,
                "class_names": list(self.class_names), "autoware_labels": list(self.autoware_labels),
                "decode_score_threshold": self.decode_score_threshold, "max_pre_nms": self.max_pre_nms,
                "max_post_nms": self.max_post_nms, "nms_iou_threshold": self.nms_iou_threshold,
                "nms_rescale_factor": self.nms_rescale_factor.tolist(), "score_threshold": self.score_threshold,
                "m_max": self.m_max, "n_max": self.n_max}


def _load_yaml(src: Any) -> Dict[str, Any]:
    if isinstance(src, Mapping):
        return copy.deepcopy(dict(src))
    import yaml

    with open(src, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def default_config() -> BEVDetConfig:
    """The deployed configuration (embedded copy of the Autoware files)."""
    return BEVDetConfig()


def class_index(names: Sequence[str], name: str) -> int:
    return list(names).index(name)