import sys
import argparse
import ConfigParser
import xml.etree.ElementTree as ET
import corpus2
import itertools
import matplotlib.pyplot as plt

from collections import defaultdict
import operator

import networkx as nx
from networkx import *

try:
    import argcomplete
except ImportError:
    argcomplete = None

__author__ = 'marlena'

class Graph:
  def __init__(self, tagset):
    self.graph_dict = None
    self.tagset = corpus2.get_named_tagset(tagset)
    self.__chan_name_list = ['np_sem_roles', 'deepened_chunker', 'chunk_agp']
    self.__chan_dict = None
    #Dokument nalezy przetrzymywac w konstruktorze, poniewaz
    #po zakonczeniu funkcji __add_nodes zostaje zniszczony, a obiekty Tokenu
    #sa usuwane z pamieci i nadpisywane przez cos innego.
    self.__document = None
    

  def __add_nodes(self, text_file):
    self.graph_dict = {}
    self.__chan_dict = {}
    interp_mask = self.tagset.parse_symbol('interp')
    try:
      cclreader = corpus2.CclRelReader(self.tagset, text_file, text_file)
      self.__document = cclreader.read()
      for paragraph in self.__document.paragraphs(): 
        for sentence in paragraph.sentences():
          #wydobycie chan z tokenow w zdaniu
          #wydobywany jest numer tokenu w ktorym znajduje sie
          #wybrany kanal anotacji oraz jego wartosc w tym kanale
          #Budowa slownika:
          #self.__chan_dict[nazwa kanalu][(id. zdania, id. dla tokenu w kanale)] = id. tokenu w zdaniu
          asent = corpus2.AnnotatedSentence.wrap_sentence(sentence)
          for chan_name in self.__chan_name_list:
            if asent.has_channel(chan_name):
              chan = asent.get_channel(chan_name)
              ann_vec = chan.make_annotation_vector()
              for ann in ann_vec:
                key = (sentence.id(), ann.seg_number)
                if self.__chan_dict.has_key(chan_name):
                  if not self.__chan_dict[chan_name].has_key(key):
                    self.__chan_dict[chan_name][key] = ann.head_index
                  else:
                    print 'W tekscie wejsciowym dla kanalu', chan, \
                        'w zdaniu', sentence.id(), \
                        'powtorzyl sie numer segmentu', ann.seg_number, '!'
                else:
                  new_chan_dict = {}
                  new_chan_dict[key] = ann.head_index
                  self.__chan_dict[chan_name] = new_chan_dict
          #dodanie do listy grafow, grafu dla kazdego zdania
          #w podanym na wejsciu tekscie
          self.graph_dict[sentence.id()] = nx.MultiDiGraph()
          for index, token in enumerate(sentence.tokens()):
            token_interp_mask = token.get_preferred_lexeme(self.tagset). \
                      tag().get_masked(interp_mask)
            if token_interp_mask != interp_mask:
              self.graph_dict[sentence.id()].add_node(index, 
                                  token=token, 
                                  sentence=sentence.id())
    except Exception, e:
      print >> sys.stderr, 'Error: ', e

  def __add_named_and_unnamed_edges(self, rel_file):
    tree = ET.parse(rel_file)
    root = tree.getroot()

    #Utworzenie listy relacji nazwanych oraz nienazwanych.
    #chan -> [[sentence_id, from_node_id, to_node_id, rel_name, extra_node_id], [], ..., []]
    rel_dict = defaultdict(list)
    for index, rel in enumerate(root.findall('rel')):
      rel_name = rel.get('name')

      if rel_name == '3args_rel_2nd_part':
        continue

      sentence_id = rel.find('from').get('sent')
      chan = rel.find('from').get('chan')
      extra_node = None
      extra_node_id = None

      if rel_name == '3args_rel_1st_part':
        from_node = int(root.findall('rel')[index].find('from').text)
        to_node = int(root.findall('rel')[index+1].find('to').text)

        extra_node = int(root.findall('rel')[index].find('to').text)
        extra_node_id = self.__chan_dict[chan][(sentence_id, extra_node)]
      else:
        from_node = int(rel.find('from').text)
        to_node = int(rel.find('to').text)

      from_node_id = self.__chan_dict[chan][(sentence_id, from_node)]
      to_node_id = self.__chan_dict[chan][(sentence_id, to_node)]

      rel_list = [sentence_id, \
            from_node_id, \
            to_node_id, \
            rel_name if chan == 'np_sem_roles' else None, \
            extra_node_id if extra_node_id else None]
      rel_dict[chan].append(rel_list)

    #Dodane krawedzi do grafu jako relacji nazwanych i nienazwanych.
    #Dodanie krawedzi nazwanych.
    for rel in rel_dict['np_sem_roles']:
      sentence_id = rel[0]
      from_node_id = rel[1]
      to_node_id = rel[2]
      rel_name = rel[3]
      if not to_node_id in self.graph_dict[sentence_id].neighbors(from_node_id):
        self.graph_dict[sentence_id].add_edge(from_node_id, 
                            to_node_id, 
                            rel='named', 
                            rel_name=rel_name)
      else:
        all_edges = self.graph_dict[sentence_id][from_node_id][to_node_id]
        named_rel_exist = False
        for index in all_edges.keys():
          if all_edges[index]['rel_name'] == rel_name:
            named_rel_exist = True
            break
        if not named_rel_exist:
          self.graph_dict[sentence_id].add_edge(from_node_id, 
                              to_node_id, 
                              rel='named', 
                              rel_name=rel_name)
    
    #Dodanie krawedzi nienzwanych.
    for rel in rel_dict['deepened_chunker']:
      sentence_id = rel[0]
      from_node_id = rel[1]
      to_node_id = rel[2]
      extra_node_id = rel[4]
      if not to_node_id in self.graph_dict[sentence_id].neighbors(from_node_id):
        if extra_node_id:
          self.graph_dict[sentence_id].add_edge(from_node_id, 
                              to_node_id, 
                              rel='unnamed', 
                              extra_node=extra_node_id)
        else:
          self.graph_dict[sentence_id].add_edge(from_node_id, 
                              to_node_id, 
                              rel='unnamed')
      else:
        if extra_node_id:
          all_edges = self.graph_dict[sentence_id][from_node_id][to_node_id]
          rel_with_extra_node_exist = False
          for index in all_edges.keys():
            if not all_edges[index].has_key('extra_node'):
              all_edges[index]['extra_node'] = extra_node_id
              rel_with_extra_node_exist = True
              break
            else:
              if all_edges[index]['extra_node'] == extra_node_id:
                rel_with_extra_node_exist = True
                break
          if not rel_with_extra_node_exist:
            self.graph_dict[sentence_id].add_edge(from_node_id, 
                                to_node_id, 
                                rel='unnamed', 
                                extra_node=extra_node_id)

  def __add_between_words_edges(self):
    for sentence_id, graph in self.graph_dict.iteritems():
      for index in range(len(graph.nodes())-1):
        from_node_id = graph.nodes()[index]
        to_node_id = graph.nodes()[index+1]
        self.graph_dict[sentence_id].add_edge(from_node_id, 
                            to_node_id, 
                            rel='between_words')

  def __add_phrase_head_edges(self):
    sorted_head_rels = sorted(self.__chan_dict['chunk_agp'].items(), 
                  key=operator.itemgetter(0))
    for index in range(len(sorted_head_rels)-1):
      from_node_id = sorted_head_rels[index][1]
      sentence_from_node = sorted_head_rels[index][0][0]
      to_node_id = sorted_head_rels[index+1][1]
      sentence_to_node = sorted_head_rels[index+1][0][0]

      if sentence_from_node == sentence_to_node:
        self.graph_dict[sentence_from_node].add_edge(from_node_id, 
                               to_node_id, 
                               rel='heads')

  def make_graph(self, text_file, rel_file):
    self.__add_nodes(text_file)
    self.__add_named_and_unnamed_edges(rel_file)
    self.__add_between_words_edges()
    self.__add_phrase_head_edges()


def aparser():
    arg_parser = argparse.ArgumentParser(description='Build graph from text.')
    arg_parser.add_argument('-T', '--text-file', dest = 'text_file', required = True, 
                help = 'path to tagged text file')
    arg_parser.add_argument('-R', '--rel-file', dest = 'rel_file', required = True, 
                help = 'path to relation file corresponding to text file')
    arg_parser.add_argument('-t', '--tagset', dest = 'tagset', default = 'nkjp',
                required = False, help = 'tagset (np. nkjp)')
    arg_parser.add_argument('-v', '--v-folder', dest = 'v_folder', required = True,
                help = 'path to graphs visualisation folder')
    if argcomplete:
        argcomplete.autocomplete(arg_parser)
    return arg_parser.parse_args()

def main():
  args = aparser()

  G = Graph(args.tagset)
  G.make_graph(args.text_file, args.rel_file)


  ######################### DODATEK: RYSOWANIE GRAFOW #########################
  #Wypisanie wszystkich relacji z grafow.
  for sentence_id, graph in G.graph_dict.iteritems():
    print '\nSentence ID:', sentence_id
    for n, nbrs in graph.adjacency_iter():
      for nbr, eattr in nbrs.items():
        print 'From node:', n, 'To node:', nbr, '\nRelations:',
        for index, rel in eattr.iteritems():
          print rel,
        print '\n'
  
  #Zapisanie wizualizacji grafow do plikow.
  for sentence_id, graph in G.graph_dict.iteritems():
    #Ponizsza petla zapisuje do slownika lematy w postaci:
    #id_wierzcholka -> lemat
    node_with_lemma_dict = {}
    for index in graph.nodes():
      lemma = graph.node[index]['token'].get_preferred_lexeme(G.tagset). \
                         lemma_utf8().lower()
      node_with_lemma_dict[index] = lemma

    A = nx.to_agraph(graph)
    
    #Dodanie lematow do wierzcholkow.
    for node in A.nodes():
      node.attr['label'] = node_with_lemma_dict[int(node)].decode('utf-8') \
                 + '(' + node + ')'

    #Dodanie kolorow do krawedzi w zaleznosci od typu relacji oraz
    #dodanie nazw krawedzi.
    #relacje nazwane -> kolor czerwony
    #relacje nienazwane -> kolor niebieski
    #relacje pomiedzy slowami -> kolor zielony
    #relacje pomiedzy glowami fraz -> kolor rozowy
    for edge in A.edges():
      if edge.attr['rel'] == 'named':
        edge.attr['color'] = 'red'
      elif edge.attr['rel'] == 'unnamed':
        edge.attr['color'] = 'blue'
      elif edge.attr['rel'] == 'between_words':
        edge.attr['color'] = 'green'
      elif edge.attr['rel'] == 'heads':
        edge.attr['color'] = 'magenta'

      if edge.attr['extra_node'] and edge.attr['rel_name']:
        edge.attr['label'] = edge.attr['rel'] \
                   + '(' + str(edge.attr['extra_node']) + '):' \
                   + edge.attr['rel_name']
      elif edge.attr['extra_node'] and not edge.attr['rel_name']:
        edge.attr['label'] = edge.attr['rel'] + '(' + str(edge.attr['extra_node']) + ')'
      elif not edge.attr['extra_node'] and edge.attr['rel_name']:
        edge.attr['label'] = edge.attr['rel'] + ':' + edge.attr['rel_name']
      elif not edge.attr['extra_node'] and not edge.attr['rel_name']:
        edge.attr['label'] = edge.attr['rel']

    #A.layout(prog='dot') #wizualizacja pionowa
    A.layout(prog='dot', args="-Grankdir=LR") #wizualizacja pozioma
    graph_file_name = args.v_folder + 'graph_' + sentence_id + '.png'
    A.draw(graph_file_name)
    print 'Do pliku', graph_file_name, \
        'zapisano wizualizacje grafu dla zdania', sentence_id, '!'


if __name__ == '__main__':
    main()
