#
# <one line to give the library's name and an idea of what it does.>
# Copyright (C) 2016  <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 graphrebuilderinterface
import grafon.constants as constants
from grafon.ebuilders import netoneedgebuilder

class NEGraphBuilder(graphrebuilderinterface.GraphRebuilderInterface):
  """
  This class handle NE into built base graph.

  It connects all tokens belongs to the same NE into one node.
  """
  def __init__(self, str_name='NEGraphBuilder'):
    super(NEGraphBuilder, self).__init__(str_name)
    self.nint = netoneedgebuilder.NEToNEEdgeBuilder()
    self.ne_channels = set()
    self.ne_segments = {}

  def rebuild_base_graph(self, bgraph):
    """
    For each grapg stored into bgraph.graph_dict try to rebuild graph
    with connecting NE into one node.

    For rebuild grapgs calls method @see rebuild_graph
    """
    if not bgraph or not bgraph.document:
      return 
    for p in bgraph.document.paragraphs():
      for sentence in p.sentences():
        asent = corpus2.AnnotatedSentence.wrap_sentence(sentence)
        self.find_ne_channels(asent)

        g = bgraph.graph_dict.get(sentence.id())
        if not g:
          logging.error(
            'There is no graph in NEGraph builder '
            'for sentence {}'.format(sentence.id()))
          continue
        for token in sentence.tokens():
          ne_anumber, channame = self.ne_annot_number(token, sentence)
          if bgraph.skip_token(token):
            continue
          for k in bgraph.node_converter.key(token):
            if not k:
              continue
            if ne_anumber != -1:
              g.node[k][constants.NODE_NE_ANNOTS].append(
                (ne_anumber, channame))
            else:
              g.node[k][constants.NODE_IS_SINGLETOK] = True
        g = self.rebuild_graph(g)
        bgraph.graph_dict[sentence.id()] = g

  def find_ne_channels(self, asent):
    """
    Find all NE channels into given annotated sentence 
    and store its into member value

    @param asent Annotated sentence with given token
    @type asent: corpus2.AnnotatedSentence
    """
    self.ne_channels.clear()
    self.ne_segments = {}
    for ch in asent.all_channels():
      if self.nint.get_ne_set_name(ch):
        self.ne_channels.add(ch)
        self.ne_segments[ch] = list(asent.get_channel(ch).segments())

  def ne_annot_number(self, token, sent):
    """
    For give token return annotation number 
    with name of ne channel in NE channel.
    If token haven't ne annotation then returns (-1, '')

    @param token Token to check annotation number
    @type token: object of corpus2.Token
    @param sent Sentence with given token
    @type sent: corpus2.Sentence

    @return Pair of {Number of annotation or -1 if no ne annotation 
            and name of channel (or empty string)}
    @rtype: (int, str)
    """
    idx = list(sent.tokens()).index(token)
    if idx == -1:
      return (-1, '')
    for channame, segments in self.ne_segments.iteritems():
      if segments[idx] != 0:
        return (segments[idx], channame)
    return (-1, '')

  def rebuild_graph(self, g):
    """
    This function reguilds given graph based on the NE 
    informations into nodes. Returns rebuilded graph.

    @param g Built graph
    @type g: networkx.MultiDiGraph
    @return rebuilded graph
    @rtype: networkx.MultiDiGraph
    """
    # Slownik id anotacji na liste nodow do polaczenia
    new_nodes = {}
    for n, attr_dict in g.nodes(data=True):
      if len(g.node[n][constants.NODE_NE_ANNOTS]):
        if not g.node[n].has_key(constants.NODE_NE_ANNOTS):
          continue
        for (anum, channame) in g.node[n][constants.NODE_NE_ANNOTS]:
          if not new_nodes.has_key((anum, channame)):
            new_nodes[(anum, channame)] = []
          new_nodes[(anum, channame)].append(n)

    # W tej petli lacze ze soba NE i dodaje do grafu
    ne_nodes = {}
    for (anum, channame), nodes in new_nodes.iteritems():
      k = ""
      for n in nodes:
        k += "{} ".format(n.encode('utf8'))
      k += " ({}:{})".format(channame, anum)
      k = k.decode('utf8')
      if g.has_node(k):
        g.node[k][constants.NODE_COUNT] += 1
      else:
        g.add_node(k, attr_dict = {
          constants.NODE_LABEL        : k,
          constants.NODE_TOKENS       : [], 
          constants.NODE_COUNT        : 1,
          constants.NODE_SOURCES      : [constants.NODE_FROM_TEXT]})
      if not ne_nodes.has_key(k):
        ne_nodes[k] = []
      ne_nodes[k].extend(nodes)
      for n in nodes:
        if self.nint.exists_edge(g, k, n, channame):
          g[k][n][channame][constants.EDGE_COUNT] += 1
        else:
          g.add_edge(
            k, n, key=channame, 
            attr_dict = {
              constants.EDGE_LABEL: channame, 
              constants.EDGE_COUNT: 1})

    # Tutaj przepinam wejscia i wyjscia NE do elementow NE
    """
    ntorem = []
    for k, nodes in ne_nodes.iteritems():
      for n in nodes:
        " ""
        for nfrom, nto, d in g.in_edges(n, data=True):
          print "IN:", nfrom, nto
          g.add_edge(nfrom, nto, attr_dict = d)

        for nfrom, nto, d in g.out_edges(n, data=True):
          print "OUT:", nfrom, nto
          g.add_edge(nfrom, nto, attr_dict = d)

        " ""
        for s in g.successors(n):
          print 'new e ', k, s
          g.add_edge(s, k)
        for p in g.predecessors(n):
          g.add_edge(n, k)
        if not g.node[n][constants.NODE_IS_SINGLETOK]:
          ntorem.append(n)
    #g.remove_nodes_from(ntorem)
    """
    return g