import onnx
from onnx.tools import update_model_dims

def update_inputs_outputs_dims(model, input_dims, output_dims):
    """
        This function updates the sizes of dimensions of the model's inputs and outputs to the values
        provided in input_dims and output_dims. if the dim value provided is negative, a unique dim_param
        will be set for that dimension.
    """
    def update_dim(tensor, dim, i, j, dim_param_prefix):
        dim_proto = tensor.type.tensor_type.shape.dim[j]
        if isinstance(dim, int):
            if dim >= 0:
                dim_proto.dim_value = dim
            else:
                dim_proto.dim_param = dim_param_prefix + str(i) + '_' + str(j)
        elif isinstance(dim, str):
            dim_proto.dim_param = dim
        else:
            raise ValueError('Only int or str is accepted as dimension value, incorrect type: {}'.format(type(dim)))

    for i, input_dim_arr in enumerate(input_dims):
        for j, dim in enumerate(input_dim_arr):
            update_dim(model.graph.input[i], dim, i, j, 'in_')

    for i, output_dim_arr in enumerate(output_dims):
        for j, dim in enumerate(output_dim_arr):
            update_dim(model.graph.output[i], dim, i, j, 'out_')

    onnx.checker.check_model(model)
    return model

def remove_unused_floor(model):
    nodes = model.graph.node

    for i, n in enumerate(nodes):
        n.name = str(i)

    floor_nodes = [node for node in nodes if node.op_type=='Floor']

    for f in floor_nodes:
        in_id = f.input[0]
        out_id = f.output[0]
        in_n = [node for node in nodes if node.output == [in_id]][0]
        if in_n.op_type == 'Mul':
            out_n = [node for node in nodes if node.input == [out_id]][0]
            out_n.input[0] = in_n.output[0]
            nodes.remove(f)

    return model


if __name__ == "__main__":  
    model_path = 'maskrcnn.onnx'
    save_path = 'maskrcnn_post.onnx'
    onnx_model = onnx.load(model_path)
    onnx_model = remove_unused_floor(onnx_model)
    #onnx_model = update_inputs_outputs_dims(onnx_model, [[1,9, 512, 512]], [[1,'nbox', 4], [1,'nbox'], [1,'nbox']])
    onnx_model = update_model_dims.update_inputs_outputs_dims(onnx_model, {'image':[9, 512, 512]},{'bboxes':['nbox', 4],'labels':['nbox'],'scores':['nbox']})
    onnx.save(onnx_model,save_path)
