from pathlib import Path
import numpy as np
import onnx
from onnx import TensorProto, helper, mapping, numpy_helper
import onnxruntime
import tvm
from tvm import relay

def get_input_data_shape_dict(graph_def, input_data):
    """Get input data shape"""
    if isinstance(input_data, list):
        input_names = {}
        shape_dict = {}
        for i, _ in enumerate(input_data):
            input_names[i] = graph_def.graph.input[i].name
            input_ = input_data[i]

            if input_ is None or not hasattr(input_, "shape") or input_.shape == ():
                # Skip adding input shape data when the input data is None;
                # This is to enable optional arguments for onnx operators.
                continue

            elif isinstance(input_, list):
                shape_dict[input_names[i]] = (len(input_),)

            else:
                shape_dict[input_names[i]] = input_.shape

    else:
        input_names = graph_def.graph.input[0].name
        shape_dict = {input_names: input_data.shape}

    return input_names, shape_dict

def get_onnxruntime_output(graph_def, inputs):
    """Generic function to generate onnxruntime output"""
    # rep = onnxruntime.backend.prepare(graph_def.SerializeToString(), 'CPU', providers=['CPUExecutionProvider'])
    sess = onnxruntime.InferenceSession(
        graph_def.SerializeToString(), providers=['CPUExecutionProvider']
    )
    for x, data in zip(sess.get_inputs(), inputs):
        input_names[x.name] = data
    output_names = [out.name for out in sess.get_outputs()]
    output = sess.run(output_names, input_names)
    return output

def get_tvm_output(
    graph_def,
    input_data,
    target,
    dev,
    opset=None,
    freeze_params=False,
    convert_config=None,):
    if not isinstance(input_data, list):
        input_data = [input_data]
    _, shape_dict = get_input_data_shape_dict(graph_def, input_data)
    mod, params = relay.frontend.from_onnx(
        graph_def,
        shape_dict,
        opset=opset,
        freeze_params=freeze_params,
        convert_config=convert_config,
    )
    result = relay.create_executor("vm", mod=mod, device=dev, target=target).evaluate()(
        *input_data, **params
    )
    if isinstance(result, tvm.runtime.NDArray):
        return [result.numpy()]
    return [r.numpy() for r in result]