# BSD 3-Clause License
#
# Copyright (c) 2017 xxxx
# All rights reserved.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# ============================================================================
import time
from collections import defaultdict

_total_times = defaultdict(lambda:  0)
_start_times = defaultdict(lambda: -1)
_disabled_names = set()
_timer_stack = []
_running_timer = None
_disable_all = False

def disable_all():
	global _disable_all
	_disable_all = True

def enable_all():
	global _disable_all
	_disable_all = False

def disable(fn_name):
	""" Disables the given function name fom being considered for the average or outputted in print_stats. """
	_disabled_names.add(fn_name)

def enable(fn_name):
	""" Enables function names disabled by disable. """
	_disabled_names.remove(fn_name)

def reset():
	""" Resets the current timer. Call this at the start of an iteration. """
	global _running_timer
	_total_times.clear()
	_start_times.clear()
	_timer_stack.clear()
	_running_timer = None

def start(fn_name, use_stack=True):
	"""
	Start timing the specific function.
	Note: If use_stack is True, only one timer can be active at a time.
	      Once you stop this timer, the previous one will start again.
	"""
	global _running_timer, _disable_all
	
	if _disable_all:
		return

	if use_stack:
		if _running_timer is not None:
			stop(_running_timer, use_stack=False)
			_timer_stack.append(_running_timer)
		start(fn_name, use_stack=False)
		_running_timer = fn_name
	else:
		_start_times[fn_name] = time.perf_counter()

def stop(fn_name=None, use_stack=True):
	"""
	If use_stack is True, this will stop the currently running timer and restore
	the previous timer on the stack if that exists. Note if use_stack is True,
	fn_name will be ignored.

	If use_stack is False, this will just stop timing the timer fn_name.
	"""
	global _running_timer, _disable_all

	if _disable_all:
		return

	if use_stack:
		if _running_timer is not None:
			stop(_running_timer, use_stack=False)
			if len(_timer_stack) > 0:
				_running_timer = _timer_stack.pop()
				start(_running_timer, use_stack=False)
			else:
				_running_timer = None
		else:
			print('Warning: timer stopped with no timer running!')
	else:
		if _start_times[fn_name] > -1:
			_total_times[fn_name] += time.perf_counter() - _start_times[fn_name]
		else:
			print('Warning: timer for %s stopped before starting!' % fn_name)


def print_stats():
	""" Prints the current timing information into a table. """
	print()

	all_fn_names = [k for k in _total_times.keys() if k not in _disabled_names]

	max_name_width = max([len(k) for k in all_fn_names] + [4])
	if max_name_width % 2 == 1: max_name_width += 1
	format_str = ' {:>%d} | {:>10.4f} ' % max_name_width

	header = (' {:^%d} | {:^10} ' % max_name_width).format('Name', 'Time (ms)')
	print(header)

	sep_idx = header.find('|')
	sep_text = ('-' * sep_idx) + '+' + '-' * (len(header)-sep_idx-1)
	print(sep_text)

	for name in all_fn_names:
		print(format_str.format(name, _total_times[name]*1000))
	
	print(sep_text)
	print(format_str.format('Total', total_time()*1000))
	print()

def total_time():
	""" Returns the total amount accumulated across all functions in seconds. """ 
	return sum([elapsed_time for name, elapsed_time in _total_times.items() if name not in _disabled_names])


class env():
	"""
	A class that lets you go:
		with timer.env(fn_name):
			# (...)
	That automatically manages a timer start and stop for you.
	"""

	def __init__(self, fn_name, use_stack=True):
		self.fn_name = fn_name
		self.use_stack = use_stack

	def __enter__(self):
		start(self.fn_name, use_stack=self.use_stack)

	def __exit__(self, e, ev, t):
		stop(self.fn_name, use_stack=self.use_stack)