# SPDX-License-Identifier: Apache-2.0
"""The Autoware parameters this port reproduces and the port's COMPILE constants.

Target (PLAN.md D3, research/ptv3/SPEC.md 1.1): the weights of ``AutowareFoundation/ptv3@v4.0`` (commit
``1673104843af3c740145ac74dc4af6b8a4502bf3``) with the pre- / post-processing of the autoware_universe node at
``5eb9286d703b5de2dd8a30dfcd599d264ec7f13a`` (``perception/autoware_ptv3``; the last node commit before the breaking
#13155 and the only node revision that consumes the v4.0 graphs' inputs). Citations ``AU@5eb9286:<file>:<line>`` are
relative to ``perception/autoware_ptv3`` of that commit (exported to ``assets/ptv3/node_src_5eb9286`` in the
workspace); ``HF:<file>`` is the weights repository.

numpy only; no side effects on import.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, Tuple

import numpy as np

__all__ = [
    "WEIGHTS_REPO", "WEIGHTS_TAG", "WEIGHTS_REVISION", "NODE_REVISION", "ONNX_FILES", "ONNX_SHA256",
    "SEG_CLASSES", "DET_CLASSES", "PALETTE", "CLASS_ID_LUT", "POINTCLOUD_CLASSIFICATION", "INVALID_CLASS_ID",
    "FILTER_CLASSES", "OBJECT_LABELS", "DET_CLASS_TO_OBJECT_LABEL", "ENC_CHANNELS", "ENC_DEPTHS", "DEC_CHANNELS",
    "DEC_DEPTHS", "HEAD_DIM", "PATCH_SIZE", "ATTN_SCALE", "SPATIAL_SHAPES", "NUM_STAGES", "NUM_ORDERS",
    "BN_EPS", "LN_EPS", "BEV_SIZE", "NUM_PROPOSALS", "NEIGHBOUR_MODES", "DEFAULT_NEIGHBOUR_MODE",
    "RECONSTRUCTION_MODES", "PreprocessConfig", "DetectionConfig", "SegmentationConfig", "PTv3Config",
    "STAGE_CAPACITIES", "POOL_CHILDREN", "BEV_CELL_CANDIDATES", "CLASS_MAPPING", "stage_heads",
]

# ---------------------------------------------------------------------------------------------- artifacts
WEIGHTS_REPO = "AutowareFoundation/ptv3"
WEIGHTS_TAG = "v4.0"
WEIGHTS_REVISION = "1673104843af3c740145ac74dc4af6b8a4502bf3"
NODE_REVISION = "5eb9286d703b5de2dd8a30dfcd599d264ec7f13a"
ONNX_FILES = {"encoder": "ptv3_encoder.onnx", "seg": "ptv3_seg3d_head.onnx", "det": "ptv3_det3d_head.onnx"}
ONNX_SHA256 = {  # SPEC 1.4 (HF LFS sha256, re-hashed locally)
    "encoder": "c244cfe2e6016740f894fb24480e15d8e29a8c667c90bfcdf95cfbad8a34b9b8",
    "seg": "06dcc59e67c39e6ddf9ab26254a188cd85e4d765cfacbd48e9678c152edf5817",
    "det": "5b199d2765fcec19a4f15e2586f9dde0a85292ba8babd31d2275fd4787807abf",
}

# ---------------------------------------------------------------------------------------------- classes
# HF:ml_package_ptv3_seg3d_head.param.yaml:5-21
SEG_CLASSES: Tuple[str, ...] = (
    "car", "truck", "bus", "bicycle", "pedestrian", "traffic_cone", "barrier", "debris", "drivable_flat",
    "non_drivable_flat", "vegetation", "building", "vertical_thin", "static_clutter", "noise")
# HF:ml_package_ptv3_seg3d_head.param.yaml:22-38 (RGB per class)
PALETTE = np.array([[0, 150, 245], [160, 32, 240], [255, 255, 0], [255, 192, 203], [255, 0, 0], [255, 240, 150],
                    [255, 127, 0], [210, 180, 140], [255, 0, 255], [139, 137, 137], [0, 175, 0], [230, 230, 250],
                    [153, 50, 204], [188, 143, 143], [64, 224, 208]], np.uint8)
# autoware_core autoware_point_types types.hpp PointCloudClassification (SPEC 2, R8)
POINTCLOUD_CLASSIFICATION = {"CAR": 0, "TRUCK": 1, "BUS": 2, "MOTORCYCLE": 3, "BICYCLE": 4, "PEDESTRIAN": 5,
                             "ANIMAL": 6, "HAZARD": 7, "FLAT_SURFACE": 8, "STRUCTURE": 9, "VEGETATION": 10,
                             "NOISE": 11, "INVALID": 255}
INVALID_CLASS_ID = 255
# AU@5eb9286:config/ptv3.param.yaml:20-35 (segmentation3d.class_mapping)
CLASS_MAPPING = {"car": "CAR", "truck": "TRUCK", "bus": "BUS", "bicycle": "BICYCLE", "pedestrian": "PEDESTRIAN",
                 "traffic_cone": "HAZARD", "debris": "HAZARD", "vertical_thin": "HAZARD", "barrier": "STRUCTURE",
                 "building": "STRUCTURE", "static_clutter": "STRUCTURE", "drivable_flat": "FLAT_SURFACE",
                 "non_drivable_flat": "FLAT_SURFACE", "vegetation": "VEGETATION", "noise": "NOISE"}
# model class -> PointCloudClassification value (ptv3_config.hpp make_class_id_lut): [0,1,2,4,5,7,9,7,8,8,10,9,7,9,11]
CLASS_ID_LUT = np.array([POINTCLOUD_CLASSIFICATION[CLASS_MAPPING[c]] for c in SEG_CLASSES], np.uint8)
# AU@5eb9286:config/ptv3.param.yaml:36-39 (segmentation3d.filter.classes)
FILTER_CLASSES: Tuple[str, ...] = ("drivable_flat", "non_drivable_flat", "noise")

# HF:ml_package_ptv3_det3d_head.param.yaml:8
DET_CLASSES: Tuple[str, ...] = ("CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN", "TRAFFIC_CONE", "BARRIER")
# autoware_perception_msgs ObjectClassification (msgs 1.13.0)
OBJECT_LABELS: Tuple[str, ...] = ("UNKNOWN", "CAR", "TRUCK", "BUS", "TRAILER", "MOTORCYCLE", "BICYCLE", "PEDESTRIAN",
                                  "ANIMAL", "HAZARD")
# AU@5eb9286:lib/ros_utils.cpp:72-107 get_classification_type (TRAFFIC_CONE, BARRIER -> HAZARD)
DET_CLASS_TO_OBJECT_LABEL = np.array([1, 2, 3, 6, 7, 9, 9], np.uint8)

# ---------------------------------------------------------------------------------------------- network
ENC_CHANNELS = (32, 64, 128, 256, 512)           # HF:ml_package_ptv3_encoder.param.yaml:8
ENC_DEPTHS = (2, 2, 2, 2, 2)                     # blocks per encoder stage (graph)
DEC_CHANNELS = (64, 64, 128, 256)                # seg head decoder output channels at stages 0..3 (graph)
DEC_DEPTHS = (0, 0, 1, 1)                        # HF:ml_package_ptv3_seg3d_head.param.yaml:4
HEAD_DIM = 16                                    # heads = C / 16
PATCH_SIZE = 512                                 # serialized attention patch K (graph constant c512)
ATTN_SCALE = 0.25                                # 16 ** -0.5 (graph constant c0.25)
NUM_STAGES = 5
NUM_ORDERS = 2                                   # "z", "z-trans"; block b uses order b % 2
# GetIndicePairsImplicitGemm spatial_shape per stage (graph attributes; floor-halved 2048x2048x67)
SPATIAL_SHAPES = ((2048, 2048, 67), (1024, 1024, 33), (512, 512, 16), (256, 256, 8), (128, 128, 4))
BN_EPS = 1e-3                                    # every BN except the two bev_projector BNs (1e-5, SPEC 4.3)
LN_EPS = 1e-5
BEV_SIZE = 256                                   # det BEV grid 256 x 256 (bbox_voxel_size 0.96 = 0.12 * 2^3)
NUM_PROPOSALS = 500                              # /bbox_head/TopK_K

# ---------------------------------------------------------------------------------------------- host knobs
NEIGHBOUR_MODES = ("alias-first", "exact", "alias-last")   # SPEC 4.1 / R2; PLAN D3: alias-compat default
DEFAULT_NEIGHBOUR_MODE = "alias-first"
RECONSTRUCTION_MODES = ("full", "partial", "none")          # AU@5eb9286:config/ptv3.param.yaml:18 default full

# ---------------------------------------------------------------------------------------------- port COMPILE constants
# Static per-stage capacities of the device graph (PLAN.md 0.2: grid-independent constants chosen from data).
# Functional port (PLAN.md 2.7): ONE bucket at the node maximum of 256,000 voxels at stage 0. Pooled levels have
# no separate bound in node @5eb9286 (<= N0). Their capacities are 256,000 times the largest measured pooling
# ratio per stage plus margin (N1/N0 <= 0.662, N2/N1 <= 0.547, N3/N2 <= 0.472, N4/N3 <= 0.44 over the 4 T4 and 16
# public golden frames -> factors 0.70 / 0.60 / 0.50 / 0.50), rounded up to multiples of PATCH_SIZE. A frame whose
# pooled count exceeds a capacity is refused by the host check (``overflow``), never truncated silently. Measured
# peaks: T4 132,229 / 72,880 / 33,113 / 13,090 / 4,609 (SPEC 3), PandaSet <= 74,523 voxels, nuScenes 10,256-15,564.
# More (smaller) buckets: optimization phase.
STAGE_CAPACITIES = (256_000, 179_200, 107_520, 53_760, 27_136)
POOL_CHILDREN = 8                                # stride 2 in 3-D: <= 8 children per parent
BEV_CELL_CANDIDATES = 9                          # stage-3 voxels per BEV cell: z in 0..8 (z = 8 is out of grid)


@dataclass(frozen=True)
class PreprocessConfig:
    """``encoder.*`` of HF:ml_package_ptv3_encoder.param.yaml and AU@5eb9286:config/ptv3.param.yaml:7-12."""

    point_cloud_range: Tuple[float, ...] = (-122.88, -122.88, -3.0, 122.88, 122.88, 5.0)
    voxel_size: Tuple[float, ...] = (0.12, 0.12, 0.12)
    pooling_strides: Tuple[int, ...] = (2, 2, 2, 2)
    num_orders: int = NUM_ORDERS
    min_num_voxels: int = 16_000                 # voxels_num[0]: frames below are rejected (ptv3_trt.cpp:940-946)
    max_num_voxels: int = 256_000                # voxels_num[2]: clip to the first 256k in Morton order (:947-954)
    cloud_capacity: int = 2_000_000              # config/ptv3.param.yaml:7

    @property
    def coord_min(self) -> Tuple[int, int, int]:
        """``floor(min_range / voxel)`` per axis in float32 (ptv3_config.hpp:100-107): (-1024, -1024, -25)."""
        f32 = np.float32
        return tuple(int(np.floor(f32(self.point_cloud_range[i]) / f32(self.voxel_size[i]))) for i in range(3))

    @property
    def grid_size(self) -> Tuple[int, int, int]:
        """``grid_cells`` of ptv3_config.hpp:97-107: (2048, 2048, 67)."""
        f32 = np.float32
        out = []
        for i in range(3):
            lo, hi, vs = f32(self.point_cloud_range[i]), f32(self.point_cloud_range[i + 3]), f32(self.voxel_size[i])
            out.append(int(np.floor(np.nextafter(hi, lo) / vs) - np.floor(lo / vs)) + 1)
        return tuple(out)

    @property
    def serialization_depth(self) -> int:
        """``ceil(log2(max grid size))`` = 11 (ptv3_config.hpp:108-109)."""
        return int(np.ceil(np.log2(np.float32(max(self.grid_size)))))


@dataclass(frozen=True)
class DetectionConfig:
    """``detection3d.*`` of HF:ml_package_ptv3_det3d_head.param.yaml and AU@5eb9286:config/ptv3.param.yaml:40-86."""

    class_names: Tuple[str, ...] = DET_CLASSES
    bbox_voxel_size: Tuple[float, float, float] = (0.96, 0.96, 8.0)
    min_range: Tuple[float, float] = (-122.88, -122.88)          # encoder point_cloud_range min x / y
    has_twist: bool = True
    num_proposals: int = NUM_PROPOSALS
    post_center_range: Tuple[float, ...] = (-132.88, -132.88, -5.0, 132.88, 132.88, 12.0)
    distance_bin_upper_limits: Tuple[float, ...] = (50.0, 90.0, 121.0, 200.0)
    # [distance bin][class] (ptv3_node.cpp:154-172 lays them out bin-major)
    min_confidence_scores: Tuple[Tuple[float, ...], ...] = tuple((0.1,) * 7 for _ in range(4))
    yaw_norm_thresholds: Tuple[float, ...] = (0.1, 0.1, 0.1, 0.1, 0.0, 0.0, 0.1)
    iou_nms_search_distance_2d: float = 10.0
    iou_nms_threshold: float = 0.1
    # detection_class_remapper matrices (7 x 7, indexed by ObjectClassification label value: SPEC 5)
    allow_remapping_by_area_matrix: Tuple[int, ...] = tuple(1 if i == 1 else 0 for i in range(49))
    min_area_matrix: Tuple[float, ...] = tuple(12.1 if i == 1 else 0.0 for i in range(49))
    max_area_matrix: Tuple[float, ...] = tuple(36.0 if i == 1 else 0.0 for i in range(49))

    def score_thresholds(self) -> np.ndarray:
        """``(bins, classes)`` float32."""
        return np.asarray(self.min_confidence_scores, np.float32).reshape(len(self.distance_bin_upper_limits),
                                                                         len(self.class_names))

    def with_score_threshold(self, value: float) -> "DetectionConfig":
        """Every bin and class set to ``value`` (the API's ``score_threshold``)."""
        rows = tuple((float(value),) * len(self.class_names) for _ in self.distance_bin_upper_limits)
        return DetectionConfig(**{**self.__dict__, "min_confidence_scores": rows})


@dataclass(frozen=True)
class SegmentationConfig:
    """``segmentation3d.*`` of AU@5eb9286:config/ptv3.param.yaml:13-39 and the seg head yaml."""

    class_names: Tuple[str, ...] = SEG_CLASSES
    source_reconstruction: str = "full"
    filter_classes: Tuple[str, ...] = FILTER_CLASSES
    apply_filter_to_segmentation: bool = False

    def filter_indices(self) -> np.ndarray:
        return np.asarray([self.class_names.index(c) for c in self.filter_classes], np.int64)


@dataclass(frozen=True)
class PTv3Config:
    pre: PreprocessConfig = field(default_factory=PreprocessConfig)
    det: DetectionConfig = field(default_factory=DetectionConfig)
    seg: SegmentationConfig = field(default_factory=SegmentationConfig)
    neighbour_mode: str = DEFAULT_NEIGHBOUR_MODE

    def describe(self) -> Dict[str, object]:
        return {"weights": f"{WEIGHTS_REPO}@{WEIGHTS_TAG} ({WEIGHTS_REVISION[:12]})",
                "node": f"autoware_universe@{NODE_REVISION[:7]}", "neighbour_mode": self.neighbour_mode,
                "reconstruction": self.seg.source_reconstruction, "min_num_voxels": self.pre.min_num_voxels,
                "max_num_voxels": self.pre.max_num_voxels}


def _check() -> None:  # cheap invariants of the constants above (also run by the host tests)
    pre = PreprocessConfig()
    assert pre.coord_min == (-1024, -1024, -25) and pre.grid_size == (2048, 2048, 67)
    assert pre.serialization_depth == 11
    assert CLASS_ID_LUT.tolist() == [0, 1, 2, 4, 5, 7, 9, 7, 8, 8, 10, 9, 7, 9, 11]
    assert all(c % PATCH_SIZE == 0 for c in STAGE_CAPACITIES)


def stage_heads(stage: int) -> int:
    """Attention heads of encoder stage ``stage`` (C / 16)."""
    return ENC_CHANNELS[stage] // HEAD_DIM