import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis


def LS_1(w,M,r_1,r_2):
    
    LS_1_value_one = np.sum(np.sign(np.dot(w,M.T)))
    LS_1_value_two = np.sum(np.sign(np.dot(-w,M.T)))
    LS_1_value = max(LS_1_value_one,LS_1_value_two)/(r_1*r_2)
    
    return LS_1_value

def LS_2(w,M,r_1,r_2):
    
    L2_up = np.abs(np.sum(np.dot(w,M.T)))
    L2_down = np.sum(np.abs(np.dot(w,M.T)))
    LS_2_value = L2_up/L2_down
    
    return LS_2_value
    
def J_w(w,M,r_1,r_2):
    
    J_w_up = np.square(np.sum(np.dot(w,M.T)))
    J_w_down = np.sum(np.square(np.dot(w,M.T)))
    J_w_value = J_w_up/J_w_down
    
    return J_w_value

def LDA(w,new_data_A,new_data_B,r_1,r_2):
    
    w = w.T
    
#     A_class = new_data_A.numpy()
#     B_class = new_data_B.numpy()
    A_class = new_data_A.T
    B_class = new_data_B.T

    nu_a = np.mean(A_class,axis=1).reshape(-1,1)
    nu_b = np.mean(B_class,axis=1).reshape(-1,1)

    A_c = A_class-np.repeat(nu_a,A_class.shape[1],axis=1)
    B_c = B_class-np.repeat(nu_b,B_class.shape[1],axis=1)

    S_w = np.dot(A_c,A_c.T)+np.dot(B_c,B_c.T)
    
    S_b = np.dot(nu_a-nu_b,(nu_a-nu_b).T)


    LDA_value_up = np.dot(np.dot(w.T,S_b),w)
    LDA_value_down = np.dot(np.dot(w.T,S_w),w)
    LDA_value = LDA_value_up/LDA_value_down
    
    return LDA_value



def W(X,Y):

    X1 = X[np.where(Y==1)[0]]
    X2 = X[np.where(Y==0)[0]]
    r_1 = len(X1)
    r_2 = len(X2)
    
    
    M = np.zeros((r_1*r_2,X.shape[1]),dtype=np.float32)

    index_base = np.arange(r_2)
    
    for i in range(r_1):
        M[index_base+i*r_2,:]=X1[i]-X2

    m = np.sum(M,axis=0).reshape(1,-1)
    w = m/np.linalg.norm(m)
    
    #print(w.shape)
    LS_1_value = LS_1(w,M,r_1,r_2)
    LS_2_value = LS_2(w,M,r_1,r_2)
    J_w_value = J_w(w,M,r_1,r_2)
    LDA_value = LDA(w,X1,X2,X,Y)
    
    return LS_1_value,LS_2_value,J_w_value,LDA_value


def get_layer_name(model):
    
    layer_name_list = []
    number_squence = [0,'1st','2nd','3rd','4th','5th','6th','7th','8th']

    name_conts = np.ones((6,),dtype='int')
    
    for layer in model.layers: 

        if 'conv2d' in layer.name: 
            name_conts_layer = name_conts[0]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'conv2d')
            name_conts[0]+=1

        if 'max_pooling' in layer.name: 
            name_conts_layer = name_conts[1]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'max_pooling')
            name_conts[1]+=1

        if 'flatten' in layer.name: 
            name_conts_layer = name_conts[2]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'flatten')    
            name_conts[2]+=1

        if 'dense' in layer.name:
            name_conts_layer = name_conts[3]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'dense')  
            name_conts[3]+=1

        if 'batch_normalization' in layer.name: 
            name_conts_layer = name_conts[4]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'batch_normalization')     
            name_conts[4]+=1

        if 'dropout' in layer.name: 
            name_conts_layer = name_conts[5]
            layer_name_list.append(number_squence[name_conts_layer]+'_'+'dropout')  
            name_conts[5]+=1
        
    return layer_name_list


def plot_Separability_figure(layer_name_list,x_plot,LS_1_squence,LS_2_squence,J_w_squence,LDA_squence,LS_1_squence_base,LS_2_squence_base,J_w_squence_base,LDA_squence_base):
    
    
    #figure = plt.figure(1,figsize=(8,8))
    figure = plt.figure(figsize=(16,10))
    #layer_name_list = get_layer_name(model)[reserved_layers:]
    
    
    # LS_0
    ax1 = plt.subplot(2,2,1)

    for i in range(len(LS_1_squence)):
#         if 'flatten' not in layer_name_list[i]:
           plt.plot(x_plot,LS_1_squence[i,:],label=layer_name_list[i])
    plt.plot(x_plot,LS_1_squence_base[0,:],'-k',label='base_line')
    plt.title('$\mathrm{LS}_1$')
    plt.legend(bbox_to_anchor=(1,0),loc=3)


    # LS_1
    ax2 = plt.subplot(2,2,2)

    for i in range(len(LS_2_squence)):
#         if 'flatten' not in layer_name_list[i]:
           plt.plot(x_plot,LS_2_squence[i,:],label=layer_name_list[i])
    plt.plot(x_plot,LS_2_squence_base[0,:],'-k',label='base_line')
    plt.title('$\mathrm{LS}_2$')
    plt.legend(bbox_to_anchor=(1,0),loc=3)


    # LS_2
    ax3 = plt.subplot(2,2,3)

    for i in range(len(J_w_squence)):
#         if 'flatten' not in layer_name_list[i]:
           plt.plot(x_plot,J_w_squence[i,:],label=layer_name_list[i])
    plt.plot(x_plot,J_w_squence_base[0,:],'-k',label='base_line')
    plt.title('$\mathrm{J}_w$')
    plt.legend(bbox_to_anchor=(1,0),loc=3)

    # J_w
    ax4 = plt.subplot(2,2,4)

    for i in range(len(LDA_squence)):
#         if 'flatten' not in layer_name_list[i]:
           plt.plot(x_plot,LDA_squence[i,:],label=layer_name_list[i])
    plt.plot(x_plot,LDA_squence_base[0,:],'-k',label='base_line')
    plt.title('$\mathrm{LDA}$')
    plt.legend(bbox_to_anchor=(1,0),loc=3)
    
    plt.tight_layout()
    
    
    
    return figure
    
    
def plot_net_figure(layer_name_list,x_plot,train_loss_squence,train_accuracy_squence,test_loss_squence,test_accuracy_squence):

    
    #figure = plt.figure(1,figsize=(8,8))
    figure = plt.figure(figsize=(16,10))
    #layer_name_list = get_layer_name(model)[reserved_layers:]
    
    ax1 = plt.subplot(2,2,1)
    plt.plot(x_plot,train_loss_squence)
    plt.title('Train Loss')

    ax2 = plt.subplot(2,2,2)
    plt.plot(x_plot,train_accuracy_squence)
    plt.title('Train Accuracy')

    ax3 = plt.subplot(2,2,3)
    plt.plot(x_plot,test_loss_squence)
    plt.title('Test Loss')


    ax4 = plt.subplot(2,2,4)
    plt.plot(x_plot,test_accuracy_squence)
    plt.title('Test Accuracy')

    
    plt.tight_layout()
    
    
    
    return figure

