import onnx
import numpy as np
import numpy.testing as npt
import onnxruntime as ort
from onnx import helper


def test_model(onnx_float_type, np_float_type):
    shape = (3,1)
    input_tensor = helper.make_tensor_value_info(name='input', elem_type=onnx_float_type, shape=shape)
    output_tensor = helper.make_tensor_value_info(name='output', elem_type=onnx_float_type, shape=shape)

    min_tensor = helper.make_tensor(
        name='min_tensor',
        data_type=onnx_float_type,
        dims=(1,),
        vals=[np_float_type(0.0)]
    )

    min_node = helper.make_node(
        'Constant',
        inputs=[],
        outputs=['min'],
        value=min_tensor
    )

    clip_node = helper.make_node(
        'Clip',
        inputs=['input', 'min'],  # no 'max' input
        outputs=['output'],
        name='clip'
    )

    graph = helper.make_graph(
        nodes=[min_node, clip_node],
        name='clip_model',
        inputs=[input_tensor],
        outputs=[output_tensor]
    )

    model = helper.make_model(graph)

    model_path = 'clip_model.onnx'
    onnx.save_model(model, model_path)
    session = ort.InferenceSession(model_path)

    test_input = np.array([[-1.0], [0.5], [1.5]], dtype=np_float_type)
    expected_output = np.array([[0.0], [0.5], [1.5]], dtype=np_float_type)

    outputs = session.run(None, {"input": test_input})

    print("Output:", outputs[0])
    npt.assert_array_equal(outputs[0], expected_output)


test_model(onnx.TensorProto.FLOAT, np.float32)
test_model(onnx.TensorProto.FLOAT16, np.float16)