# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["torch==2.12.0", "onnx2torch", "onnxsim", "onnx", "onnxruntime", "numpy"]
# ///
"""Exports an ONNX model to TorchScript for LibTorch (a trace of the PyTorch module sources.py gives,
with the ONNX model's state as tensors) and checks it against ONNX Runtime over consecutive calls with
the state fed back. torch 2.12.0 is the LibTorch anira's backends ship.

    uv run --no-project export_libtorch.py <model.onnx> <out.pt> [<description.model.json>]
                                           [--guitar-lstm <GuitarLSTM-dynamic.pt>]
"""

import sys

import torch

from check import compare
from sources import load, options


def main():
    onnx_path, out_path = sys.argv[1:3]
    module, input_names, output_names, shapes = load(onnx_path, *options(sys.argv[3:]))
    torch.jit.trace(module, tuple(torch.zeros(*shapes[n]) for n in input_names)).save(out_path)
    traced = torch.jit.load(out_path).eval()

    def run(values):
        with torch.no_grad():
            out = traced(*[torch.from_numpy(v) for v in values])
        return [o.numpy() for o in (out if isinstance(out, (list, tuple)) else [out])]

    compare("LibTorch", onnx_path, input_names, output_names, shapes, run)


if __name__ == "__main__":
    main()