"""
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
from functools import partial

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)

            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['weight']
            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['weight']
            for other, edge in graph[node].items()
            if other not in node.community
        )

    def unique_interactions(self, node):
        """Return 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]['weight']
            for a, b in product(community_a, community_b)
            if graph.has_edge(a, b)
        )

    def nodal_contribution(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]['weight']
            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."""
        # Fixes Compatability for python-louvain
        # Nodes cannot be deepcloned because they are generators. To fix this
        # issue we use id(Node) as the nodes for a new graph.
        graph = nx.Graph()
        nodes = {
            id(node): node
            for node in self.graph_old.nodes()
        }
        graph.add_nodes_from(nodes)
        for node_a, node_b, weight in self.graph_old.edges(data='weight'):
            graph.add_edge(id(node_a), id(node_b), {'weight': weight})

        partitions = louvain_partition(graph)

        communities = defaultdict(set)
        for node, community in partitions.items():
            # get actual node from it's id
            communities[community].add(nodes[node])

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


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

    def __init__(self, community=None, **options):
        """Create a community based node."""
        super().__init__(**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
        ]


def _decide(node, others, key):
    """
    Make a decision on which node to send to best is decided based on key.

    If the best is better than node, return best.
    If node is better than the best, return node.
    If the best and node are equal, return None.
    """
    best = max(others, key=key)
    best_key = key(best)
    node_key = key(node)

    if best_key > node_key:
        return best

    if best_key < node_key:
        return node

    return None


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

        if self.in_community_neighbours:
            target = _decide(self,
                             self.in_community_neighbours,
                             self._community.local_popularity)
            if not (target is None or target is self):
                return {target: 'local-popularity'}

        elif self.out_community_neighbours:
            target = _decide(self,
                             self.out_community_neighbours,
                             self._community.global_popularity)
            if not (target is None or target is self):
                return {target: 'global-popularity'}

        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

        dest = packet.destination

        if self.community != dest.community:
            if self.out_community_neighbours:
                cbc = partial(self._community.community_betweenness, dest)
                target = _decide(self,
                                 self.out_community_neighbours,
                                 cbc)
                if not (target is None or target is self):
                    return {target: 'community-betweenness'}

            if self.in_community_neighbours:
                ncf = partial(self._community.nodal_contribution, dest)
                target = _decide(self,
                                 self.in_community_neighbours,
                                 ncf)
                if not (target is None or target is self):
                    return {target: 'nodal-contribution'}

        if self.in_community_neighbours:
            target = _decide(self,
                             self.in_community_neighbours,
                             self._community.unique_interactions)
            if not (target is None or target is self):
                return {target: 'unique-interactions'}

            target = _decide(self,
                             self.in_community_neighbours,
                             self._community.local_popularity)
            if not (target is None or target is self):
                return {target: 'local-popularity'}

        return {}