# SPDX-License-Identifier: Apache-2.0
"""End-to-end fp32 CPU reference of METEOR's runtime on the released graph: the oracle of every TT gate.

Data flow per frame (SPEC sections 2-5):

    eight cameras + calibration + ego speed -> host.inputs / host.preprocess -> MeteorFrame (the ONNX feed)
      -> reference.model.MeteorNet (torch fp32, per-module taps) -> the 19 outputs
      -> host.result.build_output (3D / 2D decode, NMS, stationary, futures, unk2d, plan; per-stream seg fusion,
         yaw smoothing, mode hysteresis) -> host.outputs.MeteorOutput (the /predict JSON)

:meth:`MeteorReference.run_frame` is the network alone (a ``MeteorFrame`` -> numpy outputs, optional taps: what the
goldens hold). :meth:`MeteorReference.__call__` takes the same inputs as ``api.METEOR.__call__`` (``images``,
``calibration``, ``ego_speed``, ``stream``, runtime params) and keeps the host state per stream id, so the served
model and this reference can be compared request by request.

torch is imported when the reference is built. ``threads`` caps torch's intra-op threads (the workspace rule: <= 4).
"""
from __future__ import annotations

import time
from typing import Any, Dict, Mapping, Optional

import numpy as np

from ..host.inputs import prepare_request
from ..host.outputs import MeteorOutput
from ..host.postprocess import PostConfig
from ..host.preprocess import MeteorFrame
from ..host.result import build_output
from ..host.temporal import StreamState
from ..ttaw.golden import NULL_TAPS, TapRegistry
from .config import MeteorConfig, default_config
from .model import OUTPUT_NAMES, MeteorNet
from .weights import MeteorWeights

__all__ = ["MeteorReference", "to_numpy_outputs"]


def to_numpy_outputs(out: Mapping[str, Any]) -> Dict[str, np.ndarray]:
    return {k: (v.numpy() if hasattr(v, "numpy") else np.asarray(v)) for k, v in out.items()}


class MeteorReference:
    """``MeteorReference(weights_dir=None, cfg=None, threads=4)``: weights (``reference.weights``), the torch network
    and per-stream host state."""

    def __init__(self, weights_dir: Optional[Any] = None, cfg: Optional[MeteorConfig] = None, *,
                 threads: Optional[int] = 4, post: PostConfig = PostConfig(), weights: Optional[MeteorWeights] = None):
        import torch

        if threads:
            torch.set_num_threads(int(threads))
        self.cfg = (cfg or default_config()).check()
        self.weights = weights or MeteorWeights(weights_dir)
        self.net = MeteorNet(self.weights, self.cfg)
        self.post = post
        self.streams: Dict[str, StreamState] = {}

    def run_frame(self, frame: MeteorFrame, taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """The 19 outputs (numpy, ONNX names / dtypes / shapes) of one frame."""
        out = self.net(frame.imgs, frame.K, frame.T_cam_ego, frame.v0, present=frame.present, taps=taps)
        return to_numpy_outputs({k: out[k] for k in OUTPUT_NAMES})

    def stream(self, stream: Optional[Mapping[str, Any]]) -> StreamState:
        sid = str((stream or {}).get("id", "default"))
        st = self.streams.get(sid)
        if st is None:
            st = self.streams[sid] = StreamState()
        elif (stream or {}).get("reset"):
            st.reset()
        return st

    def __call__(self, images: Any = None, *, calibration: Optional[Mapping[str, Any]] = None, ego_speed: Any = None,
                 stream: Optional[Mapping[str, Any]] = None, output_heads: bool = False, **params: Any) -> MeteorOutput:
        t0 = time.perf_counter()
        frame = prepare_request(images, calibration, ego_speed, stream)
        t1 = time.perf_counter()
        outputs = self.run_frame(frame)
        t2 = time.perf_counter()
        cfg = self.post.with_params(**params)
        res = build_output(outputs, frame, cfg, self.stream(stream), heads=output_heads)
        t3 = time.perf_counter()
        res.timing_ms.update({"preprocess": 1e3 * (t1 - t0), "device": 1e3 * (t2 - t1), "postprocess": 1e3 * (t3 - t2),
                              "total": 1e3 * (t3 - t0)})
        res.meta["backend"] = "cpu-reference"
        return res