# SPDX-License-Identifier: Apache-2.0
"""ONNX Runtime golden: the deployed ``frnet_<variant>.onnx`` of AutowareFoundation/lidar_frnet@v2.0 run as-is (CPU
execution provider; re-runs are bit-identical, S:frnet:332) on the four unpadded network inputs, exactly as
the node binds them (``points`` float32 (N, 4), ``coors`` int64 (N, 3), ``voxel_coors`` int64 (M, 3),
``inverse_map`` int64 (N,)).

With ``intermediates=True`` the in-memory graph also exposes the ONNX tensors that hold the same values as the torch
reference's taps (the mapping is ``FRNetWeights.onnx_taps``, recovered from the graph by ``reference/weights.py``),
so ``tests/test_reference_cpu.py`` compares module by module. Adding graph outputs changes nothing that is
computed. Choose which taps to expose with ``tap_patterns`` (fnmatch); OT128 module taps at 146k points are ~1 GB.

onnxruntime is a research / test dependency only (not in the image): it is imported lazily, and the tests skip
where it is absent.
"""
from __future__ import annotations

import fnmatch
from typing import Any, Dict, Optional, Sequence

import numpy as np

from ..ttaw.golden import NULL_TAPS, TapRegistry
from .config import variant_files
from .weights import MODULE_TAPS, FRNetWeights, load_weights

__all__ = ["OrtFRNet", "onnxruntime_available"]


def onnxruntime_available() -> bool:
    import importlib.util

    return importlib.util.find_spec("onnxruntime") is not None and importlib.util.find_spec("onnx") is not None


def _session(model_bytes: bytes, threads: int):
    import onnxruntime as ort

    so = ort.SessionOptions()
    so.intra_op_num_threads = int(threads)
    so.inter_op_num_threads = 1
    so.log_severity_level = 3
    return ort.InferenceSession(model_bytes, so, providers=["CPUExecutionProvider"])


def _with_outputs(path: Any, tensors: Sequence[str]) -> bytes:
    """The ONNX model at ``path`` with ``tensors`` added as graph outputs (in memory)."""
    import onnx
    from onnx import TensorProto, helper

    model = onnx.load(str(path))
    existing = {o.name for o in model.graph.output}
    for t in tensors:
        if t not in existing:
            model.graph.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None))
            existing.add(t)
    return model.SerializeToString()


class OrtFRNet:
    """The deployed graph of one variant under ONNX Runtime (module docstring)."""

    def __init__(self, weights_dir: Any, variant: str, *, threads: int = 2, intermediates: bool = False,
                 tap_patterns: Sequence[str] = MODULE_TAPS, weights: Optional[FRNetWeights] = None):
        self.variant = variant
        self.path = variant_files(weights_dir, variant)["onnx"]
        self.weights = weights or load_weights(weights_dir, variant)
        self.taps: Dict[str, str] = {}
        if intermediates:
            self.taps = {name: tensor for name, tensor in self.weights.onnx_taps.items()
                         if name != "pred_probs" and any(fnmatch.fnmatchcase(name, p) for p in tap_patterns)}
            model_bytes = _with_outputs(self.path, list(self.taps.values()))
        else:
            model_bytes = self.path.read_bytes()
        self.session = _session(model_bytes, threads)
        self.output_names = [o.name for o in self.session.get_outputs()]

    def __call__(self, *args: Any, **kwargs: Any) -> np.ndarray:
        return self.forward(*args, **kwargs)

    def forward(self, points: np.ndarray, coors: np.ndarray, voxel_coors: np.ndarray, inverse_map: np.ndarray,
                taps: TapRegistry = NULL_TAPS) -> np.ndarray:
        feeds = {"points": np.ascontiguousarray(points, np.float32), "coors": np.ascontiguousarray(coors, np.int64),
                 "voxel_coors": np.ascontiguousarray(voxel_coors, np.int64),
                 "inverse_map": np.ascontiguousarray(inverse_map, np.int64)}
        outs = dict(zip(self.output_names, self.session.run(None, feeds)))
        for name, tensor in self.taps.items():
            if taps.wants(name):
                taps.tap(name, outs[tensor])
        probs = outs["pred_probs"]
        taps.tap("pred_probs", probs)
        return probs