#!/usr/bin/env python3

# used to generate model: onnx-tests/tests/div/div_shape_tensor.onnx

import numpy as np
import onnx
from onnx import helper, TensorProto
from onnx.reference import ReferenceEvaluator

# ONNX opset version to use for model generation
OPSET_VERSION = 16

def main():
    # Create a graph that tests Shape/Tensor and Tensor/Shape operations
    # Input tensors
    input_tensor = helper.make_tensor_value_info('input_tensor', TensorProto.FLOAT, [10, 20, 30])
    input_tensor_1d = helper.make_tensor_value_info('input_1d', TensorProto.INT64, [3])
    
    # Shape node - extract shape of input
    shape_node = helper.make_node(
        'Shape',
        inputs=['input_tensor'],
        outputs=['shape']
    )
    
    # Divide shape by tensor (Shape / Tensor)
    div_shape_tensor = helper.make_node(
        'Div',
        inputs=['shape', 'input_1d'],
        outputs=['shape_div_tensor']
    )
    
    # Divide tensor by shape (Tensor / Shape)
    div_tensor_shape = helper.make_node(
        'Div',
        inputs=['input_1d', 'shape'],
        outputs=['tensor_div_shape']
    )
    
    # Outputs
    output1 = helper.make_tensor_value_info('shape_div_tensor', TensorProto.INT64, [3])
    output2 = helper.make_tensor_value_info('tensor_div_shape', TensorProto.INT64, [3])
    
    # Create the graph
    graph_def = helper.make_graph(
        [shape_node, div_shape_tensor, div_tensor_shape],
        'shape_tensor_div_test',
        [input_tensor, input_tensor_1d],
        [output1, output2],
    )
    
    # Create the model
    model_def = helper.make_model(
        graph_def, 
        producer_name='onnx-tests',
        opset_imports=[helper.make_operatorsetid("", OPSET_VERSION)]
    )
    
    # Save the model
    onnx_name = "div_shape_tensor.onnx"
    onnx.save(model_def, onnx_name)
    print("Finished exporting model to {}".format(onnx_name))
    
    # Test the model with sample data
    test_input = np.random.randn(10, 20, 30).astype(np.float32)
    test_1d = np.array([2, 4, 5], dtype=np.int64)
    
    print(f"\nTest input shape: {test_input.shape}")
    print(f"Test 1d tensor: {test_1d}")
    
    # Run the model using ReferenceEvaluator
    session = ReferenceEvaluator(onnx_name, verbose=0)
    outputs = session.run(None, {"input_tensor": test_input, "input_1d": test_1d})
    
    shape_div_tensor, tensor_div_shape = outputs
    
    print(f"\nTest output shape_div_tensor: {repr(shape_div_tensor)}")
    print(f"Test output tensor_div_shape: {repr(tensor_div_shape)}")
    
    # Verify results - shape/tensor should be [5, 5, 6] and tensor/shape should be [0, 0, 0] (integer division)
    expected_shape_div = np.array([10, 20, 30]) // np.array([2, 4, 5])
    expected_tensor_div = np.array([2, 4, 5]) // np.array([10, 20, 30])
    
    assert np.array_equal(shape_div_tensor, expected_shape_div), f"Expected {expected_shape_div}, got {shape_div_tensor}"
    assert np.array_equal(tensor_div_shape, expected_tensor_div), f"Expected {expected_tensor_div}, got {tensor_div_shape}"

if __name__ == '__main__':
    main()