# SPDX-License-Identifier: Apache-2.0
"""The Autoware ``autoware_bevfusion`` parameters this port reproduces, and its compile-time constants.

Sources (autoware_universe main ``9ceaccf``, ``perception/autoware_bevfusion`` = ``PKG/``; weights
``AutowareFoundation/bevfusion@v2.0`` = commit ``e1bf164b``, research/bevfusion/SPEC.md sections 2-5):

- from the weights repo: ``ml_package_bevfusion_lidar.param.yaml`` (class names, ``voxels_num`` [1, 128000, 256000],
  range, voxel size, ``num_proposals``, ``out_size_factor``, ``max_points_per_voxel``, ``use_intensity``),
  ``detection_class_remapper.param.yaml`` (the 8x8 area remapper) and ``deploy_metadata.yaml`` (``version: v2.0``);
  the launch file loads these from the model directory, not the package (``PKG/launch/bevfusion.launch.xml:12-13``);
- node-wide, :data:`NODE_PARAMS` = ``PKG/config/bevfusion_lidar.param.yaml`` (value-identical to autoware_launch's
  copy, S:bevfusion:35);
- the node's sanitization (``PKG/include/autoware/bevfusion/bevfusion_config.hpp:129-158``): distance bins must be
  ascending and are squared, score thresholds outside [0, 1) -> 0, ``circle_nms_dist_threshold <= 0`` keeps 0.5,
  yaw-norm thresholds outside [0, 1) -> 0; the grid ``int64((max - min) / size)`` in **float32** (1440, 1440, 40;
  float64 of the float32-rounded values would give 39 in z, S:bevfusion:156).

The range, voxel size and ``num_proposals`` also live in the ONNX (the sparse encoder's ``spatial_shape``, the
``bev_pos`` table, TopK ``K``): the weights loader checks that the yaml and the graph agree and refuses a mismatch.

Port constants (PLAN.md 2.6; COMPILE): :data:`MAX_VOXELS` 256,000 (Autoware's clip), the sparse level grids, the
per-level capacity buckets :data:`CAPACITY_BUCKETS` (the TT trace shapes; grid-independent, PLAN.md 0.2) and the
device top-k :data:`TOPK_PADDED` (500 -> 512). ``yaml`` is imported when a file is read, never at import.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple

import numpy as np

from ..ttaw.nms import ClassRemapper

__all__ = [
    "VARIANTS", "DEFAULT_VARIANT", "SUPPORTED_VERSION", "ONNX_FILE", "ML_PACKAGE_FILE", "REMAPPER_FILE",
    "METADATA_FILE", "NODE_PARAMS", "NUM_CLASSES", "NUM_POINT_FEATURES", "NUM_BOX_VALUES", "MAX_VOXELS",
    "SPARSE_LEVELS", "SPARSE_SHAPE", "BEV_SHAPE", "CAPACITY_BUCKETS", "TOPK_PADDED", "POOLED_CLASSES",
    "DROP_ZERO_POINTS", "BEVFusionConfig", "grid_size_f32",
]

VARIANTS = ("lidar",)            # the camera-lidar profile is planned (PLAN.md D5, 4.1 step 14)
DEFAULT_VARIANT = "lidar"
SUPPORTED_VERSION = "v2.0"       # deploy_metadata.yaml of AutowareFoundation/bevfusion@e1bf164b

ONNX_FILE = "bevfusion_lidar.onnx"
ML_PACKAGE_FILE = "ml_package_bevfusion_lidar.param.yaml"
REMAPPER_FILE = "detection_class_remapper.param.yaml"
METADATA_FILE = "deploy_metadata.yaml"

# PKG/config/bevfusion_lidar.param.yaml (autoware_universe 9ceaccf)
NODE_PARAMS: Dict[str, Any] = {
    "trt_precision": "fp16",                       # TensorRT only (the CPU fp32 run is the oracle)
    "cloud_capacity": 2_000_000,
    "densification_num_past_frames": 1,
    "densification_world_frame_id": "map",
    "circle_nms_dist_threshold": 0.5,
    "iou_nms_search_distance_2d": 10.0,
    "iou_nms_threshold": 0.1,
    "yaw_norm_thresholds": (0.3, 0.3, 0.3, 0.3, 0.0, 0.0, 0.0),        # by class index (7 listed, 5 used)
    "distance_bin_upper_limits": (50.0, 90.0, 121.0, 200.0),
    "min_confidence_score": 0.10,                  # every class and distance bin (min_confidence_scores.<CLASS>)
}

NUM_CLASSES = 5                  # CAR, TRUCK, BUS, BICYCLE, PEDESTRIAN (the HF ml_package; S:bevfusion:47)
NUM_POINT_FEATURES = 4           # x, y, z, time_lag (use_intensity: false; bevfusion_config.hpp:68-74)
NUM_BOX_VALUES = 10              # bbox_pred rows: cx, cy, z, log l, log w, log h, sin, cos, vx, vy
MAX_VOXELS = 256_000             # voxels_num max: Autoware clips above it (PKG/lib/bevfusion_trt.cpp:550-581)
POOLED_CLASSES = (0, 1, 2, 3)    # the local max pools CAR..BICYCLE only; PEDESTRIAN passes (S:bevfusion:258)
TOPK_PADDED = 512                # device top-k k (a multiple of 16 for topk_large_indices; first 500 kept)
DROP_ZERO_POINTS = True          # (0, 0, 0) rows are organized-cloud fillers (test.pcd), as the other lidar ports

# Sparse encoder grids (x, y, z) as the ONNX GetIndicePairsImplicitGemm attributes give them (S:bevfusion:226-235):
# level 1 is the voxel grid with z = 40 + 1 (mmdet3d's sparse_shape adds one), levels 2-4 follow from the strided
# convs, the conv_out level is the 180 x 180 x 2 to-dense input.
SPARSE_SHAPE = (1440, 1440, 41)
SPARSE_LEVELS = {"l1": (1440, 1440, 41), "l2": (720, 720, 21), "l3": (360, 360, 11), "l4": (180, 180, 5),
                 "out": (180, 180, 2)}
BEV_SHAPE = (180, 180)           # BEV feature grid (rows = x, columns = y), 256 channels (c * 2 + z)

# Per-level capacity buckets of the TT port (PLAN.md 2.6 "v1 ships per-level capacity buckets"): each bucket is a
# shape variant of the trace (active voxels of l1, l2, l3, l4, out); a frame runs in the smallest bucket that holds
# all five counts. PROVISIONAL (frozen at M1 from the device cost): sized from the 121 public golden frames
# (nuScenes max 22,314 / 27,281 / 15,517 / 6,169 / 4,569; PandaSet max 98,251 / 109,850 / 59,443 / 23,561 / 15,232;
# rosbag #144 98,334 / 97,787 / 45,023 / 15,499 / 11,493) with ~1.3x margin, plus the full bucket of PLAN.md 2.6
# (256k / 256k / 128k / 64k; conv_out 64k). l2 can exceed l1 (1.43x on PandaSet 019).
CAPACITY_BUCKETS: Dict[str, Tuple[int, int, int, int, int]] = {
    "s": (32_768, 40_960, 24_576, 8_192, 8_192),
    "m": (131_072, 147_456, 81_920, 32_768, 24_576),
    "full": (262_144, 262_144, 131_072, 65_536, 65_536),
}


def _ros_params(path: Path) -> Dict[str, Any]:
    import yaml

    with open(path) as f:
        doc = yaml.safe_load(f)
    try:
        return doc["/**"]["ros__parameters"]
    except (KeyError, TypeError):
        raise ValueError(f"{path}: not a ROS 2 parameter file (no /**: ros__parameters)") from None


def grid_size_f32(range_min: float, range_max: float, size: float) -> int:
    """``static_cast<int64_t>((max - min) / size)`` with float32 operands (``bevfusion_config.hpp:159-161``)."""
    return int(np.trunc(np.float32(np.float32(range_max) - np.float32(range_min)) / np.float32(size)))


@dataclass(frozen=True)
class BEVFusionConfig:
    """The ml_package / node parameters (COMPILE values marked in PORT_LOG.md's parameter table)."""

    version: str
    class_names: Tuple[str, ...]
    voxels_num: Tuple[int, int, int]                 # (min, opt, max) TensorRT profile; max = Autoware's clip
    point_cloud_range: Tuple[float, ...]
    voxel_size: Tuple[float, ...]
    num_proposals: int
    out_size_factor: int
    max_points_per_voxel: int
    use_intensity: bool
    remap_lists: Tuple[tuple, tuple, tuple] = field(repr=False, default=((), (), ()))
    node: Mapping[str, Any] = field(default_factory=lambda: dict(NODE_PARAMS))
    model_dir: Optional[str] = None

    # ---- construction --------------------------------------------------------------------------------------
    @classmethod
    def from_model_dir(cls, model_dir: Any, node: Optional[Mapping[str, Any]] = None) -> "BEVFusionConfig":
        """Read the ml_package, remapper and metadata yaml files of an ``AutowareFoundation/bevfusion`` snapshot."""
        d = Path(model_dir)
        if not (d / ML_PACKAGE_FILE).is_file():
            raise FileNotFoundError(f"no {ML_PACKAGE_FILE} under {d}")
        p = _ros_params(d / ML_PACKAGE_FILE)
        version = SUPPORTED_VERSION
        if (d / METADATA_FILE).is_file():
            import yaml

            with open(d / METADATA_FILE) as f:
                version = str((yaml.safe_load(f) or {}).get("version", ""))
        if version != SUPPORTED_VERSION:
            raise ValueError(f"deploy_metadata version {version!r} is not supported (this port reproduces "
                             f"{SUPPORTED_VERSION}: 5 classes, mean VFE, [V, 10, 4] voxels)")
        q = _ros_params(d / REMAPPER_FILE)
        remap = (tuple(int(v) for v in q["allow_remapping_by_area_matrix"]),
                 tuple(float(v) for v in q["min_area_matrix"]), tuple(float(v) for v in q["max_area_matrix"]))
        cfg = cls(version=version, class_names=tuple(p["class_names"]),
                  voxels_num=tuple(int(v) for v in p["voxels_num"]),
                  point_cloud_range=tuple(float(v) for v in p["point_cloud_range"]),
                  voxel_size=tuple(float(v) for v in p["voxel_size"]), num_proposals=int(p["num_proposals"]),
                  out_size_factor=int(p["out_size_factor"]), max_points_per_voxel=int(p["max_points_per_voxel"]),
                  use_intensity=bool(p["use_intensity"]), remap_lists=remap,
                  node={**NODE_PARAMS, **dict(node or {})}, model_dir=str(d))
        cfg.validate()
        return cfg

    @classmethod
    def from_weights(cls, weights_dir: Any, variant: str = DEFAULT_VARIANT,
                     node: Optional[Mapping[str, Any]] = None) -> "BEVFusionConfig":
        if variant not in VARIANTS:
            raise ValueError(f"variant={variant!r}: expected one of {VARIANTS} (camera_lidar is planned)")
        return cls.from_model_dir(weights_dir, node)

    def validate(self) -> None:
        if len(self.point_cloud_range) != 6 or len(self.voxel_size) != 3:
            raise ValueError("point_cloud_range needs 6 values and voxel_size 3")
        if len(self.voxels_num) != 3 or not (0 < self.voxels_num[0] <= self.voxels_num[1] <= self.voxels_num[2]):
            raise ValueError(f"voxels_num must be [min, opt, max], got {self.voxels_num}")
        if self.voxels_num[2] > MAX_VOXELS:
            raise ValueError(f"voxels_num max {self.voxels_num[2]} exceeds the port's capacity {MAX_VOXELS}")
        if len(self.class_names) != NUM_CLASSES:
            raise ValueError(f"the network has {NUM_CLASSES} classes, the yaml lists {len(self.class_names)}")
        if self.use_intensity:
            raise ValueError("use_intensity: true is a different network (conv_input Cin 5); v2.0 has false")
        if self.grid != (SPARSE_SHAPE[0], SPARSE_SHAPE[1], SPARSE_SHAPE[2] - 1):
            raise ValueError(f"voxel grid {self.grid} does not match the sparse encoder {SPARSE_SHAPE}")
        if (self.grid[0] // self.out_size_factor, self.grid[1] // self.out_size_factor) != BEV_SHAPE:
            raise ValueError("out_size_factor does not give the 180 x 180 BEV")
        if self.num_proposals <= 0 or self.num_proposals > TOPK_PADDED:
            raise ValueError(f"num_proposals must be in (0, {TOPK_PADDED}]")
        bins = self.node["distance_bin_upper_limits"]
        if not bins or list(bins) != sorted(bins):
            raise ValueError("distance_bin_upper_limits must be non-empty and ascending (bevfusion_config.hpp:129)")
        if len(self.node["yaw_norm_thresholds"]) < NUM_CLASSES:
            raise ValueError("yaw_norm_thresholds needs one value per class")

    # ---- derived values ------------------------------------------------------------------------------------
    @property
    def grid(self) -> Tuple[int, int, int]:
        """The voxel grid (x, y, z) in float32 arithmetic: (1440, 1440, 40)."""
        r, v = self.point_cloud_range, self.voxel_size
        return tuple(grid_size_f32(r[i], r[i + 3], v[i]) for i in range(3))  # type: ignore[return-value]

    @property
    def max_voxels(self) -> int:
        return int(self.voxels_num[2])

    @property
    def min_voxels(self) -> int:
        """Autoware aborts a frame with fewer voxels (``bevfusion_trt.cpp:550-581``; ``V < 1``)."""
        return int(self.voxels_num[0])

    @property
    def distance_bins_sq(self) -> np.ndarray:
        """float32 squared distance-bin upper limits (``bevfusion_config.hpp:133-137``: ``limit * limit``)."""
        b = np.asarray(self.node["distance_bin_upper_limits"], dtype=np.float32)
        return (b * b).astype(np.float32)

    def score_thresholds(self, value: Optional[float] = None) -> np.ndarray:
        """``(bins, classes)`` float32 table (one value everywhere: the shipped yaml's 0.10, or ``value``); entries
        outside [0, 1) -> 0 (``bevfusion_config.hpp:144-146``)."""
        v = float(self.node["min_confidence_score"] if value is None else value)
        v = v if 0.0 <= v < 1.0 else 0.0
        return np.full((len(self.node["distance_bin_upper_limits"]), len(self.class_names)), v, dtype=np.float32)

    @staticmethod
    def sanitized_circle_nms(value: Optional[float]) -> float:
        """``bevfusion_config.hpp:148-150``: a distance <= 0 keeps the default 0.5 m."""
        v = float(NODE_PARAMS["circle_nms_dist_threshold"] if value is None else value)
        return v if v > 0.0 else 0.5

    @property
    def yaw_norm_thresholds(self) -> np.ndarray:
        """float32 per class index; entries outside [0, 1) -> 0 (``bevfusion_config.hpp:151-155``)."""
        v = np.asarray(self.node["yaw_norm_thresholds"], dtype=np.float64)
        return np.where((v >= 0.0) & (v < 1.0), v, 0.0).astype(np.float32)

    @property
    def remapper(self) -> ClassRemapper:
        return ClassRemapper.from_lists(*self.remap_lists)

    def with_node(self, **overrides: Any) -> "BEVFusionConfig":
        return replace(self, node={**dict(self.node), **overrides})

    def describe(self) -> Dict[str, Any]:
        """JSON-able summary (``model.info``, golden metadata)."""
        return {"version": self.version, "class_names": list(self.class_names), "voxels_num": list(self.voxels_num),
                "point_cloud_range": list(self.point_cloud_range), "voxel_size": list(self.voxel_size),
                "grid": list(self.grid), "sparse_levels": {k: list(v) for k, v in SPARSE_LEVELS.items()},
                "bev": list(BEV_SHAPE), "num_proposals": self.num_proposals, "out_size_factor": self.out_size_factor,
                "max_points_per_voxel": self.max_points_per_voxel, "max_voxels": MAX_VOXELS,
                "capacity_buckets": {k: list(v) for k, v in CAPACITY_BUCKETS.items()},
                "node": {k: (list(v) if isinstance(v, tuple) else v) for k, v in self.node.items()}}