# -*- coding: utf-8 -*-
"""
GOBO Baseline script

This script quantizes weight matrices in a numpy array format. 
How to use this code:
1. Extract the state dictionary of a model and store the weight matrices as a numpy array. 
2. Copy and run this script in the same directory as the weights. 
3. Replace the quantized weights in your model's state dictionary and run the model.

@author: Ali Hadi Zadeh
"""


bins = 8  # 2^bits for Quantized weights
outlierTresh = -4 # Outlier Threshold



import glob
import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture


###################### Outlier #######################
def Outlier_detection (w, outlierTresh):
    w_r=np.reshape(w,(-1,1))
    mean_w = w_r.mean()
    std_w = w_r.std()
    gm = GaussianMixture(n_components = 1, covariance_type = 'full', random_state=0)
    gm.fit(w_r)
    score=gm.score_samples(w_r)
    G_Matrix_index = np.nonzero(score>=outlierTresh)
    G_weights = w_r[G_Matrix_index]
    mean_g = G_weights.mean()
    std_g = G_weights.std()
    outlier_index = np.nonzero(score<outlierTresh)
    outlier_values = w_r[outlier_index]
    total_weights = w.size
    outlier_percent = outlier_values.size/w.size*100
    return outlier_index, outlier_values, G_weights, G_Matrix_index, total_weights, outlier_percent, mean_w, mean_g, std_w, std_g 



################## GOBO Original ######################
def GOBO (G_weights , mean_g , std_g, bins ):
    ARE = np.zeros(1000) # Absolute reconstruction Error per weight per iteration
    G_weights = G_weights.reshape(-1,1)
    
    if (G_weights.shape[0]==0):
        print("NO Gaussian weight!!")

    sorted_G = np.sort(G_weights , axis = 0)
    bin_interval = int(sorted_G.shape[0]/bins)+1
    init_index = np.arange(0, sorted_G.shape[0], bin_interval)
    G_edges = sorted_G[init_index].reshape(1,-1)
    G_index=np.sum(G_weights>=G_edges,axis=1)-1 
    w_dict=np.zeros(bins)
    for i in range (0,bins):
        index=G_index==i
        temp=G_weights[index]
        w_dict[i]=np.mean(temp)

    recon_Gw = w_dict[G_index].reshape (-1,1)
    ARE [0] = np.sum(abs (G_weights - recon_Gw) )/ G_weights.size


    ###### Start Iterations #####
    for j in range (1,1000):
      back_index=np.copy(G_index)
      G_index=np.argmin(abs(G_weights- w_dict),axis=1)
      back_w_dict=np.copy(w_dict)
      w_dict=np.zeros(bins)

      for i in range (0,bins):
          index=G_index==i
          temp=G_weights[index]
          w_dict[i]=np.mean(temp)
      
      recon_Gw=w_dict[G_index].reshape (-1,1)
      ARE [j] = np.sum(abs (G_weights - recon_Gw) )/ G_weights.size

      if ARE[j]>=ARE[j-1]:
          #RESTORE
          G_index=back_index
          w_dict=back_w_dict
          recon_Gw=w_dict[G_index]
          break
    return w_dict, G_index, recon_Gw 


################## Main ##############
name_of_trace = glob.glob ('*.npy')
name_of_trace.sort()


output_file_name='bins_'+str(bins)  +'_OT_'+str(outlierTresh)
os.mkdir(output_file_name)
file = open(output_file_name+'/'+output_file_name+'Result.txt','w') 
Lcount = 0

for Lname in name_of_trace:
    w_dict = 0
    G_index = 0
    recon_Gw = 0
    
    w=np.load(Lname)
    outlier_index , _ , G_weights , G_Matrix_index, total_weights , outlier_percent , _ , mean_g , _ , std_g  =  Outlier_detection (w,outlierTresh)
    w_dict , G_index , recon_Gw    = GOBO (G_weights  , mean_g  , std_g , bins )


    ##### Rebuild the matrix ######
    w_r=np.reshape(w,(-1,1))
    recon_w_r=np.copy(w_r)
    G_Matrix_index=np.array(G_Matrix_index).reshape(-1,1)
    recon_Gw=recon_Gw.reshape(-1,1)
    np.put(recon_w_r,G_Matrix_index,recon_Gw)
    error_w_r=np.sum(abs (recon_w_r - w_r) )/ w_r.size
    print ("Reconstruction Error: ","%0.4f" %error_w_r)
    w_recon_final=recon_w_r.reshape(w.shape)
    Lcount += 1

    np.save(output_file_name+'/'+Lname,w_recon_final )

    file.write('Reconstruction of Layer: %s \n'% Lname ) 
    file.write('Number of weights: %d \n' % (total_weights)) 
    file.write('Number of Outliers: %d \n' % (len(outlier_index[0]))) 
    file.write('Percentage of Outliers: %f \n' % (outlier_percent) ) 
    file.write('Total L1: %f \n' % (error_w_r))
    file.write('-------------------------------------- \n' )
file.close()