#
# <one line to give the library's name and an idea of what it does.>
# Copyright (C) 2014    <copyright holder> <email>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.    See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.    If not, see <http://www.gnu.org/licenses/>.
# 
#

import logging
import corpus2
import edgebuiderinterface
import needgebuilderinterface
import grafon.constants as constants

class NEToNEEdgeBuilder(needgebuilderinterface.NEEdgeBuilderInterface):
  """!
  The class creates edges between named entities' (NE) nodes from the text.
  It takes into account the order of NE and connections are made according to 
  it. However, the type of NE (more precisely - a specifil channel's name) does
  not matter.
  """
  def __init__(self, str_name = 'ne2ne'):
    """!
    Create the class by calling the constructor of it's parent.

    @param str_name the name assigned to edges that are created by this class.
    @type str_name: str
    """
    super(NEToNEEdgeBuilder, self).__init__(str_name)

  def add_edges(self, base_graph):
    """!
    Create an edge from one NE to another NE, regarding the sequence.

    For each NE that occures in the text make the connections between 
    it's predecessor, this node and it's successor. The order is defined by
    the words' order in the document.

    @param base_graph the graph to be enriched by adding the edges.
    @type base_graph: grafon.graphs.BaseGraph
    """
    interp_mask = base_graph.tagset.parse_symbol('interp')
    rel_id = 1
    for paragraph in base_graph.document.paragraphs():
      for sentence in paragraph.sentences():
        g = base_graph.graph_dict.get(sentence.id(), None)
        if not g:
          logging.warning('There is no graphs for sentence {}'.format(
            sentence.id()))
          continue
        asent = corpus2.AnnotatedSentence.wrap_sentence(sentence)
        # Search for NE channels
        ne_sets = set()
        ne_channels_names = list()
        for channel in asent.all_channels():
          ne_set = self.get_ne_set_name(channel)
          if ne_set:
            ne_sets.add(ne_set)
            ne_channels_names.append(channel)
        # If there was more than 1 NE set then document may be corrupted
        if len(ne_sets) > 1:
          logging.warning(
            ("WARNING: found {:} NE sets: {:}, check whether your Liner "
            "configuration is as you expected or the ccl file is correct."
            ).format(len(ne_sets), ", ".join(ne_sets)))

        # Create an edge between tokens
        ne_channels = [asent.get_channel(ne_channel_name).segments() \
            for ne_channel_name in ne_channels_names]
        # iterate over the same segment in each channel
        tokens = []
        for i, channels_segments in enumerate(zip(*ne_channels)):
          # if at least one NE channel is active, add the token
          if sum(channels_segments) > 0:
            token = asent.tokens()[i]
            if not base_graph.skip_token(token):
              tokens.append(token)

        # Add edges between tokens
        for i, token in enumerate(tokens[:-1]):
          from_nodes = base_graph.node_converter.key(tokens[i])
          to_nodes = base_graph.node_converter.key(tokens[i + 1])
          for fn in from_nodes:
            if not g.has_node(fn):
                continue
            for tn in to_nodes:
              if not g.has_node(tn):
                continue
              if fn and tn:
                if self.exists_edge(g, fn, tn, self._str_name):
                    g[fn][tn][self._str_name][constants.EDGE_COUNT] += 1
                else:
                  g.add_edge(
                    fn, tn, key=self._str_name, attr_dict = {
                      constants.EDGE_LABEL: self._str_name, 
                      constants.EDGE_RELID: rel_id,
                      constants.EDGE_COUNT: 1})
