import numpy as np
import matplotlib.pyplot as plt
import pickle
import sys

# CTensor[alpha][alpha'][i][j]: correlation between position alpha of channel i and position alpha' of channel j
# CTensor generated from Big correlation 2D matrix where channel and input size are merged  
# layer index starts from 0 which correponds to the input layer
"""
# Net work feedward steps:
#   1. Initialize C tensor and mean activations by constructor
#   2. Extract the dimensionality list of current layer including block matrix and Big Matrix
#   3. Calculate Delta tensor
#   4. Iterate one layer ahead
#   5. repeat above steps until reaching the last layer
"""
class DMFNet:

    ############# PARAM #############
    g_w = 0.8
    rho = 0.1
    sigma_b = 0.1
    ############# PARAM #############
    
    def __init__(self, struct,kernel_size,phi = np.tanh,input_size = 30, weight = None, bias = None, input_CTensor = None, input_h_mean = None, detailed_info = False, eigen_check = False):
        self.struct = struct
        self.kernel_size = kernel_size
        self.layer_len = len(struct)
        self.currentLayerIdx = 0
        self.phi = phi
        self.detailed_info = detailed_info
        self.eigen_check = eigen_check
        # General weight formulation: weight[layer idx][Cout][Cin] is a list containing the kernel size matrix (has only one element)
        # layer idx >= 1
        if weight is None:
            self.weight = []
            self.weight.append('0')
            for n in range(self.layer_len - 1):
                C_in = struct[n]
                C_out = struct[n+1]
                weightPerLayer = [[[] for j in range(C_in)] for i in range(C_out)]
                self.weight.append(weightPerLayer) 
            self.InitialWeight()
        else:
            self.weight = weight
        # Inner formulation: for fast computation in the iteration step
        # weight_inner[layer idx][betai][betaj]:  betai/j indicates the 2D location of kernel || weight[layer idx][beta] is a matrix of size [Cout,Cin]  
        # self.weight_inner[layer idx][betai][betaj][0] is the transform matrix of size [C_out, C_in]
        self.weight_inner = ['0']
        for n in range(1,self.layer_len):
            weight_inner_layer = [[[] for betaj in range(self.kernel_size)] for betai in range(self.kernel_size)]
            for betai in range(self.kernel_size):
                for betaj in range(self.kernel_size):
                    C_in = struct[n-1]
                    C_out = struct[n]
                    wtmp = np.zeros([C_out,C_in])
                    for i in range(C_out):
                        for j in range(C_in):
                            wtmp[i][j] = self.weight[n][i][j][0][betai][betaj]
                    weight_inner_layer[betai][betaj].append(wtmp)
            self.weight_inner.append(weight_inner_layer)
        # Input corvariance tensor
        if input_CTensor is None:
            self.InitialInputCorvarianceTensor(DMFNet.rho,input_size)
        else:
            self.CTensor = input_CTensor
        # initial delta tensor
        self.UpdateDeltaTensor()
        if input_h_mean is None:
            self.InitialMeanActivation(input_size)
        else:
            self.h_mean = input_h_mean
        # initial bias, self.bias is a list of length [self.layer_len - 1] and each element is a column vector | self.bias[layer idx>=1] is a column vector 
        if bias is None:
            self.bias = ['0']
            self.InitialBias()
        else:
            self.bias = bias
        self.dimList = []   # store the dimensionality per feature map
        self.dimBigMatrix = []  # store the dimensionality of each layer's big matrix
        self.ratio = []
    
    def PrintCorvarianceHeatMap(self, savePath = './RandomNet/HeatMap/'):
        # save the heat map of big matrix   ||  called before updating the dimensionality
        path = savePath + 'Layer' + str(self.currentLayerIdx) + '.jpg'
        plt.figure()
        plt.imshow(self.CBigMat)  
        plt.colorbar()
        plt.savefig(path) 

    def InitialWeight(self):
        # lidx: layer index     cidx:   channel idx
        for n in range(self.layer_len - 1):
            C_in = self.struct[n]
            C_out = self.struct[n+1]
            for i in range(C_out):
                for j in range(C_in):
                    # weight[layer idx][Cout][Cin]
                    self.weight[n+1][i][j].append(np.random.randn(self.kernel_size,self.kernel_size)*(np.sqrt(DMFNet.g_w/(C_in*self.kernel_size*self.kernel_size))))

    def InitialBias(self):
        for n in range(1,self.layer_len):
            self.bias.append(np.sqrt(DMFNet.sigma_b)*np.random.randn(self.struct[n],1))
        if self.detailed_info:
            print('Bias initialization complete')
    
    def InitialMeanActivation(self,input_size = 30):
        # a list of length input_size*input_size and each element is a column vector of size [self.struct[0],1]
        self.h_mean = []
        h_mean_alpha = np.zeros([self.struct[0],1])
        for i in range(input_size*input_size):
            self.h_mean.append(h_mean_alpha)
        if self.detailed_info:
            print('Mean activation complete')

    def InitialInputCorvarianceTensor(self,rho,input_size = 30):
        # Big matrix 
        input_channel_size = self.struct[0] 
        alphaN = input_size*input_size
        edge_size = alphaN*input_channel_size
        C = np.zeros([edge_size,edge_size])
        for i in range(edge_size):
            for j in range(i,edge_size):
                if i == j:
                    C[i][j] = 1
                else:
                    C[i][j] = (-rho + 2*rho*np.random.rand())/np.sqrt(input_channel_size*self.kernel_size*self.kernel_size)
        # keep it symmetric 
        for i in range(edge_size):
            for j in range(i):
                C[i][j] = C[j][i]
        # transfer it to 4D tensor format
        self.CTensor = np.zeros([alphaN,alphaN,input_channel_size,input_channel_size])
        for a in range(alphaN):
            if self.detailed_info:
                alphaNfifth = int(alphaN/5)
                if a%alphaNfifth == 0:
                    print('Corvariance Tensor assembly process: {:.0f}%'.format(100*a/alphaN))
            for a_ in range(alphaN):
                for i in range(input_channel_size):
                    for j in range(input_channel_size):
                        idx_i = i*alphaN + a
                        idx_j = j*alphaN + a_
                        self.CTensor[a][a_][i][j] = C[idx_i][idx_j]
        if self.detailed_info:
            print('Corvariance Tensor assembly complete')

    def UpdateDeltaTensor(self):
        Cl = self.struct[self.currentLayerIdx]
        Clplus1 = self.struct[self.currentLayerIdx + 1]
        current_size = int(np.sqrt(self.CTensor.shape[0]))
        delta_size = (current_size - self.kernel_size + 1 )*(current_size - self.kernel_size + 1 )# alpha range
        self.Delta = np.zeros([delta_size,delta_size,Clplus1,Clplus1])
        # each loop update a matrix of size [C(l+1),C(l+1)]
        # alpha alpha_ are identified according to Natural Ordering
        for alpha in range(delta_size):
            for alpha_ in range(delta_size):
                alphai = alpha//(current_size - self.kernel_size + 1 )
                alphaj = alpha%(current_size - self.kernel_size + 1 )
                alpha_i = alpha_//(current_size - self.kernel_size + 1 )
                alpha_j = alpha_%(current_size - self.kernel_size + 1 )
                for i in range(self.kernel_size):
                    for j in range(self.kernel_size):
                         for i_ in range(self.kernel_size):
                             for j_ in range(self.kernel_size):
                                c_tensor_i = (alphai+i)*current_size + alphaj + j   # shrink 2D into 1D
                                c_tensor_j = (alpha_i+i_)*current_size + alpha_j + j_
                                self.Delta[alpha][alpha_] += np.dot(np.dot(self.weight_inner[self.currentLayerIdx+1][i][j][0],self.CTensor[c_tensor_i][c_tensor_j]),self.weight_inner[self.currentLayerIdx+1][i_][j_][0].T)
    """
    def PhiCrossMult(self,x,y,A,B,C,A_,B_,C_):
        return self.phi(A*x+B+C)*self.phi(A_*x+B_+C_)
    """
    def IterateOneLayer(self):
        if self.detailed_info:
            print('Current layer:   {}'.format(self.currentLayerIdx))
        # 前进一层
        current_size = int(np.sqrt(self.CTensor.shape[0]))
        layer_next_size = current_size - self.kernel_size + 1
        layer_next_size2 = layer_next_size*layer_next_size # shrink to 1D
        Cin = self.struct[self.currentLayerIdx]
        Cout = self.struct[self.currentLayerIdx+1]
        self.CTensorNext = np.zeros([layer_next_size2,layer_next_size2,Cout,Cout])
        self.h_mean_next = []
        for i in range(layer_next_size2):
            self.h_mean_next.append(np.zeros([self.struct[self.currentLayerIdx+1],1]))
        # 将alpha 拆成2D alphai alphaj二维坐标
        # first 4 loops for feature map's positions of layer l+1
        for alphai in range(layer_next_size):
            if self.detailed_info:
                layer_next_sizefifth = int(layer_next_size/5)
                if alphai%layer_next_sizefifth == 0:
                    print('Iterating process: {:.0f}%'.format(100*alphai/layer_next_size))
            for alphaj in range(layer_next_size):
                # update mean h at position alpha
                alpha = alphai*layer_next_size + alphaj
                for i in range(Cout):
                    x = np.random.randn(1,20000)
                    # second term in phi
                    tmp = 0
                    for m in range(self.kernel_size):
                        for n in range(self.kernel_size):
                            # i stands for ith element in layer l+1
                            # self.weight_inner[self.currentLayerIdx+1][m][n][0]: C(l+1)*C(l) size 
                            tmp += np.dot(self.weight_inner[self.currentLayerIdx+1][m][n][0],self.h_mean[(alphai+m)*current_size+(alphaj+n)])[i][0]
                    Phi1 = self.phi(np.sqrt(self.Delta[alpha][alpha][i][i])*x + tmp + self.bias[self.currentLayerIdx+1][i][0])
                    self.h_mean_next[alpha][i][0] = Phi1.sum()/20000
        for alphai in range(layer_next_size):
            for alphaj in range(layer_next_size):
                alpha = alphai*layer_next_size + alphaj
                # calculate corvariance matrix C(alpha,alpha',i,j) 
                # 将alpha' 拆成2D alpha_i alpha_j二维坐标
                for alpha_i in range(layer_next_size):
                    for alpha_j in range(layer_next_size): # above loops for neurals positions
                        # this loop stands for alpha prime index
                        alpha_ = alpha_i*layer_next_size + alpha_j
                        for i in range(Cout):
                            # i indicates the same index of output layer l+1 with m(alpha,i)
                            x = np.random.randn(1,20000)
                            for j in range(Cout):
                                y = np.random.randn(1,20000)
                                # second term in phi_
                                tmp_ = 0
                                tmp = 0
                                for m in range(self.kernel_size):
                                    for n in range(self.kernel_size):
                                        tmp += np.dot(self.weight_inner[self.currentLayerIdx+1][m][n][0],self.h_mean[(alphai+m)*current_size+(alphaj+n)])[i][0]
                                        tmp_ += np.dot(self.weight_inner[self.currentLayerIdx+1][m][n][0],self.h_mean[(alpha_i+m)*current_size+(alpha_j+n)])[j][0]
                                PHI = self.Delta[alpha][alpha_][i][j]/np.sqrt(self.Delta[alpha][alpha][i][i]*self.Delta[alpha_][alpha_][j][j])
                                if self.Delta[alpha][alpha][i][i]*self.Delta[alpha_][alpha_][j][j] < 0:
                                    print('negtive root')
                                if 1-PHI*PHI<0:
                                    print("negtive error")
                                    print(PHI)
                                A = (self.phi(np.sqrt(self.Delta[alpha][alpha][i][i])*x + tmp + self.bias[self.currentLayerIdx+1][i][0])*self.phi(np.sqrt(self.Delta[alpha_][alpha_][j][j])*(PHI*x+np.sqrt(1-PHI*PHI)*y) + tmp_ + self.bias[self.currentLayerIdx+1][j][0])).sum()/20000
                                B = self.h_mean_next[alpha][i][0]*self.h_mean_next[alpha_][j][0]
                                self.CTensorNext[alpha][alpha_][i][j] = A - B
        self.currentLayerIdx = self.currentLayerIdx + 1
        self.CTensor = self.CTensorNext.copy()
        self.h_mean = self.h_mean_next.copy()
    
    def CorEigenSignCheck(self):
        # check the ratio of the negtive eigenvalue, called before iterating ahead
        num_current_channel = self.struct[self.currentLayerIdx]
        ratioLayer = []
        for c in range(num_current_channel):
            CMat = self.CTensor[:,:,c,c]
            eigenvalue, eigenvector = np.linalg.eig(CMat)
            minu_num = 0
            for n in range(eigenvalue.shape[0]):
                if eigenvalue[n] < 0:
                    minu_num = minu_num + 1
            ratioLayer.append(minu_num/eigenvalue.shape[0])
        self.ratio.append(ratioLayer)

    # extract the dimensionality of current layer
    # dimList[layer index][channel index] is the normalized dimensionality in layer [layer index] at channel [channel index]
    def UpdateDimensionality(self):
        # dimensionality in each layer and the dimensionality of the Big Matrix
        # Called before iterating to the next layer
        num_current_channel = self.struct[self.currentLayerIdx]
        num_feature_map = self.CTensor.shape[0]
        dimPerLayer = []
        if self.eigen_check:
            self.CorEigenSignCheck()
        for c in range(num_current_channel):
            CMat = self.CTensor[:,:,c,c]
            Ctrace = np.trace(CMat)
            dim = Ctrace*Ctrace/np.trace(np.dot(CMat,CMat))/num_feature_map
            dimPerLayer.append(dim)
        self.dimList.append(dimPerLayer)
        # dimensionality of Big Matrix
        big_matrix_edge = num_current_channel*num_feature_map
        self.CBigMat = np.zeros([big_matrix_edge,big_matrix_edge])
        for i in range(num_current_channel):
            for j in range(num_current_channel):
                for alpha in range(num_feature_map):
                    for alpha_ in range(num_feature_map):
                        I = i*num_feature_map + alpha
                        J = j*num_feature_map + alpha_
                        self.CBigMat[I][J] = self.CTensor[alpha][alpha_][i][j]
        Ctrace = np.trace(self.CBigMat)
        dim = Ctrace*Ctrace/np.trace(np.dot(self.CBigMat,self.CBigMat))/big_matrix_edge
        self.dimBigMatrix.append(dim)
    
    # print out dimensionality information layer by layer
    def PrintDimInfo(self):
        print('Block matrix dimensionality:')
        for l in range(self.layer_len):
            print('Layer {}'.format(l))
            for c in range(self.struct[l]):
                print('Channel {0} has dimensionality {1} '.format(c,self.dimList[l][c]))
        print('Big Matrix dimensionality:')
        for l in range(self.layer_len):
            print('Layer {0} has total dimensionality {1}'.format(l,self.dimBigMatrix[l]))


# API linked to TrueDataTrain
# Load the trained parameters and transform it to the format required by "DMFNet" class
# return [struct,weight,bias] as constructor's parameter
def DMFReadParam(fpath):
    # [Cout,Cin,kernel_size,kernel_size] for weight and [Cout,] (not column vector) for bias  
    f = open(fpath,'rb')
    ParamList = pickle.load(f)
    # get the layers of CNN part and layer's structure
    lenOfCNN = 0
    struct = [0]
    for i in range(len(ParamList)):
        # i increase by 2 due to interleave structure of ParamList
        if ParamList[i].ndim > 2:
            lenOfCNN += 1
            struct.append((ParamList[i].shape)[0])
    if lenOfCNN > 0:
        lenOfCNN += 1 
        struct[0] = (ParamList[0].shape)[1]
    # reconstruct weight list
    if lenOfCNN == 0:
        print('Input net does not have CNN structure')
        sys.exit()
    weight = []
    weight.append('0')
    # construct the sketch
    for n in range(lenOfCNN - 1):
        C_in = struct[n]
        C_out = struct[n+1]
        weightPerLayer = [[[] for j in range(C_in)] for i in range(C_out)]
        weight.append(weightPerLayer)
    for n in range(lenOfCNN - 1):
        C_in = struct[n]
        C_out = struct[n+1]
        for cout in range(C_out):
            for cin in range(C_in):
                weight[n+1][cout][cin].append(ParamList[2*n][cout,cin,:,:])
    # reconstruct bias list
    bias = ['0']
    for n in range(lenOfCNN - 1):
        bias.append(ParamList[2*n+1].reshape(ParamList[2*n+1].shape[0],1))
    return struct, weight, bias
        
            
            




        






