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

    from tt_sceneseg import SceneSeg

    with SceneSeg.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model("code/tt_sceneseg/samples/highway_normal_1.png")  # 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_sceneseg.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.

Input: ONE camera image per call, as VisionPilot's ``scene_seg_model`` node takes one image topic: the first
positional argument (path, bytes, PIL image, RGB uint8 HxWx3 array or tensor) or ``images=[...]`` like the server
(exactly one entry; its camera name is free); no calibration. ``io.single_camera`` decodes it for the API and the
server alike; ``tt_sceneseg.host`` reproduces the node's pre-processing bit for bit (BGR8, OpenCV INTER_LINEAR squash
to 640x320) and post-processing (argmax with lowest-index ties, class 1 -> 255, INTER_NEAREST back to the input size).

Variants (load-time, PLAN.md 2.2 / D10): the input channel order, folded into the stem conv on the device. ``bgr``
(default) is the deployed VisionPilot C++ order; ``rgb`` is the training order, better on two-wheelers (model card).
Importing this module has no side effects: ttnn / torch are imported inside the hooks, never at module level.
"""
from __future__ import annotations

from typing import Any, Dict, Mapping, Optional

from . import io as tio
from .device import DEVICE_DEFAULTS
from .host.output import SceneSegMask
from .reference import config
from .ttaw.api_base import ModelBase

__all__ = ["SceneSeg", "Output"]

# The result class of this model: a ttaw.outputs.Mask2D (the published mono8 mask) with the 3-class map.
Output = SceneSegMask


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

    MODEL_NAME = "sceneseg-p150"
    ENV_PREFIX = "SCENESEG"  # prefix of the environment knobs (SERVING.md section 3.4)
    # The weights are re-hosted in this bundle's own repo under weights/ (PLAN.md D10: the upstream
    # AutowareFoundation/SceneSeg card holds no weights; the files come from VisionPilot's Google Drive links).
    DEFAULT_REPO = config.WEIGHTS_REPO
    DEFAULT_TAG = None  # the weights are a commit of this repo, not a tag
    # the commit of the weights upload (deferred until it exists: PLAN.md 4.1 gates the publish, not the port)
    DEFAULT_REVISION = "a6829b246a05dba55c356323bbf75f7b1d148499"
    ALLOW_PATTERNS = ["weights/SceneSeg_FP32.onnx", "weights/provenance.json", "weights/LICENSE", "weights/NOTICE"]
    WEIGHTS_LICENSE = config.WEIGHTS_LICENSE
    VARIANTS = ["bgr", "rgb"]  # load-time: the input channel order (folded into the stem weights on the device)
    DEFAULT_VARIANT = "bgr"
    INPUT_KIND = "camera"  # lidar | camera | multicam | lidar+multicam | planner
    CAMERA_ORDER = ()  # one image per request (one node instance per camera)
    REQUIRE_CALIBRATION = False  # the node uses no CameraInfo / TF
    POINT_FIELDS = tio.DEFAULT_POINT_FIELDS  # unused (camera model); kept for the shared /info schema
    # class names in network output order (load_data_scene_seg.py:136-141; class 1 is the published mask)
    LABELS = list(config.CLASS_NAMES)
    # Per-request knobs: name -> (type, min, max, default). The node has none (pure argmax, no threshold); this port
    # adds one output toggle. Anything that changes a device shape or the weights is a load-time option
    # (BUNDLE_CONVENTIONS.md section 9).
    RUNTIME_PARAMS = {
        "class_map": (bool, None, None, True),  # also return the 3-class map (a by-product the node does not publish)
    }
    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:
        """Load the weights (``reference.weights.load_weights``: the deployed ONNX's initializers by consuming node,
        checked against the 108-layer architecture and the pinned sha256) and build the ttnn graph of
        ``tt_sceneseg.tt`` for the selected channel order with its ``TraceRunner`` (persistent uint8 input, uint32
        class-map readback; build-time knobs ``SCENESEG_*`` of ``tt/knobs.py``). No capture here (``_warm_one``)."""
        from .reference import assets
        from .reference.weights import load_weights
        from .tt.knobs import KNOBS
        from .tt.runner import SceneSegRunner

        onnx = assets.find_in(self.weights_path, config.ONNX_FILE)
        if onnx is None:
            raise FileNotFoundError(f"{config.ONNX_FILE} not found under {self.weights_path}")
        self.knobs = KNOBS.read()
        if "numba" in (self.knobs.HOST_SQUASH, self.knobs.HOST_POST):
            from .host import fast

            fast.warm_up()                                    # compile the host twins before serving
        self.runner = SceneSegRunner(self.device, load_weights(onnx), self.variant,
                                     weight_terms=self.knobs.WEIGHT_TERMS, act_terms=self.knobs.ACT_TERMS,
                                     name=self.MODEL_NAME, input_layout=self.knobs.INPUT_LAYOUT)

    def _warm_one(self, variant: Dict[str, Any]) -> None:
        """Warm up and capture the traces (``TraceRunner.capture``: every pending variant eagerly first, then the
        captures with program-cache misses forbidden; idempotent)."""
        self.runner.capture()

    def _prepare(self, points: Any = None, *, images: Any = None, calibration: Optional[Mapping[str, Any]] = None,
                 sweeps: Any = None, stream: Any = None, inputs: Any = None, **_: Any) -> Dict[str, Any]:
        """The node's pre-processing (``tt_sceneseg.host.preprocess``): the one image (``io.single_camera``, RGB) ->
        BGR8 -> the exact OpenCV INTER_LINEAR squash to 640x320 uint8 (the device input). ``tio.InputError`` for
        client mistakes (no image, two images, an undecodable image)."""
        from .host.preprocess import prepare_rgb

        other = sorted(k for k, v in {"sweeps": sweeps, "stream": stream, "inputs": inputs}.items() if v is not None)
        if other:
            raise tio.InputError(f"this model takes one camera image, not {other}")
        camera = tio.single_camera(points, images, calibration=calibration,
                                   require_calibration=self.requires_calibration())
        return {"prepared": prepare_rgb(camera.image, channel_order=self.variant), "camera": camera.name}

    def _forward(self, prepared: Dict[str, Any]) -> Any:
        """Upload ``prepared["prepared"].device_input()`` (the uint8 squash) into the persistent device input, replay
        the trace and read the class decision back: ``{"class_map": uint8 [320, 640]}``. Runs under the model
        lock."""
        return self.runner(prepared["prepared"].device_input())

    def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output:
        """The node's post-processing (``tt_sceneseg.host.postprocess``) of the device's class map (or of logits:
        CPU reference / debug readback): class 1 -> 255, INTER_NEAREST back to the source size, + the class map."""
        from .host.output import to_output
        from .host.postprocess import postprocess

        prep = prepared["prepared"]
        res = postprocess(prep.src_hw, logits=raw.get("logits"), class_map=raw.get("class_map"))
        return to_output(res, model=self.MODEL_NAME, channel_order=self.variant, include_class_map=params["class_map"])

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

    def extra_info(self) -> Dict[str, Any]:
        """Additions to ``model.info`` and ``/info``: the channel order, the network size and the classes, plus the
        trace variants, CQs and persistent tensors of the runner."""
        runner = getattr(self, "runner", None)
        info: Dict[str, Any] = {"channel_order": getattr(self, "variant", None), "input_hw": list(config.INPUT_HW),
                                "classes": list(config.CLASS_NAMES), "mask_class": config.FOREGROUND_CLASS}
        if runner is not None:
            d = runner.describe()
            info.update({"graph": d["model"], "trace": d["trace"], "knobs": self.knobs.as_dict()})
        return info