# SPDX-License-Identifier: Apache-2.0
"""What this port reproduces: Autoware VisionPilot's SceneSeg, i.e. the network the ``scene_seg_model`` node runs
(``SceneSeg_FP32.onnx``) with the node's pre- and post-processing, every constant with its source. Pure data (no
numpy / torch import); the CPU reference, the host pre/post-processing and the TT graph read their constants here.

Citations: ``S:sceneseg:N`` = research/sceneseg/SPEC.md line N (workspace); ``MR`` = vision_pilot @ ``POV_DEPLOY_SHA``
``VisionPilot/middleware_recipes``; ``PY`` = vision_pilot @ ``POV_CODE_SHA`` ``Models``; ``PLAN`` =
research/PLAN.md section 2.2.
"""
from __future__ import annotations

from typing import Dict, Tuple

# ------------------------------------------------------------------------------------------------- artifacts
# Upstream: the network is not part of autoware_universe; VisionPilot (autowarefoundation/vision_pilot, formerly
# autoware.privately-owned-vehicles) deployed it until 2026-06-15 and removed the PyTorch code on 2026-07-06.
POV_REPO = "autowarefoundation/vision_pilot"
POV_DEPLOY_SHA = "04fa3e80d0b0a9d2491c07991a9cb48101b2c2cb"   # last tree with middleware_recipes (S:sceneseg:12)
POV_CODE_SHA = "ca58cb50f5798cfba55c3696720d398bbcfce5fb"     # last tree with Models/ (S:sceneseg:11)
AW_PACKAGE = "models"                                         # MR/ROS2/models (ROS 2 package), node scene_seg_model
AW_PACKAGE_PATH = "VisionPilot/middleware_recipes/ROS2/models"
UPSTREAM_CARD = "AutowareFoundation/SceneSeg"                 # HF card only: README.md, no weights (S:sceneseg:137-141)
UPSTREAM_CARD_REVISION = "da7e600956d2f157fb41bddec34df4a9ab014925"

# The deployed weights: SceneSeg_FP32.onnx (MR/README.md:107-109, autoseg.yaml model_path), BN folded at export.
ONNX_FILE = "SceneSeg_FP32.onnx"
ONNX_SHA256 = "8e509094acd43a0ebfbc7388fffd5177a4098827b644df5af5f064c07c4bc973"
ONNX_BYTES = 193_467_193
ONNX_DRIVE_ID = "1l-dniunvYyFKvLD7k16Png3AsVTuMl9f"            # PY/model_library/SceneSeg/README.md:20
# The PyTorch state_dict the ONNX was exported from (identical weights, S:sceneseg:34): source of truth, re-hosted as
# safetensors next to the ONNX (PLAN D10).
PTH_FILE = "iter_140215_epoch_4_step_15999.pth"
PTH_SHA256 = "9482f13130df3a0c994ed1b254a2601aec2c05eabc8eea078057e1dd7a7689fc"
PTH_BYTES = 193_782_170
PTH_DRIVE_ID = "1vCZMdtd8ZbSyHn1LCZrbNKMK7PQvJHxj"             # PY/model_library/SceneSeg/README.md:18
SAFETENSORS_FILE = "sceneseg_iter_140215_epoch_4_step_15999.safetensors"   # code/scripts/make_weights_upload.py
SAFETENSORS_SHA256 = "94297835fe5e461026f4c30086a2a90cbcc2054076f9675b70a9040bb4e060d2"   # reproducible conversion

# Re-hosted weights (PLAN D10, user decision): inside this bundle's own HF repo, under weights/. The revision pin is
# api.SceneSeg.DEFAULT_REVISION (rendered from AW_SHA, = tt-model.yaml weights.revision): the commit of that upload,
# filled when it exists (PLAN 4.1: it gates the publish, not the port).
WEIGHTS_REPO = "changh95/sceneseg-p150"
WEIGHTS_SUBDIR = "weights"
ALLOW_PATTERNS = [f"{WEIGHTS_SUBDIR}/{ONNX_FILE}", f"{WEIGHTS_SUBDIR}/provenance.json", f"{WEIGHTS_SUBDIR}/LICENSE",
                  f"{WEIGHTS_SUBDIR}/NOTICE"]
WEIGHTS_LICENSE = ("Apache-2.0 (autowarefoundation/vision_pilot LICENSE and README.md:41-42, which cover the AI model "
                   "weights; the AutowareFoundation/SceneSeg card says apache-2.0; S:sceneseg:156-159)")

# Graph I/O of the ONNX (S:sceneseg:244-245)
ONNX_INPUT = "input"            # float32 [batch_size, 3, 320, 640]; only batch 1 runs (the context Reshape)
ONNX_OUTPUT = "output"          # float32 [batch_size, 3, 320, 640] logits

# ---------------------------------------------------------------------------------------- network (COMPILE)
INPUT_HW: Tuple[int, int] = (320, 640)        # fixed: the SceneContext reshape hard-codes 10 x 20 (S:sceneseg:613)
NUM_CLASSES = 3
CONTEXT_HW: Tuple[int, int] = (10, 20)        # stride-32 grid = the reshape constant of /SceneContext/Reshape
BN_EPS = 1e-5                                 # torchvision BatchNorm2d default; verified against the ONNX fold
# torchvision efficientnet_b0 stages (expand_ratio, kernel, stride, in, out, layers), SE ratio 0.25 of the block input
B0_SETTINGS: Tuple[Tuple[int, int, int, int, int, int], ...] = (
    (1, 3, 1, 32, 16, 1),
    (6, 3, 2, 16, 24, 2),
    (6, 5, 2, 24, 40, 2),
    (6, 3, 2, 40, 80, 3),
    (6, 5, 1, 80, 112, 3),
    (6, 5, 2, 112, 192, 4),
    (6, 3, 1, 192, 320, 1),
)
STEM_CHANNELS = 32
HEAD_CHANNELS = 1280
CONTEXT_MLP: Tuple[int, ...] = (1280, 800, 800, 200)          # SceneContext Linear sizes (200 = 10 x 20)
CONTEXT_CONVS: Tuple[int, ...] = (1, 128, 256, 512, 1280)       # context_layer_3..6 channels at 10 x 20
# Decoder blocks: (upsample module, cin, skip module, skip cin, decode modules with (cin, cout))
NECK_BLOCKS = (
    ("upsample_layer_0", 1280, "skip_link_layer_0", 80, (("decode_layer_0", 1280, 768), ("decode_layer_1", 768, 768))),
    ("upsample_layer_1", 768, "skip_link_layer_1", 40, (("decode_layer_2", 768, 512), ("decode_layer_3", 512, 512))),
    ("upsample_layer_2", 512, "skip_link_layer_2", 24, (("decode_layer_4", 512, 512), ("decode_layer_5", 512, 256))),
)
HEAD_BLOCKS = (
    ("upsample_layer_3", 256, "skip_link_layer_3", 32, (("decode_layer_6", 256, 256), ("decode_layer_7", 256, 128))),
    ("upsample_layer_4", 128, None, 0, (("decode_layer_8", 128, 128), ("decode_layer_9", 128, 64))),
)
FINAL_CONV = ("decode_layer_10", 64, NUM_CLASSES)

# ------------------------------------------------------------------------ pre-processing (host, MR backends)
# MR/common/backends/onnx_runtime_backend.cpp:41-60 == tensorrt_backend.cpp:160-177 (S:sceneseg 3): cv_bridge bgr8,
# cv::resize to 640x320 (INTER_LINEAR, aspect not kept), convertTo 1/255, subtract / divide the BGR-ordered ImageNet
# constants, split into B, G, R planes. Implemented bit-exactly by the vendored ttaw.image_linear (C14).
MEAN_RGB: Tuple[float, float, float] = (0.485, 0.456, 0.406)
STD_RGB: Tuple[float, float, float] = (0.229, 0.224, 0.225)
# Channel order of the network input: LOAD-time variant (PLAN 2.2, D10): "bgr" = deployed (default), "rgb" = training
# order (better on two-wheelers: research/sceneseg/public_data/README.md section 5).
CHANNEL_ORDERS: Tuple[str, ...] = ("bgr", "rgb")
DEFAULT_CHANNEL_ORDER = "bgr"
RESIZE_INTERPOLATION = "bilinear"   # cv::resize default (training also used cv2 INTER_LINEAR; S:sceneseg:196)

# ------------------------------------------------------------------------- post-processing (host, MR node)
# MR/ROS2/models/src/run_model_node.cpp:117-188 (S:sceneseg 5): per-pixel argmax over the 3 logits, scanning c = 0, 1,
# 2 from max_score = -1e9 with a strict '>' (ties -> lowest index); mask = 255 where the class is 1, else 0;
# cv::resize(mask, source size, INTER_NEAREST); published as mono8 with the input header.
CLASS_NAMES: Tuple[str, ...] = ("background", "foreground", "road")   # load_data_scene_seg.py:136-141 (bg incl. sky)
FOREGROUND_CLASS = 1
MASK_VALUE = 255
ARGMAX_INIT = -1e9                  # run_model_node.cpp: float max_score = -1e9
INPUT_TOPIC = "/sensors/video/image_raw"       # autoseg.yaml (informational)
OUTPUT_TOPIC = "/autoseg/scene_seg/mask"

# --------------------------------------------------------------------------------- per-module taps (goldens)
# Stable names shared by the CPU reference (reference/model.py), the goldens and the TT graph, with the ONNX tensor
# that holds the same value (S:sceneseg 4.2; the l0 / l2 / l3 / l4 / l8 / context / neck / logits names are the
# research taps of research/sceneseg/scripts/crosscheck.py).
TAP_ONNX_TENSORS: Dict[str, str] = {
    "l0": "/Backbone/encoder.0/encoder.0.2/Mul_output_0",                              # 32 x 160 x 320
    "l1": "/Backbone/encoder.1/encoder.1.0/block/block.2/block.2.0/Conv_output_0",     # 16 x 160 x 320
    "l2": "/Backbone/encoder.2/encoder.2.1/Add_output_0",                              # 24 x 80 x 160
    "l3": "/Backbone/encoder.3/encoder.3.1/Add_output_0",                              # 40 x 40 x 80
    "l4": "/Backbone/encoder.4/encoder.4.2/Add_output_0",                              # 80 x 20 x 40
    "l5": "/Backbone/encoder.5/encoder.5.2/Add_output_0",                              # 112 x 20 x 40
    "l6": "/Backbone/encoder.6/encoder.6.3/Add_output_0",                              # 192 x 10 x 20
    "l7": "/Backbone/encoder.7/encoder.7.0/block/block.3/block.3.0/Conv_output_0",     # 320 x 10 x 20
    "l8": "/Backbone/encoder.8/encoder.8.2/Mul_output_0",                              # 1280 x 10 x 20
    "context.mean": "/SceneContext/ReduceMean_output_0",                               # [1, 1280]
    "context.mlp": "/SceneContext/sigmoid/Sigmoid_output_0",                           # [1, 200]
    "context.gate": "/SceneContext/GeLU_5/Mul_1_output_0",                             # 1280 x 10 x 20
    "context": "/SceneContext/Add_output_0",                                           # gate * l8 + l8
    "neck.up0": "/SceneNeck/Add_output_0",                                             # 1280 x 20 x 40
    "neck.dec1": "/SceneNeck/GeLU_1/Mul_1_output_0",                                   # 768 x 20 x 40
    "neck.up1": "/SceneNeck/Add_1_output_0",                                           # 768 x 40 x 80
    "neck.dec3": "/SceneNeck/GeLU_3/Mul_1_output_0",                                   # 512 x 40 x 80
    "neck.up2": "/SceneNeck/Add_2_output_0",                                           # 512 x 80 x 160
    "neck": "/SceneNeck/GeLU_5/Mul_1_output_0",                                        # 256 x 80 x 160
    "head.up3": "/SceneSegHead/Add_output_0",                                          # 256 x 160 x 320
    "head.dec7": "/SceneSegHead/GeLU_1/Mul_1_output_0",                                # 128 x 160 x 320
    "head.up4": "/SceneSegHead/upsample_layer_4/ConvTranspose_output_0",               # 128 x 320 x 640
    "head.dec8": "/SceneSegHead/GeLU_2/Mul_1_output_0",                                # 128 x 320 x 640
    "head.dec9": "/SceneSegHead/GeLU_3/Mul_1_output_0",                                # 64 x 320 x 640
    "logits": ONNX_OUTPUT,                                                             # 3 x 320 x 640
}
TAPS: Tuple[str, ...] = tuple(TAP_ONNX_TENSORS)
# Module boundaries of the TT port (PORT_LOG.md milestone plan): the taps every module gate compares, and the inputs
# a teacher-forced module case is fed. Stored for the gate samples (research/sceneseg/goldens).
MODULE_TAPS: Tuple[str, ...] = ("l0", "l1", "l2", "l3", "l4", "l5", "l6", "l7", "l8", "context", "neck.dec1",
                                "neck.dec3", "neck", "head.dec7", "logits")
# Taps above ~26 M values (100 MB fp32): computed on demand, never stored
LARGE_TAPS: Tuple[str, ...] = ("head.up4", "head.dec8", "head.dec9")