import networkx.algorithms.community as nx_comm
import networkx as nx

""" 
Visualize network 
input : result dict file for return network 
return : saving network in form of group : nodes  ex) 1 : 2 3 5 11 13 ...
"""

class Visualization() :

    def __init__(self, nwtk_dict, root = None, vis_name = None, netX = False, edgelist = None):

        new_pair_dict = self.new_pairs_(nwtk_dict)

        # saving the files as .dat format
        nx_mid_list = self.gen_comm_node(list(new_pair_dict.values()))
        self.saving_new_g(root, vis_name, "dat", nx_mid_list)
        print("saving success!")

        if netX is True :

            # required format for modularity on networkx
            # nx_comm_list_2 = nx_list(list(new_pair_dict.values()))
            """ 
            Check modularity with networkX
            """
            nx_G = nx.read_edgelist(edgelist)
            modul_nx = nx_comm.modularity(nx_G, nx_comm.label_propagation_communities(nx_G))
            print("Modularity from networkX %d", modul_nx)

    def new_pairs_(self,p_dict) :
        pair_dict = {}
        for idx in range(len(p_dict)) :

            if p_dict[idx] not in pair_dict.keys() :
                pair_dict[p_dict[idx]] = []
                pair_dict[p_dict[idx]].append(idx)
            else :
                pair_dict[p_dict[idx]].append(idx)

        return pair_dict

    def gen_comm_node(self,comm_vals):
        comm_list = []
        for comms in comm_vals:
            new = ""
            idx = 1
            for i in range(len(comms)):
                if idx < len(comms):
                    new += str(comms[i]) + " "
                    idx += 1
                else:
                    new += str(comms[i])
            # print(new)
            new += "\n"
            comm_list.append(new)
        return comm_list

    def saving_new_g(self,root, txt_name, extension, newList):
        #need to change here
        newG = open(root + txt_name + "." + extension, "w")
        text_list = newList
        newG.writelines(text_list)
        return print("Processed!")

    # list for networkx modularity
    def nx_list(self, comm_vals):

        nx_list = []
        for comm in comm_vals:
            int_comm = list(map(int, comm))
            new = set(int_comm)
            nx_list.append(new)

        return nx_list

if __name__ == '__main__' :

    vis = Visualization(nwtk_dict= "provide result dict", root = "/Users/macbookpro/Constrained_modularity/python/generated/", vis_name= "provide name for file")

