# SPDX-License-Identifier: Apache-2.0
"""Compile-time constants of the METEOR port (``meteor_v157c3Z.onnx``, AutowareFoundation/meteor@v1.0).

Every value here is either a static shape of the deployed graph or a constant baked into it; the source of each is
the ONNX node it comes from (``research/meteor/SPEC.md``, cited as ``S:<line>``) or METEOR's runtime
(``tier4/METEOR`` @ ``dc193a81e995``, ``deploy/cpp``). Changing one of them means a different graph, a new trace and
a new image (PLAN.md section 0.4, class COMPILE). Host post-processing defaults (class RT-host) live in
``host/postprocess.py`` (``PostConfig``); load-time options (``input_norm``, ``depth_mean_bins``) in
:class:`MeteorConfig`.
"""
from __future__ import annotations

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

__all__ = [
    "ONNX_FILE", "ONNX_SHA256", "PARAM_YAML", "HF_REPO", "HF_TAG", "HF_REVISION", "METEOR_UPSTREAM_COMMIT",
    "CAMERAS", "N_CAMS", "IMG_H", "IMG_W", "FEAT_H", "FEAT_W", "FPN_C", "CTX_C", "BEV_C", "BEV_H", "BEV_W",
    "BEV_RES", "BEV_XF", "BEV_YH", "LIFT_H", "LIFT_W", "LIFT_RES", "DET_H", "DET_W", "DET_RES", "DEPTH_BINS",
    "DEPTH_MIN", "DEPTH_STEP", "DEPTH_LOG_MAX", "LANE_CLASSES", "SEG2D_CLASSES", "DET2D_CLASSES", "DET3D_CLASSES",
    "OCC_CLASSES", "TL_CLASSES", "DET2D_STRIDES", "PAINT_CLASSES", "OCC_CROP", "RISK_CROP", "EGO_MODES", "EGO_STEPS",
    "EGO_DT", "EGO_DIM", "TRAJ_C", "DONOR", "IMAGENET_MEAN", "IMAGENET_STD", "INPUT_NORMS", "DEPTH_MEAN_BINS",
    "MeteorConfig", "default_config", "depth_bin_centres",
]

# ---- artifact pins (S:99-118) ---------------------------------------------------------------------------------------
ONNX_FILE = "meteor_v157c3Z.onnx"
ONNX_SHA256 = "50397b4df76d33df2a71e8063b38d164399255008e1f0758c2a83675a200b1a7"
PARAM_YAML = "meteor_v157.param.yaml"
HF_REPO = "AutowareFoundation/meteor"
HF_TAG = "v1.0"
HF_REVISION = "01a5f6d71df5ecbbb5853ec600825481d57b9c6b"
METEOR_UPSTREAM_COMMIT = "dc193a81e9959de57b751e269127ae9b67e3f300"   # tier4/METEOR, defines pre/post (no Autoware pkg)

# ---- cameras and image branch (param.yaml:10, S:136-140) ------------------------------------------------------------
CAMERAS: Tuple[str, ...] = ("CAM_FRONT_WIDE", "CAM_FRONT_LEFT", "CAM_FRONT_RIGHT", "CAM_BACK_WIDE", "CAM_BACK_LEFT",
                            "CAM_BACK_RIGHT", "CAM_FRONT_NARROW", "CAM_BACK_NARROW")
N_CAMS = 8
IMG_H, IMG_W = 432, 768                 # network input per camera (the graph normalises u by 767, v by 431)
FEAT_H, FEAT_W = 108, 192               # stride-4 FPN map
FPN_C = 160                             # FPN width
CTX_C = 96                              # painted context channels lifted to the BEV
# the absent-camera donor rule of training (bevlane/dataset.py:17-21): zero image + the donor's K and pose
DONOR = {"CAM_BACK_NARROW": "CAM_BACK_WIDE", "CAM_FRONT_NARROW": "CAM_FRONT_WIDE"}

# ---- BEV (model.py:34-47, 82-89, 877-879; S:155-160) ----------------------------------------------------------------
BEV_C = 96
BEV_H, BEV_W = 800, 500                 # rows: x = +80 .. -80 m, cols: y = +50 .. -50 m
BEV_RES = 0.2
BEV_XF, BEV_YH = 80.0, 50.0             # forward extent, half width
LIFT_H, LIFT_W = 400, 250               # 0.4 m lift grid (V52, LIFT_DIV 2): x 79.8 .. -79.8, y 49.8 .. -49.8
LIFT_RES = 0.4
DET_H, DET_W = 400, 250                 # 3D detection / traj / stationary / risk grids
DET_RES = 0.4

# ---- depth bins (S:325-336; SPEC section 10 risk 2) -----------------------------------------------------------------
DEPTH_BINS = 64
DEPTH_MIN, DEPTH_STEP = 1.0, 1.25       # LINEAR bins of the lift and the trained head: 1.0 + 1.25 * b metres
DEPTH_LOG_MAX = 4.378896236419678       # ln(79.75) as the float32 constant of /Div_1: depth_mean's log-spaced centres
LIFT_MIN_Z = 0.5                        # camera z clamp and validity threshold
LIFT_MAX_DIST = 90.0                    # validity: Euclidean range from the camera < 90 m
LIFT_BIN_CLIP = 62.9999                 # clip of the continuous bin index (/net/Clip_4, float32 62.99990081787109)
LIFT_W_EPS = 0.05                       # + 0.05 on the interpolated bin probability (/net/Add_10)
LIFT_DEN_EPS = 1e-4                     # max(sum_n w, 1e-4) (/net/Clip_6)

# ---- heads -----------------------------------------------------------------------------------------------------
LANE_CLASSES = ("bg", "road", "sidewalk", "crosswalk", "laneline", "stopline", "road_edge", "marking", "parking")
SEG2D_CLASSES = ("bg", "obstacle", "car", "truck", "bus", "motorcycle", "bicycle", "pedestrian", "marking",
                 "traffic_light", "sign", "road", "sidewalk", "lane_stop_line", "crosswalk", "unused", "wall",
                 "building", "vegetation", "sky", "pole")          # comlops-21cls-autolabel-2504.csv (0, 15 unused)
DET2D_CLASSES = ("obs", "car", "truck", "bus", "bicycle", "motorcycle", "pedestrian", "road_paint", "traffic_light",
                 "traffic_sign")                                    # comlops-instance-2510.csv
DET3D_CLASSES = ("vehicle", "vru")
OCC_CLASSES = ("free", "obstacle", "vehicle", "two_wheeler", "pedestrian", "road", "sidewalk", "vegetation",
               "building_wall", "pole_sign_light")                 # extract_occ.py:4-8; 16 z-bins of 0.4 m from -1 m
OCC_Z_BINS = 16
TL_CLASSES = ("none", "green", "yellow", "red")
DET2D_STRIDES = (4, 8, 16)
PAINT_CLASSES = (2, 3, 4, 5, 6, 7, 8, 13)  # seg2d classes painted into the context (/net/Gather indices; checked)
SEG2D_CLAMP = 30.0                      # +-30 clamp, ONLY before the seg2d argmax (the paint softmax is unclamped)
OCC_CROP = ((200, 600), (50, 450))      # rows, cols of the RAW BEV (/net/Slice_8, Slice_9)
RISK_CROP = ((200, 600), (125, 375))    # rows, cols of the FUSED BEV (/net/Slice_11, Slice_12)
TGATE_SCALE = 4.0                       # fused = bev + tfuse(bev * 4 * softmax(tgate)[0]) (/net/Mul_12)

# ---- planner (S:315, 4.4) --------------------------------------------------------------------------------------
EGO_MODES, EGO_STEPS, EGO_DT = 3, 6, 0.5
EGO_DIM = 42                            # [0:36] 3 paths x 6 x (x, y); [36:39] mode logits; 39 steer, 40 accel, 41 brake
TRAJ_C = 39                             # per-cell agent futures: 3 modes x 6 x (dx, dy) + 3 mode logits
V0_SCALE = 15.0                         # kin_delta input [v0 / 15, 0, ..., 0] (7 wide)
KIN_DIM = 7
DEC_SPEED_MAX = 25.0                    # decoupled decoder: v = min(softplus(.), 25)
RISK_X_HALF, RISK_Y_HALF = 40.0, 25.0   # risk grid +-40 m x +-25 m at 0.2 m (400 x 250)
E2E_LN_EPS = 1e-5                       # LayerNormalization(139) epsilon (float32 9.99999974738e-06)
E2E_SCALE = 3.0                         # ego[:36] += 3 * tanh(mlp(LN(...)))
SEG_REFINE_CLAMP, SEG_REFINE_SCALE = 20.0, 8.0     # lane_logit = clamp(lane_pre, +-20) + 8 tanh(res / 8)
BOX_HM_CLAMP, BOX_REG_CLAMP, BOX_REFINE_SCALE = 15.0, 20.0, 6.0

# ---- input normalisation (D12; S:209-229, section 10 risk 1) ---------------------------------------------------------
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
INPUT_NORMS = ("onnx", "imagenet")      # onnx: x / 255 as released (default); imagenet: (x / 255 - mean) / std
DEPTH_MEAN_BINS = ("log", "linear")     # depth_mean bin centres: log as exported (default, parity) or linear (trained)


def depth_bin_centres(kind: str = "log"):
    """The 64 bin centres of ``depth_mean`` in metres: ``log`` = exp(r * ln(79.75) / 63) in float32 exactly as the
    graph (/Mul_1 -> /Exp), ``linear`` = 1 + 1.25 r (the bins the head and the lift were trained with)."""
    import numpy as np

    r = np.arange(DEPTH_BINS, dtype=np.float32)
    if kind == "log":
        step = np.float32(DEPTH_LOG_MAX) / np.float32(DEPTH_BINS - 1)
        return np.exp(r * step + np.float32(0.0)).astype(np.float32)
    if kind == "linear":
        return (np.float32(DEPTH_MIN) + np.float32(DEPTH_STEP) * r).astype(np.float32)
    raise ValueError(f"depth_mean bins must be one of {DEPTH_MEAN_BINS}, not {kind!r}")


@dataclass(frozen=True)
class MeteorConfig:
    """Load-time options of the reference (and of the TT port, ``METEOR_INPUT_NORM`` / ``METEOR_DEPTH_MEAN_BINS``).

    ``input_norm="onnx"`` reproduces the released graph (only /255; D12 default, validated against ORT);
    ``"imagenet"`` adds the training normalisation after /255 (SPEC section 10 risk 1; an in-memory patched ONNX is
    its oracle). ``depth_mean_bins="log"`` keeps ONNX parity, ``"linear"`` gives metric depth as trained."""

    input_norm: str = "onnx"
    depth_mean_bins: str = "log"
    cameras: Tuple[str, ...] = field(default=CAMERAS)

    def check(self) -> "MeteorConfig":
        if self.input_norm not in INPUT_NORMS:
            raise ValueError(f"input_norm must be one of {INPUT_NORMS}, not {self.input_norm!r}")
        if self.depth_mean_bins not in DEPTH_MEAN_BINS:
            raise ValueError(f"depth_mean_bins must be one of {DEPTH_MEAN_BINS}, not {self.depth_mean_bins!r}")
        if tuple(self.cameras) != CAMERAS:
            raise ValueError("the camera order is fixed by training (param.yaml:10)")
        return self


def default_config(**overrides) -> MeteorConfig:
    return MeteorConfig(**overrides).check()


# sanity: the log centres span 1 .. 79.75 m (S:183)
assert abs(math.exp(DEPTH_LOG_MAX) - 79.75) < 1e-4