import functools
import time


def python_profile(func):
    """Decorator for profiling MXNet operation.
    Uses Python's time module to collect execution time information
    of the operation.

    Parameters
    ----------
    func:
        Operation to be executed and timed.

    Returns
    -------
    res, timing output. res being result returned after operator execution.
    profiler output is a dictionary with summary of operation execution.
    Example output : { "add": [{"avg_time_add": 0.4053089120425284,
                                "inputs": {
                                    "lhs": [1024, 1024],
                                    "rhs": [1024,1024]
                                }]
                     }
    """

    @functools.wraps(func)
    def python_profile_it(*args, **kwargs):
        start_time = time.perf_counter()    # 1
        res = func(*args, **kwargs)
        end_time = time.perf_counter()      # 2
        run_time = end_time - start_time    # 3

        # NOTE : same as cpp_profile_it
        if len(args) > 0:
            operator_name = args[0].__name__
        elif 'block' in kwargs:
            operator_name = kwargs['block']._op_name
        else:
            raise ValueError("Unable to identify operator name to extract profiler output!")

        profiler_output = {'avg_time_'+str(operator_name): run_time}
        return res, profiler_output
    return python_profile_it