from community import best_partition as louvain_partition

from .base import EpochCommunity

# TODO: look into other louvain implementations

class LouvainCommunity(EpochCommunity):
    def __init__(self, epoch=604800, **kwargs):
        super().__init__(epoch, **kwargs)

    def process(self, network):
        while True:
            yield self.tick()
            g = self.next_epoch(env.now)

            p = louvain_partition(g, weight='duration')
            communities = defaultdict(set)
            for node, c in louvain_partition(g, weight='duration').items():
                communities[c].add(node)

            for community in communities.values():
                community = frozenset(community)
                for node in community:
                    self.community[node] = community