# SPDX-License-Identifier: Apache-2.0
"""The Autoware ``lidar_frnet`` parameters this port reproduces, and the port's compile-time constants.

Sources (autoware_universe main ``9ceaccf``, ``perception/autoware_lidar_frnet`` = ``PKG/``; weights
``AutowareFoundation/lidar_frnet@v2.0`` = commit ``a53a1c11``; research/frnet/SPEC.md sections 1-3):

- per variant, from the weights repo: ``ml_package_frnet_<variant>.param.yaml`` (``fov_up_deg`` / ``fov_down_deg``,
  ``frustum_width`` / ``frustum_height`` (the network grid), ``interpolation_width`` / ``interpolation_height`` (the
  pre-processing grid), the 27 ``class_names``, the 81-int ``palette``, the ``num_points`` / ``num_unique_coors``
  [min, opt, max] profiles) and ``deploy_metadata.yaml`` (``version: v2.0``);
- node-wide, :data:`NODE_PARAMS` = ``PKG/config/frnet.param.yaml:3-16`` (filter threshold and classes, the filtered
  cloud's output format, the ego crop box).

Every validation the node does on these values is repeated here: palette size and 0-255 range
(``PKG/include/.../utils.hpp`` ``make_palette``), filter classes must be class names (``make_filter_class_indices``),
the crop box is enabled iff any float32 bound is non-zero (``PKG/src/lidar_frnet_node.cpp:51-66``). ``yaml`` is
imported when a file is read, never at module import; nothing here imports torch / onnx / ttnn.

The two variants are an inseparable tuple (weights, frustum H x W baked into the ONNX, FOV, interpolation grid:
S:frnet:535-538): a user selects ``ot128`` / ``qt128`` only (LOAD), every number below follows from it.
"""
from __future__ import annotations

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

import numpy as np

__all__ = [
    "VARIANTS", "DEFAULT_VARIANT", "DEPLOY_VERSION", "ONNX_SHA256", "NODE_PARAMS", "NUM_CLASSES", "STRIDES",
    "BLOCKS", "CROP_BOUND_KEYS", "FRNetConfig", "variant_files", "parse_crop_box", "parse_filter_classes",
]

VARIANTS = ("ot128", "qt128")       # PKG/launch/lidar_frnet.launch.xml:9-12 (sensor_model choices)
DEFAULT_VARIANT = "ot128"           # the launch default (Hesai OT128)
DEPLOY_VERSION = "v2.0"             # deploy_metadata.yaml of the pinned weights; Autoware's ansible pin (S:frnet:75-80)

# sha256 of the deployed graphs at the pinned revision (S:frnet:84-93): a different file is refused.
ONNX_SHA256 = {
    "ot128": "d29bf5950432360d800c22ff86a590d4d9d006acc501563ac611cfa33747c11f",
    "qt128": "647af92790e3d591a19d065fe94d3ded797916ac9ffe5bf96b5e8b8324b31ad2",
}

# PKG/config/frnet.param.yaml:3-16 (node parameters; RT-host in this port)
NODE_PARAMS: Dict[str, Any] = {
    "class_probability_threshold": 0.05,
    "filter_classes": ("drivable_surface",),
    "output_format": "",
    "ego_crop_box": {"reference_frame": "base_link", "min_x": -1.53, "min_y": -1.15, "min_z": -0.1,
                     "max_x": 6.1, "max_y": 1.15, "max_z": 3.1},
}
CROP_BOUND_KEYS = ("min_x", "min_y", "min_z", "max_x", "max_y", "max_z")

NUM_CLASSES = 27                    # conv_seg Gemm 64 -> 27; ml_package class_names (COMPILE)
STRIDES = (1, 2, 4, 8)              # AWML configs/model/frnet.py: backbone strides of layer1..4 (checked on the graph)
BLOCKS = (3, 4, 6, 3)               # BasicBlocks per stage (checked on the graph)
POSTPROCESS_MAX_CLASSES = 64        # PKG/lib/postprocess/postprocess_kernel.cu:32 (const_palette[64])
POSTPROCESS_MAX_FILTER_CLASSES = 16  # postprocess_kernel.cu:30 (const_filter_class_indices[16])


def variant_files(weights_dir: Any, variant: str) -> Dict[str, Path]:
    """The files of one variant in a weights directory (the flat HF repo layout)."""
    if variant not in VARIANTS:
        raise ValueError(f"unknown variant {variant!r}; choose one of {VARIANTS}")
    root = Path(weights_dir)
    return {"onnx": root / f"frnet_{variant}.onnx", "ml_package": root / f"ml_package_frnet_{variant}.param.yaml",
            "deploy_metadata": root / "deploy_metadata.yaml"}


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 _profile(values: Sequence[Any], name: str) -> Tuple[int, int, int]:
    v = tuple(int(x) for x in values)
    if len(v) != 3 or not (0 < v[0] <= v[1] <= v[2]):
        raise ValueError(f"{name} must be [min, opt, max] with 0 < min <= opt <= max, got {list(values)}")
    return v  # type: ignore[return-value]


@dataclass(frozen=True)
class FRNetConfig:
    """One variant's parameters. COMPILE values (device shapes): ``frustum_h`` x ``frustum_w``, ``n_cap``,
    ``m_cap``, ``num_classes``; LOAD: the variant itself; RT-host: the node parameters (filter, crop box)."""

    variant: str
    version: str
    fov_up_deg: float
    fov_down_deg: float
    frustum_w: int
    frustum_h: int
    interp_w: int
    interp_h: int
    class_names: Tuple[str, ...]
    palette: Tuple[int, ...]                       # 3 x num_classes ints, RGB per class
    num_points: Tuple[int, int, int]               # [min, opt, max] after interpolation (and of the raw input)
    num_unique: Tuple[int, int, int]               # [min, opt, max] unique frustum cells
    class_probability_threshold: float = NODE_PARAMS["class_probability_threshold"]
    filter_classes: Tuple[str, ...] = NODE_PARAMS["filter_classes"]
    crop_bounds: Tuple[float, ...] = tuple(NODE_PARAMS["ego_crop_box"][k] for k in CROP_BOUND_KEYS)
    crop_reference_frame: str = NODE_PARAMS["ego_crop_box"]["reference_frame"]
    model_dir: Optional[str] = field(default=None, compare=False)

    # ---- construction --------------------------------------------------------------------------------------
    @classmethod
    def from_model_dir(cls, model_dir: Any, variant: str = DEFAULT_VARIANT, *,
                       check_version: bool = True) -> "FRNetConfig":
        """Read ``ml_package_frnet_<variant>.param.yaml`` (and check ``deploy_metadata.yaml``) of a weights
        directory laid out as the HF repo."""
        files = variant_files(model_dir, variant)
        p = _ros_params(files["ml_package"])
        version = DEPLOY_VERSION
        if files["deploy_metadata"].is_file():
            import yaml

            meta = yaml.safe_load(files["deploy_metadata"].read_text()) or {}
            version = str(meta.get("version", ""))
            if check_version and version != DEPLOY_VERSION:
                raise ValueError(f"deploy_metadata.yaml version {version!r} != the pinned {DEPLOY_VERSION!r}")
        cfg = cls(variant=variant, version=version, fov_up_deg=float(p["fov_up_deg"]),
                  fov_down_deg=float(p["fov_down_deg"]), frustum_w=int(p["frustum_width"]),
                  frustum_h=int(p["frustum_height"]), interp_w=int(p["interpolation_width"]),
                  interp_h=int(p["interpolation_height"]), class_names=tuple(str(c) for c in p["class_names"]),
                  palette=tuple(int(v) for v in p["palette"]), num_points=_profile(p["num_points"], "num_points"),
                  num_unique=_profile(p["num_unique_coors"], "num_unique_coors"), model_dir=str(model_dir))
        cfg.validate()
        return cfg

    def validate(self) -> None:
        """The node's own checks (utils.hpp ``Dims2d``, ``make_palette``, ``make_filter_class_indices``)."""
        for name in ("frustum_w", "frustum_h", "interp_w", "interp_h"):
            if getattr(self, name) <= 0:
                raise ValueError(f"{name} must be positive")
        n = len(self.class_names)
        if n == 0 or n > POSTPROCESS_MAX_CLASSES:
            raise ValueError(f"{n} classes; the node supports 1..{POSTPROCESS_MAX_CLASSES}")
        if len(self.palette) % 3:
            raise ValueError("Palette size must be a multiple of 3.")
        if len(self.palette) != 3 * n:
            raise ValueError("Palette size does not match class names size.")
        if any(v < 0 or v > 255 for v in self.palette):
            raise ValueError("Color values must be within 0-255 range.")
        self.filter_class_indices()
        if len(self.crop_bounds) != 6:
            raise ValueError("ego crop box needs 6 bounds (min_x, min_y, min_z, max_x, max_y, max_z)")

    def with_node_params(self, **overrides: Any) -> "FRNetConfig":
        """A copy with other node parameters (``class_probability_threshold``, ``filter_classes``,
        ``crop_bounds``); re-validated."""
        cfg = replace(self, **overrides)
        cfg.validate()
        return cfg

    # ---- derived values (each computed as the node computes it) ---------------------------------------------
    @property
    def num_classes(self) -> int:
        return len(self.class_names)

    @property
    def fov(self) -> Tuple[np.float32, np.float32, np.float32]:
        """utils.hpp ``FieldOfView``: ``(up, down, total)`` radians, computed in double, stored as float32."""
        up = np.float32(np.pi * self.fov_up_deg / 180.0)
        down = np.float32(np.pi * self.fov_down_deg / 180.0)
        total = np.float32(np.pi * abs(self.fov_up_deg - self.fov_down_deg) / 180.0)
        return up, down, total

    @property
    def crop_bounds_f32(self) -> np.ndarray:
        """``static_cast<float>`` of the six bounds (lidar_frnet_node.cpp:53-58)."""
        return np.asarray(self.crop_bounds, np.float64).astype(np.float32)

    @property
    def crop_enabled(self) -> bool:
        """Enabled iff any float32 bound is non-zero (lidar_frnet_node.cpp:59-66)."""
        return bool(np.any(self.crop_bounds_f32 != 0.0))

    def filter_class_indices(self, names: Optional[Sequence[str]] = None) -> Tuple[int, ...]:
        """Indices of ``names`` (default: the configured filter classes) in ``class_names``; an unknown name raises
        ``ValueError`` (the node throws ``Filter class '<name>' not found in class names.``)."""
        names = self.filter_classes if names is None else tuple(names)
        if len(names) > POSTPROCESS_MAX_FILTER_CLASSES:
            raise ValueError(f"at most {POSTPROCESS_MAX_FILTER_CLASSES} filter classes")
        out = []
        for name in names:
            if name not in self.class_names:
                raise ValueError(f"Filter class '{name}' not found in class names.")
            out.append(self.class_names.index(name))
        return tuple(out)

    @property
    def palette_rgb(self) -> np.ndarray:
        """``(num_classes, 3)`` uint8."""
        return np.asarray(self.palette, np.int64).reshape(-1, 3).astype(np.uint8)

    @property
    def palette_packed(self) -> np.ndarray:
        """utils.hpp ``make_palette``: ``(r << 16) | (g << 8) | b`` reinterpreted as float32 (the visualization
        cloud's ``rgb`` field), ``(num_classes,)`` float32."""
        rgb = self.palette_rgb.astype(np.uint32)
        return ((rgb[:, 0] << 16) | (rgb[:, 1] << 8) | rgb[:, 2]).astype(np.uint32).view(np.float32)

    @property
    def n_cap(self) -> int:
        """Point capacity of the device graph (COMPILE): the profile maximum, so every frame Autoware accepts fits
        (functional port; smaller buckets are optimization work)."""
        return self.num_points[2]

    @property
    def m_cap(self) -> int:
        """Frustum-cell (voxel) capacity of the device graph (COMPILE); the gather tables add one zero row."""
        return self.num_unique[2]

    @property
    def frustum_cells(self) -> int:
        return self.frustum_h * self.frustum_w

    def research_cfg(self) -> Dict[str, Any]:
        """The same values in the layout of research/frnet/scripts/frnet_preproc_ref.py ``CFGS[variant]``."""
        return dict(name=self.variant, fov_up_deg=self.fov_up_deg, fov_down_deg=self.fov_down_deg,
                    frustum_w=self.frustum_w, frustum_h=self.frustum_h, interp_w=self.interp_w,
                    interp_h=self.interp_h, num_points=self.num_points, num_unique=self.num_unique)

    def describe(self) -> Dict[str, Any]:
        return {"variant": self.variant, "version": self.version, "fov_deg": [self.fov_up_deg, self.fov_down_deg],
                "frustum_hw": [self.frustum_h, self.frustum_w], "interpolation_hw": [self.interp_h, self.interp_w],
                "num_points": list(self.num_points), "num_unique_coors": list(self.num_unique),
                "num_classes": self.num_classes, "class_probability_threshold": self.class_probability_threshold,
                "filter_classes": list(self.filter_classes), "ego_crop_box": dict(zip(CROP_BOUND_KEYS,
                                                                                      self.crop_bounds)),
                "ego_crop_reference_frame": self.crop_reference_frame}


def parse_crop_box(value: Any) -> Optional[Tuple[float, ...]]:
    """An ego crop box override: ``None`` / ``False`` (disabled = all-zero bounds, the node's own switch), six
    numbers ``[min_x, min_y, min_z, max_x, max_y, max_z]``, or a mapping with those keys (missing keys keep
    Autoware's defaults). Raises ``ValueError`` on anything else."""
    if value is None or value is False:
        return (0.0,) * 6
    if isinstance(value, Mapping):
        unknown = sorted(set(value) - set(CROP_BOUND_KEYS) - {"reference_frame"})
        if unknown:
            raise ValueError(f"ego_crop_box: unknown keys {unknown}")
        if value.get("reference_frame", "base_link") != "base_link":
            raise ValueError("ego_crop_box.reference_frame must be base_link (the transform's target frame)")
        base = NODE_PARAMS["ego_crop_box"]
        vals = [value.get(k, base[k]) for k in CROP_BOUND_KEYS]
    else:
        vals = list(value)
        if len(vals) != 6:
            raise ValueError("ego_crop_box needs 6 numbers [min_x, min_y, min_z, max_x, max_y, max_z]")
    try:
        out = tuple(float(v) for v in vals)
    except (TypeError, ValueError):
        raise ValueError("ego_crop_box bounds must be numbers") from None
    if not all(np.isfinite(out)):
        raise ValueError("ego_crop_box bounds must be finite")
    return out


_IDENT = re.compile(r"^[A-Za-z0-9_+\-]+$")


def parse_filter_classes(value: Any) -> Tuple[str, ...]:
    """``filter_classes`` as a request gives it: a comma-separated string (``""`` = no filter class) or a list of
    names. Names are checked against the class list by :meth:`FRNetConfig.filter_class_indices`."""
    if value is None:
        return ()
    if isinstance(value, str):
        names = [v.strip() for v in value.split(",") if v.strip()]
    else:
        names = [str(v).strip() for v in value]
    bad = [n for n in names if not _IDENT.match(n)]
    if bad:
        raise ValueError(f"invalid filter class names {bad}")
    return tuple(names)