import functions as aux
import statistics as st
# import pandas as pd

max_epoch = 30
max_size = 8

data_types = ['mni', 'cif']
model_types = ['fre', 'bay']

o_data_types = ['MNIST', 'CIFAR']
o_model_types = ['LeNet', 'BCNN']

gpu_ene_data = aux.load_pickle("energy_gpu.pkl")
cpu_ene_data = aux.load_pickle("energy_cpu.pkl")

gpu_exp_data = aux.load_pickle("result_gpu.pkl")

for data in data_types:
    for model in model_types:
        for size in range(1, max_size):
            gpu_ene_data[data][model][size] = \
                    list(
                        aux.split(gpu_ene_data[data][model][size], max_epoch)
                        )

for data in data_types:
    for model in model_types:
        for size in range(1, max_size):
            cpu_ene_data[data][model][size] = \
                    list(
                    aux.split(cpu_ene_data[data][model][size], max_epoch)
                    )

spl_ene_data = dict(gpu_ene_data)
for data in data_types:
    for model in model_types:
        for size in range(1, max_size):
            for i in range(0, max_epoch):
                spl_ene_data[data][model][size][i] = \
                    gpu_ene_data[data][model][size][i] +\
                    cpu_ene_data[data][model][size][i]

for data in data_types:
    for model in model_types:
        for size in range(1, max_size):
            for i in range(0, max_epoch):
                spl_ene_data[data][model][size][i] = \
                    sum(spl_ene_data[data][model][size][i])

for data in data_types:
    for model in model_types:
        for size in range(1, max_size):
            temp = []
            for i in range(0, max_epoch):
                temp.append(
                        # st.mean(spl_ene_data[data][model][size][0:i+1])
                        sum(spl_ene_data[data][model][size][0:i+1])
                        )
            spl_ene_data[data][model][size] = temp

eff_data = dict(gpu_ene_data)
for data in data_types:
    for model in model_types:
        if model == 'fre':
            alpha = 100
        elif model == 'bay':
            alpha = 1000
        for size in range(1, max_size):
            for i in range(0, max_epoch):
                eff_data[data][model][size][i] = \
                    (gpu_exp_data[data][model][size]['acc'][i] /
                     spl_ene_data[data][model][size][i]) * alpha


for data, o_data in zip(data_types, o_data_types):
    eff_data[o_data] = \
            eff_data.pop(data)


for o_data in o_data_types:
    for model, o_model in zip(model_types, o_model_types):
        eff_data[o_data][o_model] = \
                eff_data[o_data].pop(model)

# mul = pd.MultiIndex.from_product([[1, 2, 3, 4, 5, 6, 7],
#                                   ['bay', 'fre'], ['cif', 'mni']])

# eff_data = pd.DataFrame(eff_data)

aux.save_pickle("efficiency_data.pkl", eff_data)