77#
# <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 grafon.constants as constants
from collections import defaultdict

class DoubleRelEdgeBuilder(edgebuiderinterface.EdgeBuiderInterface):
  def __init__(self, str_name, chan_name = None, set_name = None):
    # The names of relations into excluded_relations will not be added to graph
    self._excluded_relations = []
    super(DoubleRelEdgeBuilder, self).__init__(str_name, chan_name, set_name)
    
  def add_edges(self, base_graph):
    """!
    General method do adds the edge beteen two nodes based on the two-elements
    relation.
    
    @param base_graph: Object of grafon.graphs.BaseGraph
    """
    # dictionary of sentence id to relations in given channel
    drels, accepted_channels_names = \
        self._get_rels_dir(base_graph.document.relations())
    for paragraph in base_graph.document.paragraphs():
      for sentence in paragraph.sentences():
        # the sentence may do not have any annotations into given channel name
        if not drels.has_key(sentence.id()):
          continue
        g = base_graph.graph_dict.get(sentence.id(), None)
        if not g:
          logging.warning('There is no graphs for sentence {}'.format(
            sentence.id()))
          continue
        rels = drels[sentence.id()]
        asent = corpus2.AnnotatedSentence.wrap_sentence(sentence)
        gen_channels = (ch for ch in asent.all_channels() \
            if ch in accepted_channels_names)
        for channel_name in gen_channels:
          channel = asent.get_channel(channel_name)
          segments = [s for s in channel.segments()]
          # beacause in Saper exists a bug with multiple adding same relations
          # now, here we must to check these situations
          annotated = []
          for r in rels:
            if r.rel_name() in self._excluded_relations:
              continue
            f_pos = r.rel_from().annotation_number()
            t_pos = r.rel_to().annotation_number()
            if (f_pos, t_pos) in annotated:
              continue
            f_index, t_index = self._get_from_and_to_tokens_indices(
              f_pos, 
              t_pos, 
              segments)
            f_tok = sentence.tokens()[f_index]
            t_tok = sentence.tokens()[t_index]
            if base_graph.skip_token(f_tok) or base_graph.skip_token(t_tok):
                continue
            from_nodes = base_graph.node_converter.key(f_tok)
            to_nodes = base_graph.node_converter.key(t_tok)
            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:
                  lbl_txt = "%s:%s" % (self._str_name, r.rel_name())
                  if self.exists_edge(g, fn, tn, lbl_txt):
                    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: lbl_txt,
                        constants.EDGE_RELNAME: str(r.rel_name()),
                        constants.EDGE_COUNT: 1})
            annotated.append((f_pos, t_pos))

  def _get_rels_dir(self, relations):
    """!
    Return a dictionary of from_sentence_id - a list of relations pairs.
    """
    drels = defaultdict(list)
    for r in relations:
      pfrom = r.rel_from()
      if pfrom.channel_name() == self._chan_name:
        drels[pfrom.sentence_id()].append(r)
    return drels, set([self._chan_name])

  def _get_from_and_to_tokens_indices(self, f_pos, t_pos, segments):
    """!
    Returns indices of the tokens (from_token and to_token) in a sentence.
    """
    return segments.index(f_pos), segments.index(t_pos)
