# SPDX-License-Identifier: Apache-2.0
"""Python API: METEOR (TIER IV, AutowareFoundation/meteor) on one Tenstorrent Blackhole p150.

    from tt_meteor import METEOR, load_sample

    with METEOR.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model(**load_sample("code/tt_meteor/samples/synthetic_8cam.json"))  # 8 cameras + calibration + speed
    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_meteor.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.

Inputs per call (``_prepare``; host code in ``tt_meteor.host``, METEOR's own runtime semantics: there is no Autoware
package):

- ``images``: the cameras ``CAM_FRONT_WIDE``, ``CAM_FRONT_LEFT``, ``CAM_FRONT_RIGHT``, ``CAM_BACK_WIDE``,
  ``CAM_BACK_LEFT``, ``CAM_BACK_RIGHT``, ``CAM_FRONT_NARROW``, ``CAM_BACK_NARROW`` (any order; raw, unrectified frames
  of at least 768x432, resized with OpenCV INTER_AREA semantics); the two narrow cameras may be absent (zero image +
  donor pose, as trained);
- ``calibration``: per camera ``intrinsics`` (of the image as sent) and ``T_ref_from_camera`` (camera optical frame ->
  ego / base_link), or ``{"preset": name}``;
- ``ego_speed``: m/s (the graph's ``v0``);
- ``stream``: ``{"id", "reset", "timestamp_s", "T_world_from_ego"}`` (or ``"pose": [x, y, yaw]``): the host temporal
  post-processing per stream (BEV seg fusion needs the pose; yaw smoothing; mode hysteresis).

Runtime params (``RUNTIME_PARAMS``): every knob of METEOR's C++ renderer (``host.postprocess.PostConfig``). Load-time
knobs: :data:`KNOBS` (``METEOR_*``): ``INPUT_NORM`` (``onnx`` = the released /255, D12; ``imagenet`` = the trained
normalisation), ``DEPTH_MEAN_BINS`` (``log`` as exported; ``linear``), ``MAX_STREAMS``.

The device graph is ``tt_meteor.tt.model.TtMETEOR``: the whole network in ONE trace (``frame``), replayed per call;
the host half (``_build_host``, ``_prepare``, ``_postprocess``) is what the host tests exercise. Importing this module
has no side effects (device code is imported in ``_build``).
"""
from __future__ import annotations

from collections import OrderedDict
from typing import Any, Dict, Mapping, Optional

from . import io as tio
from .device import DEVICE_DEFAULTS
from .host.outputs import LABELS as _LABELS
from .host.outputs import MeteorOutput
from .reference.config import CAMERAS
from .ttaw.api_base import ModelBase
from .ttaw.knobs import Knob, Knobs

__all__ = ["METEOR", "Output", "KNOBS", "load_sample"]

# The result class of this model (the multi-task MeteorOutput of host.outputs; BUNDLE_CONVENTIONS 7.4 "per head").
Output = MeteorOutput

# Load-time switches (read once at build; env METEOR_<NAME>). Each optimization adds its knob here (A/B switch).
KNOBS = Knobs("METEOR", [
    Knob("INPUT_NORM", "onnx", "input normalisation: the released graph's /255 (D12 default, ORT parity) or the "
                               "trained ImageNet mean/std (SPEC section 10 risk 1)", choices=("onnx", "imagenet")),
    Knob("DEPTH_MEAN_BINS", "log", "depth_mean bin centres: log-spaced as exported (parity) or linear 1 + 1.25 b "
                                   "(the trained bins, SPEC section 10 risk 2)", choices=("log", "linear")),
    Knob("MAX_STREAMS", 16, "host temporal states kept (seg fusion, yaw tracks, mode); the least recent is dropped"),
    Knob("IMAGE_PRECISION", "terms3", "image branch (ResNet-34, FPN, image heads): terms3 = fp32 activations, every "
                                      "conv as three bf16 terms (needed by the depth / seg2d argmax gates, PORT_LOG.md "
                                      "E20); bf16 = bf16 activations (A/B only: below the depth argmax gate, "
                                      "0.9797 < 0.99)", choices=("terms3", "bf16")),
])


def load_sample(path: Any, *, with_stream: bool = True) -> Dict[str, Any]:
    """``model(**load_sample("samples/<name>.json"))``: cameras + calibration preset + ego speed (+ stream)."""
    from .host.inputs import load_sample as _load

    return _load(path, with_stream=with_stream)


class METEOR(ModelBase):
    """METEOR (TIER IV, AutowareFoundation/meteor) on one Blackhole p150. Create it with :meth:`from_pretrained`."""

    MODEL_NAME = "meteor-p150"
    ENV_PREFIX = "METEOR"  # prefix of the environment knobs (SERVING.md section 3.4)
    DEFAULT_REPO = "AutowareFoundation/meteor"
    DEFAULT_TAG = "v1.0"  # the AutowareFoundation release tag (no Autoware ansible pin exists for METEOR)
    # the commit DEFAULT_TAG points to (pinned: tags can move)
    DEFAULT_REVISION = "01a5f6d71df5ecbbb5853ec600825481d57b9c6b"
    ALLOW_PATTERNS = ["meteor_v157c3Z.onnx", "meteor_v157.param.yaml", "LICENSE", "SHA256SUMS"]
    VARIANTS = ["default"]  # load-time: selects weights files and trace shapes
    DEFAULT_VARIANT = "default"
    INPUT_KIND = "multicam"  # lidar | camera | multicam | lidar+multicam | planner
    CAMERA_ORDER = CAMERAS  # the network's fixed camera order (meteor_v157.param.yaml:10)
    POINT_FIELDS = tio.DEFAULT_POINT_FIELDS  # unused (camera model); kept for the shared /info schema
    LABELS = list(_LABELS)  # 3D box classes (label_id 0 vehicle, 1 VRU)
    # Per-request knobs: name -> (type, min, max, default). Host-side post-processing only (SPEC section 9, RT-host;
    # render.cpp:345-376 defaults).
    RUNTIME_PARAMS = {
        "det3d_threshold": (float, 0.0, 1.0, 0.15),
        "det3d_topk": (int, 1, 1024, 64),
        "vehicle_threshold": (float, 0.0, 1.0, 0.35),
        "vru_threshold": (float, 0.0, 1.0, 0.15),
        "bev_nms_iou": (float, 0.0, 1.0, 0.3),
        "bev_nms_containment": (float, 0.0, 1.0, 0.6),
        "stationary_logit_threshold": (float, None, None, 0.0),
        "det2d_threshold": (float, 0.0, 1.0, 0.30),
        "det2d_topk": (int, 1, 1024, 48),
        "det2d_hide": (str, None, None, "7"),
        "unk2d": (bool, None, None, True),
        "unk2d_threshold": (float, 0.0, 1.0, None),
        "ground_z": (float, -5.0, 5.0, 0.0),
        "mode_hysteresis": (float, 0.0, 10.0, 0.35),
        "straight_margin": (float, 0.0, 10.0, 1.0),
        "seg_fuse": (bool, None, None, True),
        "thin_road_edge": (bool, None, None, True),
        "yaw_smoothing": (bool, None, None, True),
        "heads": (bool, None, None, False),
    }
    EXTRA_INPUTS = ("ego_speed",)
    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:
        """Weights (``reference.weights.MeteorWeights``: the ONNX parameters by consuming node) -> the device graph of
        ``tt_meteor.tt`` (``TtMETEOR``: the whole graph as ONE ``TraceRunner`` variant ``frame``, no capture yet)."""
        self._build_host()
        from .host.calib import CalibrationCache
        from .reference.weights import MeteorWeights
        from .tt.model import TtMETEOR
        from .tt.params import MeteorParams

        params = MeteorParams(MeteorWeights(self.weights_path), self.cfg)
        self.calib_cache = CalibrationCache(points=params.ground_points())
        self.tt = TtMETEOR(self.device, params, num_command_queues=self.device_info.get("num_command_queues"),
                           image_precision=str(self.knobs.IMAGE_PRECISION))
        self.runner = self.tt.runner

    def _warm_one(self, variant: Any) -> None:
        """Warm-up of the ``frame`` variant (compiles every program) and its capture."""
        self.runner.capture()

    def _build_host(self) -> None:
        """The host half of the model: load-time knobs and the per-stream host temporal states (no device, no
        weights; the host tests call it directly)."""
        from .host.postprocess import PostConfig
        from .reference.config import default_config

        params = {k.upper(): v for k, v in getattr(self, "compile_params", {}).items() if k.upper() in KNOBS.knobs}
        self.knobs = KNOBS.read(**params)
        self.cfg = default_config(input_norm=str(self.knobs.INPUT_NORM),
                                  depth_mean_bins=str(self.knobs.DEPTH_MEAN_BINS))
        self.post = PostConfig()
        self._streams: "OrderedDict[str, Any]" = OrderedDict()

    def _stream_state(self, stream: Optional[Mapping[str, Any]]):
        from .host.temporal import StreamState

        sid = str((stream or {}).get("id", "default"))
        st = self._streams.get(sid)
        if st is None:
            st = self._streams[sid] = StreamState()
            while len(self._streams) > int(self.knobs.MAX_STREAMS):
                self._streams.popitem(last=False)
        elif (stream or {}).get("reset"):
            st.reset()
        self._streams.move_to_end(sid)
        return sid, st

    def _prepare(self, points: Any = None, *, images: Any = None, calibration: Any = None, stream: Any = None,
                 ego_speed: Any = None, **_: Any) -> Dict[str, Any]:
        """Cameras + calibration + ego speed -> the graph feed (``host.inputs.prepare_request``)."""
        from .host.inputs import prepare_request

        if points is not None:
            raise tio.InputError("this model takes eight camera images (+ calibration, ego_speed), not a point cloud")
        frame = prepare_request(images, calibration, ego_speed, stream)
        sid, state = self._stream_state(stream)
        return {"frame": frame, "stream_id": sid, "state": state}

    def _forward(self, prepared: Dict[str, Any]) -> Any:
        """Upload (cameras, v0; the lift tables when the rig changed) + replay of the ``frame`` trace + one packed
        read -> the 19 outputs (``tt.unpack``: ONNX names / layouts / dtypes)."""
        frame = prepared["frame"]
        geom = self.calib_cache.get(frame.K, frame.T_cam_ego)
        return self.tt.run_frame(frame, geom)

    def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output:
        """METEOR's C++ host decode on the 19 outputs (``host.result.build_output``) with this stream's state."""
        from .host.result import build_output

        hide = tuple(int(v) for v in str(params.get("det2d_hide", "7")).replace(" ", "").split(",") if v != "")
        knobs = {k: v for k, v in params.items() if k not in ("det2d_hide", "heads")}
        cfg = self.post.with_params(det2d_hide=hide, **knobs)
        return build_output(raw, prepared["frame"], cfg, prepared["state"], heads=bool(params.get("heads")),
                            model=self.MODEL_NAME, meta={"stream_id": prepared["stream_id"]})

    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]:
        runner = getattr(self, "runner", None)
        info: Dict[str, Any] = {"knobs": self.knobs.as_dict()} if hasattr(self, "knobs") else {}
        if runner is not None:
            info["trace"] = runner.describe()
        return info