# SPDX-License-Identifier: Apache-2.0
"""Python API: FRNet (Autoware lidar_frnet) on one Tenstorrent Blackhole p150.

    from tt_frnet import FRNet

    with FRNet.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model("code/tt_frnet/samples/synthetic_ot128.npz", calibration={"preset": "synthetic_ot128"})
    print(out.to_dict())  # the same JSON as POST /predict

The contract shared by every bundle of the Autoware collection is the vendored ``ttaw.api_base.ModelBase``
(BUNDLE_CONVENTIONS.md section 8): ``from_pretrained`` resolves the pinned weights before it claims the chip, opens it
(ETH dispatch, 12x10), builds the graph and captures every trace variant in ``warmup_variants``, so the first call is
as fast as the later ones; calls are serialised by a lock (one chip, batch 1) and fill ``timing_ms``; ``close()`` is
idempotent, also runs at interpreter exit, and closes the chip only if the model opened it. The HTTP server
(``tt_frnet.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.

This module holds the identity of the port and its hooks: host pre- and post-processing ported from Autoware live in
``tt_frnet.host`` (shared with the CPU reference ``tt_frnet.reference``, so device and reference see identical
network inputs), the ttnn graph in ``tt_frnet.tt`` (run by a ``ttaw.trace.TraceRunner``; milestones M1-M3 of
PORT_LOG.md). Importing it has no side effects: ttnn / torch are imported inside the hooks.

Inputs (``host/frame.py`` has the full contract): ``points`` -- ONE sweep in the **sensor frame**, fields x, y, z,
intensity (raw 0-255); ``calibration`` -- ``{"T_base_link_from_sensor": ...}`` (or ``{"preset": name}`` of
``calib/``) for the node's ego crop box, ``ego_crop_box`` to override or disable the box (``None``).
``variant``: ``ot128`` (Autoware's default sensor model) or ``qt128``, a load-time choice (weights + grid).

Output: :class:`FRNetSegmentation` (a ``ttaw.outputs.Segmentation3D``) in the node's output order: the kept raw
points (input order) then Autoware's interpolated points; ``label_ids`` = class_id, ``scores`` = probability,
``points``, ``filter_keep`` (the filtered cloud), ``input_index``. A frame outside the node's profiles gives an empty
output whose ``meta["skipped"]`` is the node's message (the node publishes nothing).
"""
from __future__ import annotations

from typing import Any, Dict, Optional

from . import io as tio
from .device import DEVICE_DEFAULTS
from .host.postprocess import RUNTIME_PARAMS, FRNetSegmentation
from .reference.config import DEFAULT_VARIANT, VARIANTS
from .ttaw.api_base import ModelBase

__all__ = ["FRNet", "Output", "FRNetSegmentation", "LABELS"]

# The result class of this model (ttaw.outputs: Detections3D, Detections2D, Segmentation3D, Mask2D, Trajectory).
Output = FRNetSegmentation

# ml_package_frnet_{ot128,qt128}.param.yaml class_names (identical in both variants; _build_host checks the yaml).
# Indices 22 / 23 read "debris" / "stroller" here, but the weights were trained with stroller = 22, debris = 23
# (AWML frnet_1xb8_t4dataset-ot128-seg.py:60-61; S:frnet:314-317): Autoware's names are kept for parity.
LABELS = ("drivable_surface", "other_flat_surface", "sidewalk", "manmade", "vegetation", "car", "bus",
          "emergency_vehicle", "train", "truck", "tractor_unit", "semi_trailer", "construction_vehicle", "forklift",
          "kart", "motorcycle", "bicycle", "pedestrian", "personal_mobility", "animal", "pushable_pullable",
          "traffic_cone", "debris", "stroller", "other_stuff", "noise+ghost_point", "unknown")


def _given(value: Any) -> bool:
    """A real input (``None`` and empty containers are 'not given'); never the truth value of an array."""
    if value is None:
        return False
    if isinstance(value, (list, tuple, dict, str)):
        return len(value) > 0
    return True


class _DeviceRows:
    """``_forward``'s result with ``FRNET_HOST_POST=tiles`` / ``numba``: the device's ``pred_probs`` in sorted order
    (padded; raw TILE-layout elements with ``width``, else untilized rows) and the frame plan; ``_postprocess``
    un-permutes them in its compiled pass."""

    __slots__ = ("rows", "plan", "width")

    def __init__(self, rows: Any, plan: Any, width: Any = None):
        import numpy as np

        r = np.asarray(rows, np.float32)
        self.width = width                     # None: untilized rows; else the raw TILE elements of this width
        self.rows = r.reshape(-1) if width is not None else r.reshape(-1, r.shape[-1])
        self.plan = plan


class FRNet(ModelBase):
    """FRNet (Autoware lidar_frnet) on one Blackhole p150. Create it with :meth:`from_pretrained`."""

    MODEL_NAME = "frnet-p150"
    ENV_PREFIX = "FRNET"  # prefix of the environment knobs (SERVING.md section 3.4)
    DEFAULT_REPO = "AutowareFoundation/lidar_frnet"
    DEFAULT_TAG = "v2.0"  # the Autoware ansible artifacts pin
    # the commit DEFAULT_TAG points to (pinned: tags can move)
    DEFAULT_REVISION = "a53a1c11b0b28d31fd50a03abef5db0e830a5b3c"
    ALLOW_PATTERNS = ["frnet_*.onnx", "ml_package_frnet_*.param.yaml", "deploy_metadata.yaml"]
    VARIANTS = list(VARIANTS)  # load-time: selects the ONNX, the frustum grid and the trace shapes
    DEFAULT_VARIANT = DEFAULT_VARIANT
    INPUT_KIND = "lidar"  # lidar | camera | multicam | lidar+multicam | planner
    CAMERA_ORDER = ()  # multicam models: the fixed camera order of the network
    POINT_FIELDS = tio.DEFAULT_POINT_FIELDS  # columns the network consumes, in order (x, y, z, intensity)
    LABELS = list(LABELS)
    # Per-request knobs (host post-processing only; defaults are Autoware's): host/postprocess.py
    RUNTIME_PARAMS = RUNTIME_PARAMS
    EXTRA_INPUTS = ()
    DEVICE_DEFAULTS = DEVICE_DEFAULTS  # validated open parameters (device.py)

    # ---- port-specific hooks (called by ModelBase; keep host work out of _forward) ---------------------------
    def _build(self) -> None:
        """Host side (config, weights, pre- / post-processors), then the ttnn graph of ``tt_frnet.tt``
        (``TtFRNet``: point encoder, K1 segmented max, gathers, convs, decoder, head) as the ``"default"`` variant of
        a ``ttaw.trace.TraceRunner`` with the per-frame tables as persistent inputs (no capture here; ``_warm_one``
        captures). Warm-up runs an empty frame (every index valid)."""
        from .tt.model import FrameStaging, TtFRNet, bucket_name, build_runner, default_policy

        self._build_host()
        self.graph = TtFRNet(self.device, self.frnet_weights, self.config, policy=default_policy())
        # BUCKETS (OPT round 3): the same graph at smaller N_cap / M_cap, one trace each; a frame runs in the
        # smallest one that holds it (select())
        self.buckets = [self.graph.view(n, m) for n, m in self.graph.bucket_caps()]
        self.runner = build_runner(self.graph, variants=("default",), name=f"{self.MODEL_NAME}/{self.variant}",
                                   buckets=self.buckets)
        self.staging = FrameStaging(self.graph)   # persistent host buffers of the inputs (written in place)
        self.stagings = [(bucket_name(v), FrameStaging(v, suffix="@" + bucket_name(v))) for v in self.buckets]
        self.stagings.append(("default", self.staging))
        self._host_outs: Dict[str, Any] = {}

    def select(self, plan: Any):
        """``(trace variant, FrameStaging)`` of the smallest capacity that holds the frame (``"default"``: N_cap)."""
        for name, st in getattr(self, "stagings", [("default", getattr(self, "staging", None))]):
            if name == "default" or st.fits(plan):
                return name, st
        return "default", self.staging

    def _build_host(self) -> None:
        """The host half of ``_build`` (no device): usable alone by host tests and tools."""
        from .host.frame import FRNetPreprocessor
        from .host.postprocess import FRNetPostprocessor
        from .reference.config import FRNetConfig
        from .reference.weights import check_variant_grid, load_weights

        self.config = FRNetConfig.from_model_dir(self.weights_path, self.variant)
        if tuple(self.config.class_names) != tuple(self.LABELS):
            raise ValueError(f"ml_package class_names differ from FRNet.LABELS: {self.config.class_names}")
        self.frnet_weights = load_weights(self.weights_path, self.variant)
        check_variant_grid(self.frnet_weights, self.config.frustum_h, self.config.frustum_w)
        from .tt.model import KNOBS

        knobs = KNOBS.read()
        self.pre = FRNetPreprocessor(self.config, pre_norm=(self.frnet_weights.pre_norm_scale,
                                                            self.frnet_weights.pre_norm_shift),
                                     backend=str(knobs.HOST_PRE))
        self.post = FRNetPostprocessor(self.config)
        self.host_post = str(knobs.HOST_POST)
        self._host_outs = {}                     # HOST_POST=tiles: the persistent host tensor of each variant's output
        if self.host_post in ("tiles", "numba"):
            from .host import fast

            if not fast.warm_up():              # numba missing: plan.unsort + the numpy post-processing
                self.host_post = "numpy"

    def _warm_one(self, variant: Dict[str, Any]) -> None:
        """Warm up and capture the traces of one entry of ``default_warmup_variants()``. ``TraceRunner.capture``
        runs every pending variant eagerly (kernel JIT, program cache) before any capture, then captures with
        program-cache misses forbidden; it is idempotent."""
        self.runner.capture()

    def _prepare(self, points: Any = None, *, sweeps: Optional[Any] = None, images: Any = None,
                 calibration: Any = None, stream: Any = None, inputs: Any = None, **_: Any):
        """Host pre-processing (``host/frame.py``): decoding, the ego crop, float32 projection, fp16 interpolation,
        unique cells, the node's profile checks and the device plan. Client mistakes raise ``tio.InputError``
        (HTTP 400); an out-of-profile frame is ``prepared.skipped`` (no device run, empty output)."""
        if points is None:
            raise tio.InputError("this model needs 'points'")
        extra = [name for name, value in (("sweeps", sweeps), ("images", images), ("inputs", inputs),
                                          ("stream", stream)) if _given(value)]
        if extra:
            raise tio.InputError(f"FRNet takes one sensor-frame sweep ('points') and its 'calibration' only; got "
                                 f"{extra}")
        return self.pre(points, calibration)

    def _forward(self, prepared: Any) -> Any:
        """Upload the frame plan (features, CSR offsets, gather tables) into the persistent device inputs, replay
        the trace, read ``pred_probs`` back (one D2H) and return it as ``(N, 27)`` float32 in the network-input order
        (``plan.unsort``). A skipped frame runs nothing. Runs under the model lock."""
        if prepared.skipped is not None:
            return None
        mode = getattr(self, "host_post", "numpy")
        if mode == "tiles":                     # the raw TILE bytes of the output; un-permuted in _postprocess
            flat = self.device_tiles(prepared.plan)
            if flat is not None:
                return _DeviceRows(flat, prepared.plan, width=self._tile_width)
            mode = self.host_post = "numba"     # no raw host buffer access on this ttnn: the untilized rows
        if mode == "numba":                     # sorted device rows; un-permuted in _postprocess
            name, st = self.select(prepared.plan)
            raw = self.runner(name, inputs=st.write(prepared.plan))
            return _DeviceRows(raw, prepared.plan)
        return self.device_probs(prepared.plan)

    def device_tiles(self, plan: Any) -> Any:
        """Run one frame and return the output ``pred_probs`` ``[1, 1, N_cap, 32]`` as its raw TILE-layout float32
        elements (a view of a persistent host tensor, valid until the next call; None when this ttnn exposes no
        host buffer). Call it under the model lock."""
        import numpy as np
        import ttnn

        name, st = self.select(plan)
        self.runner.run(name, st.write(plan))
        dev_out = self.runner.outputs(name)
        if dev_out.layout != ttnn.TILE_LAYOUT or dev_out.dtype != ttnn.float32:
            return None
        host_out = self._host_outs.get(name)
        if host_out is None:
            host_out = self._host_outs[name] = ttnn.allocate_tensor_on_host(dev_out.spec, self.device)
            self._tile_width = int(dev_out.padded_shape[-1])
        ttnn.copy_device_to_host_tensor(dev_out, host_out, blocking=True)
        try:
            buf = host_out.host_buffer().get_shard(ttnn.MeshCoordinate(0, 0))
            return np.from_dlpack(buf).view(np.float32)
        except Exception:  # noqa: BLE001 - another ttnn build: fall back to the untilized rows
            return None

    def device_probs(self, plan: Any) -> Any:
        """``pred_probs`` ``(N, 27)`` float32 (network-input order) of one ``host.preprocess.FramePlan`` (the device
        half of ``_forward``; tests and tools feed stored network inputs through it). Call it under the model
        lock (``model(...)`` does)."""
        import numpy as np

        name, st = self.select(plan)
        raw = self.runner(name, inputs=st.write(plan))
        probs = np.asarray(raw, np.float32).reshape(-1, raw.shape[-1])
        return plan.unsort(probs[:, : self.config.num_classes])

    def _postprocess(self, raw: Any, prepared: Any, params: Dict[str, Any]) -> Output:
        """Host post-processing (``host/postprocess.py``, ``fill_cloud_kernel``): first-max argmax, probability,
        the class-probability filter -> ``FRNetSegmentation``; ``raw`` = ``pred_probs`` or None (skipped)."""
        meta = {"variant": self.variant, **prepared.stats()}
        self.post.filter_indices(params["filter_classes"])             # an unknown class name is a 400 either way
        if raw is None:
            return self.post.empty(model=self.MODEL_NAME, frame_id=prepared.frame_id, meta=meta)
        inp = prepared.inputs
        if isinstance(raw, _DeviceRows):
            return self.post.from_rows(raw.rows, raw.plan, points=inp.points, num_points_raw=inp.num_points_raw,
                                       input_index=prepared.input_index, params=params, model=self.MODEL_NAME,
                                       frame_id=prepared.frame_id, meta=meta, tile_width=raw.width)
        return self.post(raw, points=inp.points, num_points_raw=inp.num_points_raw,
                         input_index=prepared.input_index, params=params, model=self.MODEL_NAME,
                         frame_id=prepared.frame_id, meta=meta)

    def _release(self) -> None:
        """Release the traces, persistent tensors and device weights; also called when ``from_pretrained`` fails
        half-way (tolerates a partial build)."""
        runner = getattr(self, "runner", None)
        if runner is not None:
            runner.release()
            self.runner = None
        graph = getattr(self, "graph", None)
        if graph is not None:
            graph.release()
            self.graph = None

    def extra_info(self) -> Dict[str, Any]:
        """Additions to ``model.info`` and ``/info``: the node parameters and the runner's traces."""
        out: Dict[str, Any] = {}
        config = getattr(self, "config", None)
        if config is not None:
            out["frnet"] = config.describe()
        runner = getattr(self, "runner", None)
        if runner is not None:
            out["trace"] = runner.describe()
        graph = getattr(self, "graph", None)
        if graph is not None:
            out["graph"] = graph.describe()
        return out