import hidet
import numpy as np
import onnx

from onnx import numpy_helper, TensorProto
from onnx.helper import (
    make_model, make_node, make_graph,
    make_tensor_value_info)

X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None, None])

v_min = numpy_helper.from_array(np.array(-1, dtype='float32'), name='v_min')
v_max = numpy_helper.from_array(np.array(1, dtype='float32'), name='v_max')

vmin = make_node('Constant', [], ['vmin'], value=v_min) 
vmax = make_node('Constant', [], ['vmax'], value=v_max) 

node = make_node('Clip',['X', 'vmin', 'vmax'], ['Y'])
graph = make_graph([vmin, vmax, node], 'foo', [X], [Y]) 
onnx_model = make_model(graph)

hidet.option.search_space(1)
hidet_onnx_module = hidet.graph.frontend.from_onnx(onnx_model)
hidet_onnx_module = hidet_onnx_module.to_cuda()

input_names = hidet_onnx_module.input_names

inputs = [hidet.asarray(np.array([6], dtype='float32'), device='cuda')]
symbol_data = [hidet.symbol_like(inputs[0])]
symbol_output = hidet_onnx_module(*symbol_data)
graph: hidet.FlowGraph = hidet.trace_from(symbol_output, inputs=symbol_data)
with hidet.graph.PassContext() as ctx:
    graph_opt: hidet.FlowGraph = hidet.graph.optimize(graph)
    print(graph_opt)
cuda_graph = graph_opt.cuda_graph()
outputs = cuda_graph.run(inputs)

print(outputs[0].detach().cpu().numpy())