import networkx as nx
import numpy as np
import pandas as pd
from igraph import *
from networkx.algorithms.community import *
from tqdm import trange

Layer_Num = 1           #Select the layer you want to calculate the adjacency matrix
top_Values = 50         #Top S neurons from each class
Neurons = number_Neurons = [512, 512]     #Total Neurons in the layer
number_Classes = 10     #Total Number of classes
iterations = 20         #Total iterations used

for Layer_Num in range(1, 3):
    Total_Neurons = Neurons[Layer_Num-1]
    modularity_Values = np.zeros((iterations, 4))
    for outer_Iter in trange(0, iterations):
        oute_outer = outer_Iter

        """
        --------------------------------------------------------------------------------------------------------------------
        Create Network------------------------------------------------------------------------------------------------------
        --------------------------------------------------------------------------------------------------------------------
        """
        dataset = []
        class_Number = number_Classes
        activated_nodes = np.zeros((Total_Neurons, class_Number))
        a = np.zeros((Total_Neurons, Total_Neurons))
        temp_Adjacency = -1 * np.ones((top_Values + 3, class_Number))
        adjacency_Matrix = np.zeros((Total_Neurons, Total_Neurons))
        for class_num in range(0, class_Number):
            count = 0
            file_path = os.path.join("../Adjacency_Matrix/Layer_" + str(Layer_Num) +
                                  "/" + "Top_" + str(top_Values) + "_Edge_Frequency_Iteration_" + str(outer_Iter) + "_class_" + str(class_num) + ".csv")
            dataset = pd.read_csv(file_path, header=None)
            dr, dc = dataset.shape
            activation_val = np.array(dataset)  # activation_val stores the activavtion of the neurons in a class in a layer
            adjacency_Matrix += activation_val
            temp = []
            for row in range(dr):
                for column in range(dc):
                    if activation_val[row, column] > 0:
                        temp.append(row)
                        temp.append(column)
            temp = np.unique(temp)
            for r in range(len(np.unique(temp))):
                temp_Adjacency[r, class_num] = temp[r]  # temp_Adjacency stores the activated neurons of each class

        adjacency_Matrix = 1/adjacency_Matrix
        adjacency_Matrix[adjacency_Matrix == np.inf] = 0
        unique_Node_Values = (np.unique(np.ravel(np.copy(temp_Adjacency))))
        if unique_Node_Values[0] == -1:
            unique_Node_Values = unique_Node_Values[1:]

        # adjacency_Matrix = adjacency_Matrix + adjacency_Matrix.T - np.diag(np.diag(adjacency_Matrix))
        g = Graph()
        g.add_vertices(len(unique_Node_Values))

        weights = []
        for row in range(Total_Neurons):
            for column in range(Total_Neurons):
                if adjacency_Matrix[row, column] > 0:
                    for ite in range(len(unique_Node_Values)):
                        if int(row) == int(unique_Node_Values[ite]):
                            start_Node = ite
                        if int(column) == int(unique_Node_Values[ite]):
                            end_Node = ite
                    g.add_edges([(start_Node, end_Node)])
                    weights.append(adjacency_Matrix[row, column])
        g.simplify(multiple=True, loops=True)
        g.es['weight'] = weights
        node_Membership = -1 * np.ones((len(unique_Node_Values), class_Number))
        for ite in range(len(unique_Node_Values)):
            for ite2 in range(class_Number):
                if int(unique_Node_Values[ite]) in temp_Adjacency[:, ite2] and int(unique_Node_Values[ite]) != -1:
                    node_Membership[ite, ite2] = int(ite2)
        node_Membershsip_Final = -np.sort(-node_Membership, axis=1)
        adjacency_Matrix_Both_Graph = np.array(Matrix._get_data(Graph.get_adjacency(g, type=2)))

        """
        --------------------------------------------------------------------------------------------------------------------
        NetworkX Network and Modularity Functions---------------------------------------------------------------------------
        --------------------------------------------------------------------------------------------------------------------
        """
        """L = g.layout("mds")
        plot(g, layout=L, bbox=[1200, 1200], margin=60, vertex_label_size=10)"""
        g.write_gml('Graph_Iteration_'+str(outer_Iter)+'_Layer_'+str(Layer_Num)+'.gml')
        H = nx.read_gml('Graph_Iteration_'+str(outer_Iter)+'_Layer_'+str(Layer_Num)+'.gml', label=None)

        # Kernighan Lin Bisection algorithm
        modularity_Values[oute_outer, 0] = modularity(H, kernighan_lin_bisection(H, max_iter=100, weight="weight"))

        """
        --------------------------------------------------------------------------------------------------------------------
        NO-OVERLAP, NO-WEIGHT-----------------------------------------------------------------------------------------------
        --------------------------------------------------------------------------------------------------------------------
        """
        Nnode_Number = len(g.vs())
        node_Number = len(g.vs())
        """
        :parameter
        m       = Total edges in the network
        modu    = Modularity of the network
        i, j    = Any two nodes in the network
        """
        adjacency_Matrix_Both =  adjacency_Matrix + adjacency_Matrix.T - np.diag(np.diag(adjacency_Matrix))
        adjacency_Matrix_Weighted = np.copy(adjacency_Matrix_Both)
        adjacency_Matrix_Both[adjacency_Matrix_Both>0] = 1
        modu_Normal = 0
        m = np.sum(np.sum(adjacency_Matrix_Both))
        node_Degree = np.sum(adjacency_Matrix_Both, axis=0)
        for i in range(Nnode_Number):
            for j in range(Nnode_Number):
                if i != j and adjacency_Matrix_Both[i, j] != 0:
                    for ite in range(class_Number):
                        if node_Membershsip_Final[i, ite] != -1 and node_Membershsip_Final[j, ite] != -1:
                            if node_Membershsip_Final[i, ite] in node_Membershsip_Final[j, :]:
                                modu_Normal = modu_Normal + (
                                        adjacency_Matrix_Both[i, j] - (node_Degree[i] * node_Degree[j]) / (2 * m))
        modu_Normal = modu_Normal * (1 / (2 * m))
        modularity_Values[oute_outer, 1] = modu_Normal

        """
        --------------------------------------------------------------------------------------------------------------------
        OVERLAP, NO-WEIGHT--------------------------------------------------------------------------------------------------
        --------------------------------------------------------------------------------------------------------------------
        """
        """
        :parameter
        m       = Total edges in the network
        modu    = Modularity of the network
        i, j    = Any two nodes in the network
        """
        modu_Shen = 0
        m = np.sum(np.sum(adjacency_Matrix_Both))
        node_Degree = np.sum(adjacency_Matrix_Both, axis=0)
        for i in range(node_Number):
            for j in range(node_Number):
                if i != j and adjacency_Matrix_Both[i, j] != 0:
                    for ite in range(class_Number):
                        if node_Membershsip_Final[i, ite] != -1:
                            if node_Membershsip_Final[i, ite] in node_Membershsip_Final[j, :]:
                                o_I = np.copy(node_Membershsip_Final[i, :])
                                o_I[o_I != -1] = 1
                                o_I[o_I == -1] = 0
                                o_I = np.sum(np.sum(o_I))
                                o_J = np.copy(node_Membershsip_Final[j, :])
                                o_J[o_J != -1] = 1
                                o_J[o_J == -1] = 0
                                o_J = np.sum(np.sum(o_J))
                                modu_Shen = modu_Shen + (
                                        adjacency_Matrix_Both[i, j] - ((node_Degree[i] * node_Degree[j]) / (2 * m))) / (
                                                    o_I * o_J)
        modu_Shen = modu_Shen * (1 / (2 * m))
        modularity_Values[oute_outer, 2] = modu_Shen

        """
        --------------------------------------------------------------------------------------------------------------------
        WEIGHTED, OVERLAP---------------------------------------------------------------------------------------------------
        --------------------------------------------------------------------------------------------------------------------
        """
        # Calculate Weighted Adjacency Matrix
        node_Neighbourhood = np.ones(((node_Number), node_Number)) * -1
        node_Degree_Weighted = np.sum(adjacency_Matrix_Weighted, axis=0)
        for ite in range(node_Number):
            ite3 = 0
            for ite2 in range(node_Number):
                if adjacency_Matrix_Both[ite, ite2] != 0:
                    node_Neighbourhood[ite, ite3] = ite2
                    ite3 += 1
        belonging_Coefficient = np.zeros((node_Number, class_Number))
        for ite in range(node_Number):
            for ite2 in range(node_Number):
                if node_Neighbourhood[ite, ite2] != -1:
                    for ite3 in range(class_Number):
                        node_2_Communities = np.copy(node_Membershsip_Final[int(node_Neighbourhood[ite, ite2]), ite3])
                        if node_2_Communities != -1:
                            belonging_Coefficient[ite, int(node_2_Communities)] = belonging_Coefficient[
                                                                                      ite, int(node_2_Communities)] + \
                                                                                  adjacency_Matrix_Weighted[ite, int(
                                                                                      node_Neighbourhood[ite, ite2])]

        total_Belonging_Coefficient = np.sum(belonging_Coefficient, axis=1)

        modu_Chen = 0
        m = np.sum(np.sum(adjacency_Matrix_Both))
        node_Degree = np.sum(adjacency_Matrix_Both, axis=0)
        for ite in range(node_Number):
            for ite2 in range(node_Number):
                for ite3 in range(class_Number):
                    if ite != ite2:
                        if total_Belonging_Coefficient[ite] == 0 or total_Belonging_Coefficient[ite2] == 0:
                            alpha_U = 0
                            alpha_V = 0
                        else:
                            alpha_U = belonging_Coefficient[ite, ite3] / total_Belonging_Coefficient[ite]
                            alpha_V = belonging_Coefficient[ite2, ite3] / total_Belonging_Coefficient[ite2]
                        modu_Chen += (adjacency_Matrix_Both[ite, ite2] - (
                                (node_Degree[ite] * node_Degree[ite2]) / (2 * m))) * (alpha_U * alpha_V)
        modu_Chen = modu_Chen * (1 / (2 * m))
        modularity_Values[oute_outer, 3] = modu_Chen
    file_P = os.path.join("../Results/Modularity/Layer_" + str(Layer_Num) + "/" + "All_Modularity_InverseW_Correct_Adj.csv")
    np.savetxt(file_P, modularity_Values, delimiter=",", header = "KLB, NOVERLAP, OVERLAP, WEIGHTLAP" , fmt="%f")
print('DONE!!!')