import argparse, os, glob, datetime, yaml, sys, gc
sys.path.append("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church")
sys.path.append("/home/liuxuewen/Dome/DeepCache/experiments")
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import tqdm
import torch 
import torch.nn as nn 
import torch.nn.functional as F
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns

def calculate_abCov(Y, X):
    Y_all = Y.permute(1, 0, 2, 3)
    X_all = X.permute(1, 0, 2, 3)
    a = torch.zeros(Y_all.size(0))
    b = torch.zeros(Y_all.size(0))
    for i in range(len(Y_all)):
        Y = Y_all[i]
        X = X_all[i]
        X_flat = X.contiguous().view(-1)
        Y_flat = Y.contiguous().view(-1)
        X_mean = torch.mean(X_flat)
        Y_mean = torch.mean(Y_flat)

        cov_XY = torch.mean((X_flat - X_mean) * (Y_flat - Y_mean))
        var_X = torch.var(X_flat)

        a[i] = cov_XY / var_X
        b[i] = Y_mean - a[i] * X_mean
    return a, b

def cache_abCov_statistic():
    from quant.quant_layer import UniformAffineQuantizer 
    import torch.nn.functional as F
    # act_delta, act_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/act_quantizer_params.pth", map_location=torch.device('cpu'))
    # weight_delta, weight_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_quantizer_params.pth", map_location=torch.device('cpu'))
    # weight, bias = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_params.pth", map_location=torch.device('cpu'))
    # all_cache = torch.load("/mnt/data/liuxuewen/Deepcache/feature_maps.pt")
    # fwd_kwargs = dict(stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1)
    act_delta, act_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/act_quantizer_norm_params.pth", map_location=torch.device('cpu'))
    weight_delta, weight_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_quantizer_norm_params.pth", map_location=torch.device('cpu'))
    weight, bias = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_norm_params.pth", map_location=torch.device('cpu'))
    all_cache = torch.load("/mnt/data/liuxuewen/Deepcache/feature_norm_maps.pt")
    fwd_kwargs = dict(stride=(1, 1), padding=(1, 1), dilation=(1, 1), groups=1)
    
    fwd_func = F.conv2d
    wq_params = {'n_bits': 8, 'symmetric': False, 'channel_wise': True, 'scale_method': 'mse'}
    aq_params = {'n_bits': 8, 'symmetric': False, 'channel_wise': False, 'scale_method': 'mse', 'leaf_param': True, "prob": 1.0, "num_timesteps": 250}
    act_quantizer = UniformAffineQuantizer(**aq_params)
    weight_quantizer = UniformAffineQuantizer(**wq_params)
    act_quantizer.inited = True
    weight_quantizer.inited = True
    act_quantizer.delta = act_delta
    act_quantizer.zero_point = act_zero_point
    weight_quantizer.delta = weight_delta
    weight_quantizer.zero_point = weight_zero_point
    weight = weight[:, :192, ...]
    q_weight = weight_quantizer(weight)
    '''
    with torch.no_grad():
        cache_c = all_cache[192]
        cache_r = all_cache[209]
        act_quantizer.time = 209
        q_cache_c = act_quantizer(cache_c)

        WrXr = fwd_func(cache_r, weight, bias, **fwd_kwargs)
        WrXc = fwd_func(cache_c, weight, bias, **fwd_kwargs)
        WqXcq = fwd_func(q_cache_c, q_weight, bias, **fwd_kwargs)

        a, b = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_abCov_list.pth")
        a = a[209].contiguous().view(192, 1, 1, 1).cpu()
        b = b[209].cpu()
        pre_WrXr = fwd_func(q_cache_c, a * q_weight, bias + b, **fwd_kwargs)
        newE = WrXr - pre_WrXr
        fig = plt.figure(figsize=(10, 5))
        data_for_boxplot = newE.permute(1, 0, 2, 3).reshape(newE.size(1), -1).detach()
        plt.boxplot(data_for_boxplot, showfliers=False)
        plt.xlabel('Channels (dim=1)')
        plt.ylabel('Values')
        plt.savefig("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/cache_plox_newECov_distribution.png")
    '''
    # '''
    interval_seq = range(0,100,5)
    interval_seq = [0, 10, 20, 27, 33, 38, 42, 47, 52, 56, 60, 64, 69, 74, 79, 83, 88, 92, 96, 99]
    Cova_list = []
    Covb_list = []
    with torch.no_grad():
        for i in range(len(all_cache)):
            if i in interval_seq:
                cache_c = all_cache[i]
                a = torch.ones_like(bias)
                b = torch.zeros_like(bias)
                Cova_list.append(a)
                Covb_list.append(b)
            else:
                cache_r = all_cache[i]
                act_quantizer.time = i
                q_cache_c = act_quantizer(cache_c)

                WrXr = fwd_func(cache_r, weight, bias, **fwd_kwargs)
                WrXc = fwd_func(cache_c, weight, bias, **fwd_kwargs)
                WqXcq = fwd_func(q_cache_c, q_weight, bias, **fwd_kwargs)

                Y_all = WrXr.permute(1, 0, 2, 3)
                X_all = WqXcq.permute(1, 0, 2, 3)
                a = torch.zeros_like(bias)
                b = torch.zeros_like(bias)
                for i in range(len(Y_all)):
                    Y = Y_all[i]
                    X = X_all[i]
                    X_flat = X.contiguous().view(-1)
                    Y_flat = Y.contiguous().view(-1)
                    X_mean = torch.mean(X_flat)
                    Y_mean = torch.mean(Y_flat)

                    cov_XY = torch.mean((X_flat - X_mean) * (Y_flat - Y_mean))
                    var_X = torch.var(X_flat)

                    a[i] = cov_XY / var_X
                    b[i] = Y_mean - a[i] * X_mean
                Cova_list.append(a)
                Covb_list.append(b)

    Cova_list = [a.cuda() for a in Cova_list]
    Covb_list = [b.cuda() for b in Covb_list]
    torch.save((Cova_list, Covb_list), "/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_norm_abCov_my_list.pth") #_norm
    # '''


def cache_pred_cacheerr_abCov_statistic(interval):
    all_cache = torch.load("/mnt/data/liuxuewen/Deepcache/feature_maps.pt")
    # interval_seq, all_cali_data, all_t, all_cali_t, all_cache = \
    #         torch.load("/mnt/data/liuxuewen/Deepcache/church{}_cache{}_{}.pth".format(100, interval, "my_opt"))
    # del all_cali_data, all_t, all_cali_t, all_cache
    interval_seq = range(0,100,interval)
    print(list(interval_seq))
    Cova_list = []
    Covb_list = []
    with torch.no_grad():
        for i in range(len(all_cache)):
            if i in interval_seq:
                cache_c = all_cache[i]
                a = torch.ones(cache_c.size(1))
                b = torch.zeros(cache_c.size(1))
                Cova_list.append(a)
                Covb_list.append(b)
            else:
                cache_r = all_cache[i]
                a, b = calculate_abCov(cache_r, cache_c)
                Cova_list.append(a)
                Covb_list.append(b)

    Cova_list = [a.cuda() for a in Cova_list]
    Covb_list = [b.cuda() for b in Covb_list]
    torch.save((Cova_list, Covb_list), f"/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_cacheerr_abCov_interval{interval}_list.pth")
    # '''
class GroupNorm32(nn.GroupNorm):
    def forward(self, x):
        return super().forward(x.float()).type(x.dtype)
def normalization(channels):
    return GroupNorm32(16, channels)
def cache_pred_quanterr_abCov_statistic(weight_bit, norm, interval):
    from quant.quant_layer import UniformAffineQuantizer 
    import torch.nn.functional as F
    in_layers = nn.Sequential(
        normalization(192),
        nn.SiLU(),
    )
    all_cache = torch.load("/mnt/data/liuxuewen/Deepcache/feature_maps.pt")
    all_cache_after_cacheadd = []
    (a_list, b_list) = torch.load(f"/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_cacheerr_abCov_interval{interval}_list.pth")
    # interval_seq, all_cali_data, all_t, all_cali_t, all_cache = \
    #         torch.load("/mnt/data/liuxuewen/Deepcache/church{}_cache{}_{}.pth".format(100, interval, "my_opt"))
    # del all_cali_data, all_t, all_cali_t, all_cache
    interval_seq = range(0,100,interval)
    print(list(interval_seq))
    with torch.no_grad():
        for i in range(len(all_cache)):
            if i in interval_seq:
                cache_c = all_cache[i]
                all_cache_after_cacheadd.append(cache_c)
            else:
                a = a_list[i].contiguous().view(1, a_list[i].size(0), 1, 1).cpu()
                b = b_list[i].contiguous().view(1, b_list[i].size(0), 1, 1).cpu()
                cache_after_cacheadd = a * cache_c + b
                all_cache_after_cacheadd.append(cache_after_cacheadd)
    if norm:
        save_path = f"/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_norm_quanterr_abCov_weight{weight_bit}_interval{interval}_list.pth"
        # del all_cache
        # gc.collect()
        # all_cache = torch.load("/mnt/data/liuxuewen/Deepcache/feature_norm_maps.pt")
        group_weight, group_bias = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/groupnorm_norm_params.pth")
        in_layers[0].weight = torch.nn.Parameter(group_weight[:192].cpu())
        in_layers[0].bias = torch.nn.Parameter(group_bias[:192].cpu())
        with torch.no_grad():
            all_cache_after_cacheadd = [in_layers(cache) for cache in all_cache_after_cacheadd]
        act_delta, act_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/act_quantizer_norm_params_aftercacheadd.pth", map_location=torch.device('cpu'))
        weight_delta, weight_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_quantizer_norm_params_aftercacheadd.pth", map_location=torch.device('cpu'))
        weight, bias = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_norm_params.pth", map_location=torch.device('cpu'))
        fwd_kwargs = dict(stride=(1, 1), padding=(1, 1), dilation=(1, 1), groups=1)
    else:
        save_path = f"/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/pre_quanterr_abCov_weight{weight_bit}_interval{interval}_list.pth"
        act_delta, act_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/act_quantizer_params_aftercacheadd.pth", map_location=torch.device('cpu'))
        weight_delta, weight_zero_point = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_quantizer_params_aftercacheadd.pth", map_location=torch.device('cpu'))
        weight, bias = torch.load("/home/liuxuewen/Dome/DeepCache/experiments/err_add/church/weight_params.pth", map_location=torch.device('cpu'))
        fwd_kwargs = dict(stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1)
    
    fwd_func = F.conv2d
    wq_params = {'n_bits': weight_bit, 'symmetric': False, 'channel_wise': True, 'scale_method': 'mse'}
    aq_params = {'n_bits': 8, 'symmetric': False, 'channel_wise': False, 'scale_method': 'mse', 'leaf_param': True, "prob": 1.0, "num_timesteps": 250}
    act_quantizer = UniformAffineQuantizer(**aq_params)
    weight_quantizer = UniformAffineQuantizer(**wq_params)
    act_quantizer.inited = True
    weight_quantizer.inited = True
    act_quantizer.delta = act_delta
    act_quantizer.zero_point = act_zero_point
    weight_quantizer.delta = weight_delta
    weight_quantizer.zero_point = weight_zero_point
    weight = weight[:, :192, ...]
    q_weight = weight_quantizer(weight)

    Cova_list = []
    Covb_list = []
    with torch.no_grad():
        for i in range(len(all_cache)):
            cache_c = all_cache_after_cacheadd[i]
            cache_r = all_cache[i]
            act_quantizer.time = i
            q_cache_c = act_quantizer(cache_c)

            # WrXr = fwd_func(cache_r, weight, bias, **fwd_kwargs)
            WrXc = fwd_func(cache_c, weight, bias, **fwd_kwargs)
            WqXcq = fwd_func(q_cache_c, q_weight, bias, **fwd_kwargs)

            a, b = calculate_abCov(WrXc, WqXcq)

            Cova_list.append(a)
            Covb_list.append(b)

    Cova_list = [a.cuda() for a in Cova_list]
    Covb_list = [b.cuda() for b in Covb_list]
    torch.save((Cova_list, Covb_list), save_path)

if __name__ == "__main__":
    from quant.utils import seed_everything
    seed_everything(1234+9)


    # cache_abCov_statistic()
    # cache_pred_cacheerr_abCov_statistic(interval=2)
    cache_pred_quanterr_abCov_statistic(weight_bit=8, norm=False, interval=2)
    cache_pred_quanterr_abCov_statistic(weight_bit=8, norm=True, interval=2)
