# -*- coding: utf-8 -*-
import numpy
import numpy as np
import torch
import sys
from collections import OrderedDict
import scipy.stats
import pandas as pd
import os
import os.path
#
d1 = sys.argv[1] # bit
d2 = sys.argv[2] # mode
d3 = sys.argv[3] # n_exp
# d1=2
# d2 = sys.argv[2] # epoch
# d1=2
# d2=3


sum=0
flag=0


# CNN FLOPs = Cout * Hout * Wout * (2 * Cin * K * K )  是考虑bias  否则-1
# FCN FLOPs = Cout * Cin  是考虑bias 否则-1


# 把相关的relu，pool也考虑进去了
# MAdd
# weight0 =np.array( [ 705600.0+4704.0+ 3528.0 ,  480000.0+ 1600.0 + 1200.0 , 95880.0 + 120.0,
#            20076.0 + 84.0 , 1670.0 ])
# weight1=np.array([705,600.0 ,  480,000.0,+  95,880.0 ,
#            20,076.0  , 1,670.0 ])

# flops
weight_f0= np.array([357504+4704+4704, 241600+1600+1600,48000+120,10080+84,840])
weight_f1=np.array([357504, 241600,48000,10080,840])



summary_quan_dict=OrderedDict()
summary_base_dict=OrderedDict()
# 最外层：不同epoch的字典 内层：各个网络层的grads

flag = 0
dirpath = './project/p/qat_analysis_data/mode' + str(d2)

if not os.path.isdir(dirpath):
    os.makedirs(dirpath, mode=0o777)
    os.chmod(dirpath, mode=0o777)

if int(d2) == 1 or int(d2) == 2:
    csvpath = dirpath + '/scratch_loss.csv'
else:
    csvpath = dirpath + '/scratch_loss_' + str(d3) + '.csv'


if os.path.exists(csvpath):
    flag = 1

if flag == 0:  # 还没有csv
    df = pd.DataFrame()

else:  # 已有csv
    df = pd.read_csv(csvpath, index_col=0)
    df2 = pd.DataFrame()




for epoch in ([5, 10, 15, 20, 25, 30]):
    sums = []
    total_quan_list = []
    total_base_list = []
    for i in range(int(epoch)):
        if int(d2) == 1:
            total_quan_list.append(torch.load(
                './project/p/checkpoint/cifar-10_lenet_bn_quant/scratch/' + str(d1) + '/ckpt_cifar-10_lenet_bn_quant_' + str(
                    i + 1) + '.pth'))
        elif int(d2) == 2:
            total_quan_list.append(torch.load(
                './project/p/checkpoint/cifar-10_lenet_bn_quant/scratch/mode' + str(d2) + '/' + str(
                    d1) + '/ckpt_cifar-10_lenet_bn_quant_' + str(
                    i + 1) + '.pth'))
        else:
            total_quan_list.append(torch.load(
                './project/p/checkpoint/cifar-10_lenet_bn_quant/scratch/mode' + str(d2) + '_' + str(d3) + '/' + str(
                    d1) + '/ckpt_cifar-10_lenet_bn_quant_' + str(
                    i + 1) + '.pth'))
        sum_loss = 0
        loss = total_quan_list[i]['losses']
        # print(len(loss))
        # 每个epoch的不同batch的
        for j in range(len(loss)):
            sum_loss += loss[j].cpu()
        # print(sum_loss)
        sum_loss /= j
        sums.append(sum_loss)
        # print(sums)

    #print(sums[0] - sums[int(d1) - 1])

    if flag == 0:
        df[str(epoch)] = [(sums[0] - sums[int(epoch) - 1]).detach().numpy()]

    else:
        df2[str(epoch)] = [(sums[0] - sums[int(epoch) - 1]).detach().numpy()]



if flag == 0:
    # df.to_csv('./qat_analysis_data/scratch_loss.csv')
    df.to_csv(csvpath)
else:
    df = df.append(df2)
    # df.to_csv('./qat_analysis_data/scratch_loss.csv')
    df.to_csv(csvpath)