import onnx
import onnx.helper
import numpy as np


# Make inputs and outputs
input_value = onnx.helper.make_tensor_value_info(
    'input',
    onnx.TensorProto.FLOAT,
    (1, 8, 32, 32)
)

output_value = onnx.helper.make_tensor_value_info(
    'output',
    onnx.TensorProto.FLOAT,
    (1, 8, 30, 30)
)

nodes = []

# Make a MaxPool node.
MP_node = onnx.helper.make_node(
    'MaxPool',
    ['input'],
    ['output'],
    name='MP0',
    kernel_shape = [3, 3]
)
nodes.append(MP_node)

# Make an ArgMax node.
ArgMax_node = onnx.helper.make_node(
    'ArgMax',
    ['input'],
    ['AM0'],
    name='argMax0'
)
nodes.append(ArgMax_node)

# Make a Dropout node.
Dropout_node = onnx.helper.make_node(
    'Dropout',
    ['input'],
    ['Dropout0'],
    name='Dropout0',
    ratio=0.1
)
nodes.append(Dropout_node)

# Make an GatherND node.
indices_node0 = onnx.helper.make_node(
    'Constant',
    [],
    ['indices0'],
    name='indices0',
    value=onnx.helper.make_tensor(
        name='indices0',
        data_type=onnx.TensorProto.INT64,
        dims=[4, 3],
        vals=np.array([[0, 1, 2],[0, 2, 3],[0, 1, 4],[0, 2, 5]])
    )
)
nodes.append(indices_node0)

GND_node = onnx.helper.make_node(
    'GatherND',
    ['Dropout0', 'indices0'],
    ['GND0'],
    name='GND0'
)
nodes.append(GND_node)

# Make a Cast node.
Cast_node0 = onnx.helper.make_node(
    'Cast',
    ['GND0'],
    ['cast0'],
    name='Cast0',
    to=onnx.TensorProto.INT32
)
nodes.append(Cast_node0)

# Make an DequantizeLinear node.
scale_node0 = onnx.helper.make_node(
    'Constant',
    [],
    ['scale0'],
    name='scale0',
    value=onnx.helper.make_tensor(
        name='scale0',
        data_type=onnx.TensorProto.FLOAT,
        dims=[1],
        vals=[1]
    )
)
nodes.append(scale_node0)

DL_node = onnx.helper.make_node(
    'DequantizeLinear',
    ['cast0', 'scale0'],
    ['DL0'],
    name='DL0'
)
nodes.append(DL_node)

# Make a LogSoftmax node.
LSM_node0 = onnx.helper.make_node(
    'LogSoftmax',
    ['Dropout0'],
    ['LSM0'],
    name='LSM0'
)
nodes.append(LSM_node0)

# Make a QuantizeLinear node.
scale_node1 = onnx.helper.make_node(
    'Constant',
    [],
    ['scale1'],
    name='scale1',
    value=onnx.helper.make_tensor(
        name='scale1',
        data_type=onnx.TensorProto.FLOAT,
        dims=[1],
        vals=[2]
    )
)
nodes.append(scale_node1)

QL_node0 = onnx.helper.make_node(
    'QuantizeLinear',
    ['LSM0', 'scale1'],
    ['QL0'],
    name='QL0'
)
nodes.append(QL_node0)

# Make a ReduceSum node.
RS_node0 = onnx.helper.make_node(
    'ReduceSum',
    ['LSM0'],
    ['RS0'],
    name='RS0',
    axes=[1]
)
nodes.append(RS_node0)

# Make a Softmax node.
SM_node0 = onnx.helper.make_node(
    'Softmax',
    ['RS0'],
    ['SM0'],
    name='SM0'
)
nodes.append(SM_node0)

# Make a Split node.
Split_node0 = onnx.helper.make_node(
    'Split',
    ['SM0'],
    ['Split00', 'Split01'],
    name='Split0',
    axis=2
)
nodes.append(Split_node0)

# Make a Squeeze node.
Squeeze_node0 = onnx.helper.make_node(
    'Squeeze',
    ['Split00'],
    ['Squeeze0'],
    name='Squeeze0',
    axes=[1]
)
nodes.append(Squeeze_node0)

# Make a Squeeze node.
Squeeze_node1 = onnx.helper.make_node(
    'Squeeze',
    ['Split01'],
    ['Squeeze1'],
    name='Squeeze1'
)
nodes.append(Squeeze_node1)

# Make a Squeeze node.
Unsqueeze_node1 = onnx.helper.make_node(
    'Unsqueeze',
    ['Squeeze0'],
    ['Unsqueeze1'],
    name='Unsqueeze1',
    axes=[1]
)
nodes.append(Unsqueeze_node1)

# Make a Resize node.
scale_node2 = onnx.helper.make_node(
    'Constant',
    [],
    ['scale2'],
    name='scale2',
    value=onnx.helper.make_tensor(
        name='scale2',
        data_type=onnx.TensorProto.FLOAT,
        dims=[4],
        vals=[2, 1, 1, 1]
    )
)
nodes.append(scale_node2)

roi_node0 = onnx.helper.make_node(
    'Constant',
    [],
    ['roi0'],
    name='roi0',
    value=onnx.helper.make_tensor(
        name='roi0',
        data_type=onnx.TensorProto.FLOAT,
        dims=[8],
        vals=[0, 0, 0, 0, 1, 1, 1, 1]
    )
)
nodes.append(roi_node0)

Resize_node0 = onnx.helper.make_node(
    'Resize',
    ['QL0', 'roi0', 'scale2'],
    ['Resize0'],
    name='Resize0',
    coordinate_transformation_mode="tf_half_pixel_for_nn"
)
nodes.append(Resize_node0)

# Make model.
graph_def = onnx.helper.make_graph(
    nodes,
    'test-model',
    [input_value],
    [output_value]
)
model_def = onnx.helper.make_model(graph_def, producer_name='onnx-example', opset_imports=[onnx.helper.make_opsetid("", 11)])
onnx.save(model_def, 'update_test.onnx')