#!/usr/bin/python3
# Project: performance tools
# Author: syx10
# Time 2021/1/30:12:05
import math
import time
import os
import pickle


class TimerPerformance:
    """
    这个类用于统一记录函数运行时间， 可以自己定制不同的标签保存函数运行时间的数据，并输出详细信息。
    使用例子在最下方。
    """
    def __init__(self):
        super(TimerPerformance).__init__()
        self.timer_perf = {}
        self.save_dir = os.getcwd() + '/performance/'

    def reset(self):
        self.timer_perf.clear()

    def __set_tag(self, tag):
        if not self.timer_perf.get(tag):
            self.timer_perf[tag] = []

    def set_timer(self, tag, timer):
        if not self.timer_perf.get(tag):
            self.__set_tag(tag)
        self.timer_perf[tag].append(timer)

    def timer_wrapper(self, tag):
        def wrapper(func):
            def inner(*args, **kwargs):
                start = time.perf_counter()
                res = func(*args, **kwargs)
                end = round(time.perf_counter() - start, 6)
                self.set_timer(tag, end)
                return res
            inner()
            return inner
        return wrapper

    def show_perf(self):
        import numpy as np
        max_mean = {}
        the_max_max = 0, ''
        the_max_mean = 0, ''
        for ele in self.timer_perf.keys():
            max_mean[ele] = [max(self.timer_perf[ele]), np.mean(self.timer_perf[ele])]
            print("=========================="+str(ele)+"===============================")
            print("Length of " + ele + " is " + str(len(self.timer_perf[ele])))
            print("Max " + ele + " timer: " + str(max_mean[ele][0]))
            print("Mean " + ele + " timer: " + str(max_mean[ele][1]))
            print("All time consumption for " + str(ele) + " is :" + str(sum(self.timer_perf[ele])))
        for ele in max_mean:
            if max_mean[ele][0] > the_max_max[0]:
                the_max_max = max_mean[ele][0], ele

            if max_mean[ele][1] > the_max_mean[0]:
                the_max_mean = max_mean[ele][1], ele
        print("===========Synths============")
        print("The max time conso ====>" + str(the_max_max))
        print("The max mean time conso =====>" + str(the_max_mean))

    def set_save_dir(self, directory):
        if os.path.exists(directory):
            self.save_dir = directory
        else:
            return 'Directory: ' + str(directory) + ' not exist!'

    def save_perf(self):
        if not os.path.exists(self.save_dir):
            os.mkdir(self.save_dir)
        if len(self.timer_perf) > 0:
            try:
                filename = time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime()) + '_perf.pkl'
                perf_file = open(self.save_dir + filename, 'wb')
                pickle.dump(self.timer_perf, perf_file)
                return perf_file.name
            except Exception as e:
                raise e
        else:
            return 'Performance object has no data to save!'


def performance_viewer(perf_obj=None, filepath=''):
    if perf_obj is None and filepath == '':
        raise ValueError(' arguments exception.')
    if type(perf_obj) == TimerPerformance:
        if len(perf_obj.timer_perf) == 0:
            raise Exception('Performance object is empty')
        else:
            timer_perf = perf_obj.timer_perf
    elif os.path.exists(filepath):
        try:
            f = open(filepath, 'rb')
            timer_perf = pickle.load(f)
            return timer_perf
        except Exception as e:
            raise e
        finally:
            if f:
                f.close()
    else:
        raise ValueError(' arguments exception.')

    import matplotlib.pyplot as plt
    for ele in timer_perf:
        plt.figure()
        plt.plot(timer_perf[ele])
        plt.title(ele + " timer")
        plt.show()


def timer_wrapper():
    """
    此函数用于一次性获取函数运行时间，没有标签系统，直接在目标函数上加@timer_wrapper即可
    :return:
    """
    def wrapper(func):
        def deco(*args, **kwargs):
            start = time.perf_counter()
            res = func(*args, **kwargs)
            end = round(time.perf_counter() - start, 6)
            print("Time consumption for '", func.__name__, "' is: ", end)
            return res
        return deco
    return wrapper

# ============================Example=======================================
perf = TimerPerformance()

@perf.timer_wrapper('LOOP')
def loops():
    alist = [2]
    for i in range(1, 4000):
        for j in range(0, 200):
            alist.append(math.sqrt(i + j))

@perf.timer_wrapper('LOOP2')
def loops2():
    alist = [2]
    for i in range(1, 4000):
        for j in range(0, 200):
            alist.append(i/(j+1))


for i in range(0, 10):
    loops()
    loops2()

perf.show_perf()
# Answer
# ==========================LOOP===============================
# Length of LOOP is 11
# Max LOOP timer: 0.23034
# Mean LOOP timer: 0.22398136363636362
# All time consumption for LOOP is :2.4637949999999993
# ==========================LOOP2===============================
# Length of LOOP2 is 11
# Max LOOP2 timer: 0.149719
# Mean LOOP2 timer: 0.14588063636363635
# All time consumption for LOOP2 is :1.604687
# ===========Synths============
# The max time conso ====>(0.23034, 'LOOP')
# The max mean time conso =====>(0.22398136363636362, 'LOOP')

# ============================Example=======================================