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

    from tt_transfusion import TransFusion

    with TransFusion.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model("code/tt_transfusion/samples/test_pcd.npz")  # see PYTHON.md for every input form
    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_transfusion.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_transfusion.host`` (shared with the CPU reference ``tt_transfusion.reference``, so device and reference see
identical voxel and feature tensors), the ttnn graph in ``tt_transfusion.tt`` (one trace, run by a
``ttaw.trace.TraceRunner``). Importing it has no side effects: ttnn / torch are imported inside the hooks.

Inputs (``points``, plus optional ``sweeps`` or ``stream``; host/preprocess.py has the full contract):

- a cloud in ``base_link`` with fields x, y, z, intensity (0-255, Autoware's uint8 value; float values are clipped
  and truncated like a uint8 cast);
- a column named ``time_lag`` marks an already densified cloud (Autoware's 2-sweep input), used as is;
- ``sweeps=[{"points", "time_lag_s", "T_current_from_sweep"}]``: past sweeps densified on the host (stateless);
- ``stream={"id", "timestamp_s", "T_world_from_ego", "reset"}``: Autoware's stateful densification per stream id
  (``densification_num_past_frames: 1``), a missing pose refuses the frame without caching it (Autoware's TF failure);
- ``sampling="clean"`` (default: every point, first 20 per pillar in input order) or ``"autoware_compat"`` /
  ``{"mode": "autoware_compat", "offset": int}`` (the deployed node's capacity-wide shuffle emulated, and its skips of
  frames with < 5,000 or > 60,000 pillars); an input keyword (``EXTRA_INPUTS``), also a top-level field of
  ``/predict``.

Output: ``Detections3D`` with Autoware ``DetectedObjects`` semantics: ObjectClassification label ids (``LABELS``),
``existence_probability`` as the score, the network yaw (counter-clockwise from +x, z at the box centre),
(length, width, height); ``meta`` holds the pillar counts (overflow above 65,536 pillars, Autoware's skip verdict),
the sampling, the per-stage detection counts, the per-object orientation availability and proposal index.
"""
from __future__ import annotations

from typing import Any, Dict, Optional

from . import io as tio
from .device import DEVICE_DEFAULTS
from .host.postprocess import LABELS, RUNTIME_PARAMS, TransFusionPostprocessor
from .host.preprocess import Prepared, TransFusionPreprocessor
from .reference.config import DEFAULT_VARIANT, VARIANTS, TransFusionConfig
from .ttaw.api_base import ModelBase
from .ttaw.outputs import Detections3D

__all__ = ["TransFusion", "Output"]

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


class TransFusion(ModelBase):
    """TransFusion-L (Autoware lidar_transfusion) on one Blackhole p150. Create it with :meth:`from_pretrained`."""

    MODEL_NAME = "transfusion-p150"
    ENV_PREFIX = "TRANSFUSION"  # prefix of the environment knobs (SERVING.md section 3.4)
    DEFAULT_REPO = "AutowareFoundation/lidar_transfusion"
    DEFAULT_TAG = "v2.1"  # the Autoware ansible artifacts pin
    # the commit DEFAULT_TAG points to (pinned: tags can move)
    DEFAULT_REVISION = "8f07eb93038bfd62e35d1f78e56a6b3fb76b12aa"
    ALLOW_PATTERNS = ["transfusion.onnx", "transfusion_ml_package.param.yaml", "detection_class_remapper.param.yaml",
                      "deploy_metadata.yaml"]
    VARIANTS = list(VARIANTS)  # one network; sampling is a per-request host knob, not a variant
    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  # default column names of an unnamed (N, 4) array
    # output labels: autoware_perception_msgs ObjectClassification (label_id = the message value); the network's five
    # classes (CAR, TRUCK, BUS, BICYCLE, PEDESTRIAN) map onto them, the remapper adds TRAILER
    LABELS = list(LABELS)
    # Per-request knobs (host post-processing only; defaults are Autoware's): host/postprocess.py
    RUNTIME_PARAMS = RUNTIME_PARAMS
    # Per-request pre-processing keyword (handed to _prepare): host/preprocess.py parse_sampling
    EXTRA_INPUTS = ("sampling",)
    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, pre- / post-processors), then the ttnn graph of ``tt_transfusion.tt``
        (``TtTransFusion``: PFN, gather-form scatter, SECOND, SECONDFPN, dense head, local max, query head) as the
        ``"default"`` variant of a ``ttaw.trace.TraceRunner`` with two persistent inputs (no capture here;
        ``_warm_one`` captures), plus the persistent host buffers of those inputs (``ttaw.models.pillars.FrameStaging``)."""
        from .reference.weights import load_weights
        from .tt.model import FrameStaging, TtTransFusion, build_runner, stage_inputs

        self._build_host()
        # precision policy and knobs (TRANSFUSION_PRECISION, TRANSFUSION_BACKBONE_WEIGHT_TERMS) read once here
        self.graph = TtTransFusion(self.device, load_weights(self.weights_path, self.config), self.config)
        self.runner = build_runner(self.graph, variants=("default",), name=f"{self.MODEL_NAME}/{self.variant}")
        self.staging = self.graph.make_staging()   # C24 FrameStaging (slot) or tt/pfn.PairFrameStaging (pair)
        self._stage = stage_inputs

    def _build_host(self, *, max_streams: int = 16) -> None:
        """The host half of ``_build`` (no device): usable alone by host tests and tools."""
        self.config = TransFusionConfig.from_weights(self.weights_path, self.variant)
        from .tt.model import KNOBS

        k = KNOBS.read()
        self.pre = TransFusionPreprocessor(self.config, max_streams=max_streams, voxelizer=str(k.HOST_VOXELIZER),
                                           densify=str(k.HOST_DENSIFY))
        if "numba" in (k.HOST_VOXELIZER, k.HOST_DECORATE, k.HOST_DENSIFY, k.HOST_INDEX, k.HOST_NMS):
            from .host import fast

            fast.DECORATE_THREADS = max(1, int(k.HOST_THREADS))   # H1b: parallel decoration (bit-identical)
            fast.warm_up()                                       # compile the host twins at load, not on frame 1
        self.post = TransFusionPostprocessor(self.config, nms=str(k.HOST_NMS))

    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, stream: Optional[Any] = None,
                 images: Any = None, calibration: Any = None, inputs: Any = None, sampling: Any = None,
                 **_: Any) -> Prepared:
        """Host pre-processing (``host/preprocess.py``): densification, sampling, deterministic voxelization, the
        float32 decoration and the canvas gather index. Client mistakes raise ``tio.InputError`` (HTTP 400)."""
        if points is None:
            raise tio.InputError("this model needs 'points'")
        if images or inputs is not None:
            raise tio.InputError("TransFusion takes a point cloud only (no 'images' / 'inputs')")
        return self.pre(points, sweeps=sweeps, stream=stream, sampling=sampling)

    def _forward(self, prepared: Prepared) -> Any:
        """Upload the frame's decorated features and canvas index into the persistent device inputs, replay the
        trace, read the packed outputs back (one D2H) and unpack them into ``{cls_score0, bbox_pred0,
        dir_cls_pred0}``. A frame Autoware skips (``autoware_compat``) runs nothing. Runs under the model lock."""
        if prepared.skipped is not None:
            return None
        inputs = self._stage(self.staging, prepared)
        variant = getattr(self.staging, "variant", "default")    # PFN_BUCKETS: the frame's pillar capacity
        return self.graph.assemble(self.runner(variant, inputs=inputs))

    def _postprocess(self, raw: Any, prepared: Prepared, params: Dict[str, Any]) -> Output:
        """Host post-processing (``host/postprocess.py``): decode, score threshold, circle NMS, ROS conversion,
        IoU-BEV NMS, class remapper -> ``Detections3D``; ``raw`` = ``{cls_score0, bbox_pred0, dir_cls_pred0}`` or None
        for a skipped frame (no objects, as the node publishes nothing)."""
        meta = {"variant": self.variant, **prepared.stats()}
        if raw is None:
            return self.post.empty(model=self.MODEL_NAME, frame_id=prepared.frame_id, meta={**meta, "counts": {}})
        return self.post(raw, 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 config, the host state and the runner."""
        info: Dict[str, Any] = {}
        if getattr(self, "config", None) is not None:
            info["config"] = self.config.describe()
        if getattr(self, "pre", None) is not None:
            info["preprocess"] = self.pre.describe()
        runner = getattr(self, "runner", None)
        if runner is not None:
            info["trace"] = runner.describe()
        graph = getattr(self, "graph", None)
        if graph is not None:
            info["graph"] = graph.describe()
        return info