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

``AW/`` = autoware_universe ``perception/autoware_tensorrt_bevformer`` @ 9ceaccf (package 0.53.0); ``DH/`` =
DerryHub/BEVFormer_tensorrt @ 303d3140 (the ``forward_trt`` export graph Autoware's ONNX came from); ``BF/`` =
fundamentalvision/BEVFormer @ 66b65f3a (training); ``CFG`` = ``DH/configs/bevformer/bevformer_small.py`` (identical
to ``BF/projects/configs/bevformer/bevformer_small.py``); ``S:<line>`` = ``research/bevformer/SPEC.md``.

Parameter classes (PLAN.md section 0.4): every constant here is COMPILE (it fixes a device shape or the weights'
meaning) unless its comment says RT-host; the RT-host values are only the defaults of ``api.BEVFormer``'s runtime
parameters. Pure Python + numpy: importing this module never imports torch / ttnn.
"""
from __future__ import annotations

from typing import Dict, Tuple

import numpy as np

__all__ = [
    "CAMERA_ORDER", "CLASS_NAMES", "AUTOWARE_LABELS", "AUTOWARE_LABEL_OF", "PC_RANGE", "POST_CENTER_RANGE", "BEV_H",
    "BEV_W", "NUM_QUERY_BEV", "EMBED", "NUM_HEADS", "HEAD_DIM", "FFN_DIM", "ENC_LAYERS", "DEC_LAYERS",
    "TSA_POINTS", "TSA_QUEUE", "SCA_POINTS", "PILLAR_POINTS", "DEC_POINTS", "NUM_QUERY", "NUM_CLASSES", "CODE_SIZE",
    "NUM_CAMS", "NODE_HW", "SCALE", "PAD_DIVISOR", "IMG_HW", "FEAT_HW", "FEAT_LEN", "MEAN_BGR", "STD", "TO_RGB",
    "ROTATE_CENTER", "GRID_LENGTH", "Z_ANCHORS", "BN_EPS", "LN_EPS", "MSDA_EPS", "INV_SIG_EPS", "SCORE_THRESHOLD",
    "MAX_NUM", "SCA_CAPACITY", "SCA_MAX_CAMS_PER_QUERY", "RESNET_LAYERS", "RESNET_WIDTHS", "RESNET_STRIDES",
    "RESNET_DCN", "backbone_conv_specs", "ConvSpecT",
]

# ---- cameras, classes, labels ---------------------------------------------------------------------------------
# Network camera order: the nuScenes info order the node builds its inputs in (AW/src/bevformer_node.cpp:398-399,
# 526-527, 570). It is baked into the learned cams_embeds[6] (one embedding per slot), so the host must permute
# its cameras into exactly this order (S:73).
CAMERA_ORDER: Tuple[str, ...] = ("CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_FRONT_LEFT", "CAM_BACK", "CAM_BACK_LEFT",
                                 "CAM_BACK_RIGHT")
NUM_CAMS = 6
# Network class order = AW/config/bevformer.param.yaml post_process_params.class_names = CFG class_names
CLASS_NAMES: Tuple[str, ...] = ("car", "truck", "construction_vehicle", "bus", "trailer", "barrier", "motorcycle",
                                "bicycle", "pedestrian", "traffic_cone")
# The ObjectClassification label the node publishes per class name (AW/src/ros_utils.cpp:70-94): bus, trailer and
# construction_vehicle become TRUCK (BEVDet in the same repo emits BUS / TRAILER), everything unlisted UNKNOWN.
AUTOWARE_LABEL_OF: Dict[str, str] = {"car": "CAR", "truck": "TRUCK", "construction_vehicle": "TRUCK",
                                     "bus": "TRUCK", "trailer": "TRUCK", "bicycle": "BICYCLE",
                                     "motorcycle": "MOTORCYCLE", "pedestrian": "PEDESTRIAN"}
AUTOWARE_LABELS: Tuple[str, ...] = tuple(AUTOWARE_LABEL_OF.get(c, "UNKNOWN") for c in CLASS_NAMES)
NUM_CLASSES = 10
CODE_SIZE = 10   # (cx, cy, log w, log l, cz, log h, sin, cos, vx, vy)

# ---- geometry (all fixed by training: S:96-104, S:390-392) ----------------------------------------------------
PC_RANGE: Tuple[float, ...] = (-51.2, -51.2, -5.0, 51.2, 51.2, 3.0)       # CFG point_cloud_range
POST_CENTER_RANGE: Tuple[float, ...] = (-61.2, -61.2, -10.0, 61.2, 61.2, 10.0)   # AW yaml (RT-host filter)
BEV_H = BEV_W = 150                                                        # CFG bev_h_ / bev_w_
NUM_QUERY_BEV = BEV_H * BEV_W                                              # 22,500 BEV queries (row q = y * 150 + x)
GRID_LENGTH = (PC_RANGE[4] - PC_RANGE[1]) / BEV_H                          # 0.68267 m (DH transformer.py:262-294)
# Pillar z-anchors of the spatial cross-attention: linspace(0.5, 7.5, 4) / 8 * 8 - 5 m (DH encoder.py:169-195; S:419)
Z_ANCHORS: Tuple[float, ...] = tuple(float(z) for z in np.linspace(0.5, 7.5, 4) / 8.0 * 8.0 + PC_RANGE[2])
# Centre of the prev_bev rotation: the PerceptionTransformer default (100, 100) px, NOT the centre of this 150x150
# grid; the model was trained with it (DH transformer.py:26, BF transformer.py:50; S:418)
ROTATE_CENTER: Tuple[float, float] = (100.0, 100.0)

# ---- images (AW/config/bevformer.param.yaml data_params; S:116-125) -----------------------------------------------
NODE_HW: Tuple[int, int] = (900, 1600)          # input_height / input_width: the node resizes every image to it
SCALE = 0.8                                     # scale_factor
PAD_DIVISOR = 32                                # pad_divisor
IMG_HW: Tuple[int, int] = (736, 1280)           # network input after x0.8 and the /32 pad (UV normaliser)
FEAT_HW: Tuple[int, int] = (23, 40)             # C5 / FPN map (stride 32): the SCA value is 6 x 920 tokens
FEAT_LEN = FEAT_HW[0] * FEAT_HW[1]
MEAN_BGR: Tuple[float, float, float] = (103.530, 116.280, 123.675)   # BGR, caffe-style weights
STD: Tuple[float, float, float] = (1.0, 1.0, 1.0)
TO_RGB = False

# ---- transformer (CFG lines 38-187) ------------------------------------------------------------------------------
EMBED = 256              # _dim_
NUM_HEADS = 8
HEAD_DIM = EMBED // NUM_HEADS
FFN_DIM = 512            # _ffn_dim_
ENC_LAYERS = 3
DEC_LAYERS = 6           # nb_dec: outputs_classes / outputs_coords hold all 6 layers; the node uses the last
TSA_POINTS = 4           # TemporalSelfAttention num_points (1 level)
TSA_QUEUE = 2            # num_bev_queue: [prev_bev; current] (baked into the sampling_offsets weights)
SCA_POINTS = 8           # MSDeformableAttention3D num_points (point p uses pillar anchor p % 4)
PILLAR_POINTS = 4        # num_points_in_pillar
DEC_POINTS = 4           # CustomMSDeformableAttention num_points
NUM_QUERY = 900

# ---- numerics of the deployed graph ---------------------------------------------------------------------------------
BN_EPS = 1e-5            # frozen BN2d (CFG norm_cfg BN2d, eps default)
LN_EPS = 1e-5            # nn.LayerNorm default (the ttnn default is 1e-12: pass it explicitly)
MSDA_EPS = 1e-5          # point sampling: depth > 1e-5, uv / max(depth, 1e-5) (DH encoder.py:197-259)
INV_SIG_EPS = 1e-5       # inverse_sigmoid clamps (DH decoder.py:24-37; mmdet inverse_sigmoid in the head)

# ---- post-processing (RT-host defaults; AW/config/bevformer.param.yaml, AW/src/bevformer_node.cpp:651) --------------
SCORE_THRESHOLD = 0.2    # score_thre
MAX_NUM = 300            # max_num (get_parameter_or default; CFG bbox_coder max_num)

# ---- spatial cross-attention rebatching (the TT port; PLAN.md 2.11, P11) ---------------------------------------------
# Static per-camera capacity of the rebatched (camera, query) rows. The SPEC's 5,376 (S:335, from sample0's maximum
# 5,371) is too small: nuScenes mini scene-0103 reaches 5,460 (CAM_BACK) and the shipped PandaSet rig 5,704
# (CAM_BACK_RIGHT). 6,144 = 192 tiles leaves 7.7 % margin over every rig measured (PORT_LOG.md decision E7); a
# calibration above it uses the dense SCA path. A fixed constant, never derived from the core grid (PLAN.md 0.2).
SCA_CAPACITY = 6144
# Most cameras that see one BEV query (gather-form scatter-add of the SCA output): 2 on nuScenes rigs, 3 allowed
SCA_MAX_CAMS_PER_QUERY = 3

# ---- backbone: ResNet-101, caffe style, DCNv2 in stages 3-4 (CFG lines 49-64) ----------------------------------------
RESNET_LAYERS: Tuple[int, ...] = (3, 4, 23, 3)
RESNET_WIDTHS: Tuple[int, ...] = (64, 128, 256, 512)       # bottleneck width; output = 4x
RESNET_STRIDES: Tuple[int, ...] = (1, 2, 2, 2)
RESNET_DCN: Tuple[bool, ...] = (False, False, True, True)   # stage_with_dcn: conv2 is DCNv2 (deform_groups 1)

# (module, kind, stride, padding, bn module) of every backbone conv, in graph order. kind: "conv" (BN folded),
# "dcn" (the DCNv2 main weight, BN folded, no conv bias) or "offset" (the DCN offset conv: bias, no BN).
ConvSpecT = Tuple[str, str, int, int, str]


def backbone_conv_specs(prefix: str = "img_backbone") -> Tuple[ConvSpecT, ...]:
    """Every convolution of the backbone. Caffe style: the stride sits on the 1x1 ``conv1`` of the first block of a
    stage (mmdet ResNet ``style='caffe'``), the 3x3 ``conv2`` (plain or DCN) keeps stride 1 / padding 1."""
    specs = [(f"{prefix}.conv1", "conv", 2, 3, f"{prefix}.bn1")]
    for s, (n, stride, dcn) in enumerate(zip(RESNET_LAYERS, RESNET_STRIDES, RESNET_DCN)):
        for b in range(n):
            p = f"{prefix}.layer{s + 1}.{b}"
            st = stride if b == 0 else 1
            specs.append((f"{p}.conv1", "conv", st, 0, f"{p}.bn1"))
            if dcn:
                specs.append((f"{p}.conv2.conv_offset", "offset", 1, 1, ""))
                specs.append((f"{p}.conv2", "dcn", 1, 1, f"{p}.bn2"))
            else:
                specs.append((f"{p}.conv2", "conv", 1, 1, f"{p}.bn2"))
            specs.append((f"{p}.conv3", "conv", 1, 0, f"{p}.bn3"))
            if b == 0:
                specs.append((f"{p}.downsample.0", "conv", st, 0, f"{p}.downsample.1"))
    return tuple(specs)