import onnx
from onnx import helper, TensorProto, version_converter
import numpy as np
import onnxruntime as ort

def create_softmax_model_opset13():
    """
    Create a simple ONNX model with a Softmax node in opset 13.
    """
    # Create inputs and outputs
    input_tensor = helper.make_tensor_value_info('input', TensorProto.FLOAT, [3, 4])
    output_tensor = helper.make_tensor_value_info('output', TensorProto.FLOAT, [3, 4])

    # Create a Softmax node
    softmax_node = helper.make_node(
        'Softmax',
        inputs=['input'],
        outputs=['output'],
        axis=-1
    )

    # Create the graph
    graph = helper.make_graph(
        nodes=[softmax_node],
        name='SoftmaxGraph',
        inputs=[input_tensor],
        outputs=[output_tensor]
    )

    # Create the model (opset 13)
    model = helper.make_model(
        graph,
        opset_imports=[helper.make_operatorsetid("", 13)],
        producer_name='test_softmax_opset13'
    )

    # Explicitly set IR version to 10
    model.ir_version = 10

    return model

def run_model_with_ort(model_path, input_data):
    """
    Run the ONNX model with ONNX Runtime and return the outputs.
    """
    # Create an ONNX Runtime session
    session = ort.InferenceSession(model_path)

    # Run the session
    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: input_data})

    return outputs

def test_softmax_conversion_with_runtime():
    """
    Test the conversion of a Softmax model from opset 13 to opset 12
    and validate the results using ONNX Runtime.
    """
    # Step 1: Create the model
    model_opset13 = create_softmax_model_opset13()

    # Step 2: Save the model for inspection
    onnx.save(model_opset13, "softmax_opset13.onnx")

    # Step 3: Convert the model from opset 13 to opset 12
    converted_model = version_converter.convert_version(model_opset13, 12)

    # Step 4: Save the converted model for inspection
    onnx.save(converted_model, "softmax_opset12.onnx")

    # Step 5: Prepare dummy input
    input_data = np.random.rand(3, 4).astype(np.float32)

    # Step 6: Run the original model with ONNX Runtime
    original_outputs = run_model_with_ort("softmax_opset13.onnx", input_data)

    # Step 7: Run the converted model with ONNX Runtime
    converted_outputs = run_model_with_ort("softmax_opset12.onnx", input_data)

    # Step 8: Verify the outputs are identical
    np.testing.assert_allclose(
        original_outputs[0],
        converted_outputs[0],
        rtol=1e-5,
        atol=1e-6,
        err_msg="The outputs of the original and converted models do not match!"
    )

    # Step 9: Additional verification
    # Check that the converted model has the correct opset version
    assert converted_model.opset_import[0].version == 12, "Opset version is not 12!"

    print("Test passed: Softmax conversion from opset 13 to 12 is correct, and outputs match.")

if __name__ == "__main__":
    test_softmax_conversion_with_runtime()