"""
pydtn community module.

Implements the following community detection schemes on epochs:
- Louvain community detection
- K-Clique community detection

"""

__all__ = [
    "Community",
    "KCliqueCommunity",
    "LouvainCommunity",

    "CommunityNode",
    "BubbleNode",
    "HCBFNode",
]
__version__ = '0.1'
__author__ = 'Jarrod Pas <j.pas@usask.ca>'

from collections import defaultdict
from itertools import product

import networkx as nx
from community import best_partition as louvain_partition

from pydtn import Node


class Community:
    """
    Base class for community detection algorithms.

    Implements partitioning on epoch turn over.
    To implement your own community detection algorithm inherit from this class
    and implement `partition`. You may need to extend other methods.
    """

    def __init__(self, epoch):
        """Create a community detector."""
        self.network = None
        self.epoch = epoch

        self.graph = nx.Graph()
        self.graph_old = None

        self._community = {}

    def start(self, network):
        """
        Start community detection loop.

        If it has already been started do nothing.
        """
        if self.network is not None:
            return

        self.network = network
        self.network.env.process(self.epoch_loop())

    def epoch_loop(self):
        """Transition current epoch to next epoch."""
        env = self.network.env
        while True:
            yield env.timeout(self.epoch)
            print('epoch')

            new = nx.Graph()
            for node_a, node_b, start in self.graph.edges(data='start'):
                if start is not None:
                    self.leave(node_a, node_b)
                    self.join(node_a, node_b, graph=new)

            self.graph_old = self.graph
            self.graph = new

            for community in self.partition():
                for node in community:
                    self._community[node] = community

    def partition(self):
        """Partition nodes into communities."""
        for node in self.graph_old.node():
            yield frozenset([node])

    def join(self, node_a, node_b, graph=None):
        """
        Node a and b join each other's neighbourhoods.

        This is an idempotent operation.
        """
        if graph is None:
            if self.graph is None:
                return
            graph = self.graph

        if not graph.has_edge(node_a, node_b):
            graph.add_edge(node_a, node_b, {
                'start': None,
                'weight': 0,
            })

        edge = graph[node_a][node_b]
        if edge['start'] is None:
            edge['start'] = self.network.env.now

    def leave(self, node_a, node_b, graph=None):
        """
        Node a and b leave each other's neighbourhoods.

        This is an idempotent operation.
        """
        if graph is None:
            if self.graph is None:
                return
            graph = self.graph

        if graph.has_edge(node_a, node_b):
            edge = graph[node_a][node_b]
            if edge['start'] is not None:
                edge['weight'] += self.network.env.now - edge['start']
                edge['start'] = None

    def __getitem__(self, node):
        """
        Return the community of a node.

        If the node has no community it is it's own all alone.
        """
        if node not in self._community:
            self._community[node] = frozenset(list(node))
        return self._community[node]

    def local_popularity(self, node):
        """Return local popularity of a node."""
        graph = self.graph_old

        if node not in graph:
            return 0

        return sum(
            edge['duration']
            for other, edge in graph[node].items()
            if other in node.community
        )

    def global_popularity(self, node):
        """Return global popularity of a node."""
        graph = self.graph_old

        if node not in graph:
            return 0

        return sum(
            edge['duration']
            for other, edge in graph[node].items()
            if other not in node.community
        )

    def unique_interactions(self, node):
        """Unique interactions for a node within it's community."""
        graph = self.graph_old

        if node not in graph:
            return 0

        return len(
            other
            for other in node.community
            if graph.has_edge(node, node)
        )

    def community_betweenness(self, node_a, node_b):
        """Return community betweenness for 2 nodes."""
        graph = self.graph_old
        community_a = self[node_a]
        community_b = self[node_b]

        if community_a == community_b:
            return float('inf')

        return sum(
            graph[a][b]['duration']
            for a, b in product(community_a, community_b)
            if graph.has_edge(a, b)
        )

    def nodal_contribution_factor(self, node_a, node_b):
        """Return nodal contribution factor of node_a in node_b's community."""
        graph = self.graph_old
        community_b = self[node_b]

        if node_a not in graph or node_a in community_b:
            return 0

        return sum(
            graph[node_a][b]['duration']
            for b in community_b
            if graph.has_edge(node_a, b)
        )


class KCliqueCommunity(Community):
    """K-clique community detection."""

    def __init__(self, epoch, k):
        """Create a k-clique community detector."""
        super().__init__(epoch)
        self.k = k

    def partition(self):
        """Partition graph by k-clique communities method."""
        return nx.k_clique_communities(self.graph_old, k=self.k)


class LouvainCommunity(Community):
    """Louviain community detection."""

    def partition(self):
        """Partition graph with Louvain algorithm."""
        partitions = louvain_partition(self.graph_old)

        communities = defaultdict(set)
        for node, community in partitions.items():
            communities[community].add(node)

        for community in communities.values():
            yield frozenset(community)

class CommunityNode(Node):
    """Base node for community based forwarding heuristics."""

    def __init__(self, name, community=None, **options):
        """Create a community based node."""
        super().__init__(name, **options)
        if community is None:
            raise ValueError('No community set')
        self._community = community

    def start(self, network):
        """Start event loop for node and community detector."""
        super().start(network)
        self._community.start(network)

    def join(self, node):
        """
        Approach the neighbouhood of this node.

        This is an idempotent operation.
        """
        super().join(node)
        self._community.join(self, node)

    def leave(self, node):
        """
        Leave the neighbouhood of this node.

        This is an idempotent operation.
        """
        super().leave(node)
        self._community.leave(self, node)

    @property
    def community(self):
        """Return my community."""
        return self._community[self]

    @property
    def in_community_neighbours(self):
        """Return nodes in my community."""
        return [
            neighbour
            for neighbour in self.neighbours
            if neighbour in self.community
        ]

    @property
    def out_community_neighbours(self):
        """Return nodes not in my community."""
        return [
            neighbour
            for neighbour in self.neighbours
            if neighbour not in self.community
        ]


class BubbleNode(CommunityNode):
    """Node with BubbleRap forwarding."""

    def forward(self, packet):
        """Forward packet with the BubbleRap heuristic."""
        # direct forwarding
        forward = super().forward(packet)
        if forward:
            return forward

        return {}


class HCBFNode(CommunityNode):
    """Node with Hybrid Community Based forwarding."""

    def forward(self, packet):
        """Forward packet with Hybrid Community Based heuristic."""
        # direct forwarding
        forward = super().forward(packet)
        if forward:
            return forward

        return {}