"""
IGraph wrapper for add node easily
"""

import igraph
import leidenalg


class IGraphWrapper:

    def __init__(self):
        self.g = igraph.Graph(directed=False)
        self.id2zid = {}
        self.zid2id = {}

    def add_node(self, zid):
        self.add_vertex(zid, quiet=True)

    @property
    def nodes(self):
        for zid in self.zid2id:
            yield zid

    def add_vertex(self, zid, quiet=False):
        if zid in self.zid2id:
            if quiet:
                return
            else:
                raise ValueError(f"{zid} was in graph")

        self.g.add_vertex(zid)
        gid = len(self.id2zid)
        self.id2zid[gid] = zid
        self.zid2id[zid] = gid

    def add_vertex_from(self, zids, quiet=False):
        for zid in zids:
            self.add_vertex(zid, quiet)

    def add_edge(self, src_id, dst_id):
        if src_id not in self.zid2id:
            raise ValueError(f"Node {src_id} is not in graph")

        if dst_id not in self.zid2id:
            raise ValueError(f"Node {dst_id} is not in graph")

        return self.g.add_edge(self.zid2id[src_id], self.zid2id[dst_id])

    def add_edge_from(self, edges):
        self.g.add_edges([(self.zid2id[x[0]], self.zid2id[x[1]]) for x in edges])

    def neighbors(self, zid):
        if zid not in self.zid2id:
            raise ValueError(f"Can not found vertex id: {zid}")

        neighbors = self.g.neighbors(self.zid2id[zid])
        return [self.id2zid[x] for x in neighbors]

    def run_personalize_pagerank(self, reset_vertices, damping_factor=0.85, niter=1000):
        root_nodes = [self.zid2id[zid] for zid in reset_vertices if zid in self.zid2id]
        result = self.g.personalized_pagerank(reset_vertices=root_nodes,
                                              directed=False,
                                              damping=damping_factor,
                                              niter=niter)

        ppr = {}
        for i in range(len(result)):
            ppr[self.id2zid[i]] = result[i]

        return ppr

    def has_node(self, zid):
        return zid in self.zid2id

    def has_edge(self, zid_src, zid_dst):
        return self.g.get_eid(self.zid2id[zid_src], self.zid2id[zid_dst], False, error=False) > -1

    def number_of_edges(self):
        return len(self.g.get_edgelist())

    def number_of_nodes(self):
        return len(self.zid2id)

    def louvain(self):
        communities = leidenalg.find_partition(self.g, leidenalg.SignificanceVertexPartition, n_iterations=2)
        communities = [[self.id2zid[x] for x in community] for community in communities]
        return communities

    def remove_redundant_graph(self):
        delete_vertices = []
        for node in self.id2zid:
            if self.g.degree(node) == 1:
                delete_vertices.append(node)

        self.g.delete_vertices(delete_vertices)
        print("Delete vertices: {0}".format(delete_vertices))

if __name__ == "__main__":
    g = IGraphWrapper()
    g.add_vertex(1)
    g.add_vertex(2)
    g.add_vertex(3)
    g.add_vertex(5)
    g.add_vertex(4)
    g.add_vertex(6)
    g.add_edge(1, 2)
    g.add_edge(2, 6)

    g.add_edge(3, 5)
    g.add_edge(4, 5)
    g.add_edge(3, 4)

    a = g.louvain()
    print(a[0])
    print(a[1])
    print(g.has_edge(1, 2))