# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import time

import onnx

from onnxscript import optimizer


def timeit(f, message):
    def timed(*args, **kw):
        ts = time.time()
        result = f(*args, **kw)
        te = time.time()
        print(f"{message} time: {te - ts}")
        return result

    return timed


load = timeit(onnx.load, "Load")

save = timeit(onnx.save, "Save")

infer = timeit(onnx.shape_inference.infer_shapes, "Infer")

fold_constants = timeit(optimizer.fold_constants, "Fold Constants")

remove_unused = timeit(optimizer.remove_unused_nodes, "Remove Unused")

optimize = timeit(optimizer.optimize, "Optimize")

# rewrite = timeit(all_rules.apply_to_model, "Rewrite")