"""!
A module that contains graph reducers which use different clustering 
techniques to do so.
The goal is to cut a connected subgraph with reduced number of nodes and edges.
The idea of using clustering is to find groups of nodes, that are closer to each
other than the rest of the graph, and look for the path of a minimal cost which
goes through all the nodes from the text. Then, all clusters must be connected
in order to guarantee that the graph is connected.
"""
import abc
import networkx
import scipy
import grafon.constants as constants
from sklearn.cluster import KMeans
from collections import defaultdict
from itertools import combinations
from reducerinterface import ReducerInterface

class ClusteringReducerInterface(ReducerInterface):
  """!
  The interface for a clustering reducer class.
  """
  __metaclass__ = abc.ABCMeta

  DATA_CLUSTER = 'cluster'

  def __init__(self, options, reducer_name):
    """!
    The class in an abstract class. It isn't possible to create exactly it's 
    instance.

    @param options extra parameters read from the config file. It is used by 
      children classes.
    @type options: dict(str: str)
    @param reducer_name the name of the reducer that is actually created.
    @type reducer_name: str
    """
    super(ClusteringReducerInterface, self).__init__(options, reducer_name)

  @abc.abstractmethod
  def reduce(self, digraph, text_nodes):
    """!
    Return reduced representation of the graph containing all text's nodes.

    @param graph the directed graph to be reduced.
    @type graph: networkx.DiGraph()
    @param text_nodes the list of nodes' identifiers from a text.
    @type text_nodes: list
    """
    pass
    

class KMeansClusteringReducer(ClusteringReducerInterface):
  """!
  The class uses the KMeans algorithm for finding clusters of nodes that
  are closer to each other than to the rest of the graph.   
  """
  def __init__(self, options, reducer_name = "KMeansClusteringReducer"):
    """!
    Create the instance of the class with given parameters.

    The input reducer's name is passed to the parent's constructor. Options 
    should contain two values: k and iteration_number. K is the number of 
    clusters that will be created. The number of iterations defines how many
    time the KMeans algorithm will be restarted.
    If something is missing in the options dictionary then the following values
    are assigned to the attributes:
    - k: 4
    - iteration_number: 10

    @param options a dictionary that should contain 'k' and 'iteration_number'
    @type options: dict(str: str)|dict(str: unicode)
    @param reducer_name the name of the reducer that is actually created.
    @type reducer_name: str
    """    
    super(KMeansClusteringReducer, self).__init__(options, reducer_name)
    self.k = options.get('k', 4)
    self.n_iter = options.get('iteration_number', 10)

    # If options are taken from the config then they may be str or unicode
    # convert them to int
    if isinstance(self.k, (str, unicode)):
      self.k = int(self.k)
    if isinstance(self.n_iter, (str, unicode)):
      self.n_iter = int(self.n_iter)

    self.kmeans = KMeans(
      n_clusters = self.k, 
      n_init = self.n_iter, 
      copy_x = False)

  def reduce(self, digraph, text_nodes):
    """!
    Find the shortest paths between the nodes within the same cluster.

    The method finds k clusters based on an adjacency matrix of the original
    graph. The adjacency matrix is constructed as follows:
    1) if nodes are connected then add to the cell [from_node, to_node] one,
    2) if the from_node is in the nodes from the text, add one again,
    3) if the to_node is in the nodes from the text, add one again to that cell.
    The adjacency matrix is the input for the K-Means algorithm and labels for
    each node is the output. 
    Having the labels, it is possible to assign weights to the edges of the 
    graph. The default weight is 4, but it decreases in the following 
    situations:
    1) decrease the weight of the edge by one for each it's node that is in 
      a cluster,
    2) decrease the weight of the edge by one for each it's node that belongs
      to the original text.
    Due to the fact that K-Means assign every observation to a cluster, 
    the weight of an edge is at most 3. Clusters with no nodes from the original
    text are not considered.
    In the next step all the shortest paths between the nodes from the text
    that belongs to the same cluster are found. 
    The final step is to connect the clusters. To do so, the nodes that have
    the highest number of connected edges are selected as center nodes. 
    Then the shortest path between each pair of the center nodes is found 
    and added to the output graph.    

    @param graph the directed graph to be reduced.
    @type graph: networkx.DiGraph()
    @param text_nodes nodes from a text.
    @type text_nodes: grafon.structures.textnodes.TextNodes

    @return the reduced graph.
    @rtype: networkx.DiGraph()
    """
    # If the input graph has less nodes than the defined number of clusters
    # return just the copy of digraph
    if digraph.number_of_nodes() <= self.kmeans.n_clusters:
      return networkx.DiGraph(digraph.edges(data=True))

    # Create sparse adjacency matrix for the graph
    txt_nodes = text_nodes.graph_keys
    g_weighted = networkx.Graph()
    adjacency_matrix = self._create_sparse_adjacency_matrix(
      digraph, 
      txt_nodes, 
      g_weighted)

    # Clustering with the K-Means algorithm
    labels = self.kmeans.fit_predict(adjacency_matrix)

    # Set labels to the nodes and group the nodes from the text according
    # to their labels.
    clusters = defaultdict(set)
    for node, label in zip(g_weighted.nodes(), labels):
      g_weighted.node[node][self.DATA_CLUSTER] = label
      if node in txt_nodes:
        clusters[label].add(node)

    # Add weights to the g_weighted's edges
    self._set_weights(g_weighted, txt_nodes)

    # Find the shortest paths between the nodes from the text 
    # within the groups
    out_graph = networkx.DiGraph()
    self._add_edges_within_cluster(
      digraph, 
      out_graph, 
      g_weighted, 
      clusters, 
      text_nodes)
    # and between the groups
    self._add_edges_between_clusters(
      digraph, 
      out_graph, 
      g_weighted, 
      clusters, 
      text_nodes)
    
    # Not all text nodes must be in the input digraph. Add those that are not.
    for txt_node in txt_nodes:
      if not txt_node in out_graph:
        out_graph.add_node(txt_node)

    return out_graph

  def _create_sparse_adjacency_matrix(self, digraph, txt_nodes, weighted_graph):
    """!
    Create the sparse adjacency matrix of the directed graph.

    The adjacency matrix is constructed as follows:
    1) if nodes are connected then add to the cell [from_node, to_node] one unit,
    2) if the from_node is in the nodes from the text, add one unit again,
    3) if the to_node is in the nodes from the text, add one unit again to that 
      cell.
    One unit equals 1./3.
    At the same time connections to the weighted graph are added.

    @param digraph the directed graph to find the adjacency matrix for.
    @type digraph: networkx.DiGraph()
    @param txt_nodes the set of nodes' identifiers from a text.
    @type txt_nodes: set
    @param weighted_graph a temporary graph that keeps weights of the 
      connections based on the adjacency matrix of the original directional 
      graph.
    @type weighted_graph: networkx.DiGraph

    @return sparse adjacency matrix for the input graph.
    @rtype: scipy.sparse.lil_matrix
    """
    # Make a simple mapping between the node's identifier and it's position
    # on the list of the nodes in the graph. It is used for creating 
    # the sparse adjacency matrix.
    mapping = dict((n, v) for v, n in enumerate(digraph.nodes()))
    # Set adjacency weights to edges based on the occurence of the nodes 
    # in the text.
    adjacency_matrix = scipy.sparse.lil_matrix(
      (digraph.number_of_nodes(), digraph.number_of_nodes()))
    # Use 1./3 in order to avoid the normalization of the matrix
    # TODO: Is it still necessary?
    for node_from, node_to in digraph.edges_iter():
      adjacency = 1./3
      if node_from in txt_nodes:
        adjacency += 1./3
      if node_to in txt_nodes:
        adjacency += 1./3
      weighted_graph.add_edge(
        node_from, 
        node_to)
      adjacency_matrix[mapping[node_from], mapping[node_to]] = adjacency
    return adjacency_matrix

  def _set_weights(self, graph, txt_nodes):
    """!
    Set weights to the edges of the input graph.

    Set weights to the edges in order to calculate the shortest paths.
    The weight is 4, but some situations decrease it:
    1) -2 if connected nodes are in the same cluster.
    2) -1 if the first node belongs to one cluster and the second node to 
        another
    3) -1 for each node that belongs to the original text.

    @param graph a graph whose nodes have 'cluster' assigned.
    @type graph: networkx.Graph()
    @param txt_nodes the set of nodes' identifiers from a text.
    @type txt_nodes: set
    """
    # Set weights to the edges in order to calculate the shortest paths.
    # The weight is 4, but some situations decrease it:
    # -1 for each node in the same group
    # -1 for each node that belongs to the original text 
    for node_from, node_to in graph.edges_iter():
      # K-means put every node to a group, there are no outliers.
      # Thus, start from the weight equals 3
      weight = 3
      node_from_data = graph.node[node_from]
      node_to_data = graph.node[node_to]
      group = node_from_data[self.DATA_CLUSTER]
      # Nodes are in the same group
      if node_to_data[self.DATA_CLUSTER] == group:
        # If the initial weight is 4, then here will be -2
        weight -= 1

      # Is the first node in the original text?
      if node_from in txt_nodes:
        weight -= 1
      # How about the second node?
      if node_to in txt_nodes:
        weight -= 1
      # Assign the weight to the edge
      graph[node_from][node_to][constants.EDGE_WEIGHT] = weight

  def _add_edges_within_cluster(self, digraph, out_graph, weighted_graph, 
      clusters, text_nodes):
    """!
    Add edges within the nodes of the same cluster that belong to the text.

    For each pair of nodes from the original text the shortest path between
    them is found and added to the out graph.

    @param digraph the original, directional graph.
    @type digraph: networkx.DiGraph
    @param out_graph the reduced graph that is being created
    @type out_graph: networkx.DiGraph
    @param weighted_graph a temporary graph that keeps weights of the 
      connections based on the adjacency matrix of the original directional 
      graph.
    @type weighted_graph: networkx.DiGraph
    @param clusters contains identifiers of nodes from the original text that
      belong to each cluster.
    @type clusters: dict( int: set(str) )
    @param text_nodes nodes from a text.
    @type text_nodes: grafon.structures.textnodes.TextNodes
    """
    for nodes in clusters.itervalues():
      for source, target in text_nodes.get_combinations(nodes):
        shortest_path = networkx.shortest_path(
          weighted_graph, 
          source, 
          target, 
          weight = constants.EDGE_WEIGHT)
        # Retrieve the direction of the edges in the shortest path
        out_graph.add_edges_from(
          self._get_directed_edges_of_path(shortest_path, digraph))

  def _add_edges_between_clusters(self, digraph, out_graph, weighted_graph, 
      clusters, text_nodes):
    """!
    Add edges between a pair of the most representative nodes of each cluster.

    The method finds nodes within a cluster that have already been connected and
    have the largest number of edges. They are called representative nodes
    (also center nodes).
    For each pair of center nodes the shortest path between them is found and 
    added to the out graph.

    @param digraph the original, directional graph.
    @type digraph: networkx.DiGraph
    @param out_graph the reduced graph that is being created
    @type out_graph: networkx.DiGraph
    @param weighted_graph a temporary graph that keeps weights of the 
      connections based on the adjacency matrix of the original directional 
      graph.
    @type weighted_graph: networkx.DiGraph
    @param clusters contains identifiers of nodes from the original text that
      belong to each cluster.
    @type clusters: dict( int: set(str) )
    @param text_nodes nodes from a text.
    @type text_nodes: grafon.structures.textnodes.TextNodes
    """
    if len(clusters) > 1:
      # The connections within each group were created. However, there are still
      # no the connections between the groups. Thus, they need to be calculated.
      center_nodes = dict(
        max(out_graph.degree(cluster_nodes).items(), key = lambda (n, d): d) \
        if len(cluster_nodes)> 1 else (cluster_nodes.pop(), 0)
          for cluster_nodes in clusters.itervalues()
        )
      center_nodes = center_nodes.keys()
      for source, target in text_nodes.get_combinations(center_nodes):
        shortest_path = networkx.shortest_path(
          weighted_graph, 
          source, 
          target, 
          weight = constants.EDGE_WEIGHT)
        out_graph.add_edges_from(
          self._get_directed_edges_of_path(shortest_path, digraph))