import numpy as np
import onnx
import onnxsim
import onnx.reference
import cv2 as cv

shape_ini = (4, 5)
sim_net = True
N = 2
data = np.arange(0, 1 * N * 3 *shape_ini[0]* shape_ini[1]).reshape((1, N, 3, shape_ini[0], shape_ini[1])).astype(dtype=np.float32)    
shape_ini = data.shape 
select_axes = np.array([2], dtype=np.int64)
keepdims = 0
print("input DATA ")
print(data)
onnx_name = "reduceMean2inp"
out1 = onnx.helper.make_tensor_value_info('out1', onnx.TensorProto.FLOAT, [None, None, None, None])
inp0 = onnx.helper.make_tensor_value_info('inp0', onnx.TensorProto.FLOAT, [1, 2, 3 ,4 , 5])
axis = onnx.helper.make_tensor_value_info('axis', onnx.TensorProto.INT64, [1])
node1 = onnx.helper.make_node( "ReduceMean", inputs=["inp0", "axis"], outputs=["out1"], keepdims=keepdims )
graph = onnx.helper.make_graph([node1], onnx_name,  [inp0, axis], [out1])
onnx_model = onnx.helper.make_model(graph)

feeds = {'inp0': data, 'axis': select_axes}
sess = onnx.reference.ReferenceEvaluator(onnx_model)
res_onnx = sess.run(["out1"], feeds)
print("ONNX result")
print(res_onnx[0])
print("Writting model")
with open(onnx_name + ".onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())
if sim_net:
    print("Reading model")
    onnx_model = onnx.load(onnx_name+'.onnx')
    simplified_model, _ = onnxsim.simplify(onnx_model)
    onnx.save(simplified_model, onnx_name+'_sim.onnx')

net = cv.dnn.readNet(onnx_name+'.onnx')
net.enableWinograd(False)
print("Set opencv input DATA ")

net.setInput(data, "inp0")
net.setInput(data, "axis")
res_ocv = net.forward()
print("Opencv result. Shape = ", res_ocv.shape)
print(res_ocv)
print("Quadratic error for node  : ", np.mean((res_ocv -  res_onnx[0])**2))
print("Max error for  node : ", np.max((res_ocv -  res_onnx[0])**2))