# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["litert-torch", "onnx2torch", "onnxsim", "onnx", "onnxruntime", "numpy"]
# ///
"""Exports an ONNX model to LiteRT (.tflite, also run by anira's TensorFlow Lite engine) from the
PyTorch module sources.py gives, so that it computes what the ONNX model does, and checks it against
ONNX Runtime over consecutive calls with the state fed back.

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

import sys

import litert_torch
import numpy as np
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:]))
    converted = litert_torch.convert(module, tuple(torch.zeros(*shapes[n]) for n in input_names))
    converted.export(out_path)

    def run(values):
        out = converted(*[torch.from_numpy(v) for v in values])
        return [np.asarray(o) for o in (out if isinstance(out, (list, tuple)) else [out])]

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


if __name__ == "__main__":
    main()