# -*- coding: utf-8 -*-

# <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/>.
#

"""
This module is a handler for entailment data and other etailment things.
"""

import os, csv
import networkx
import logging
try:
  import argcomplete
except ImportError:
  argcomplete = None

from grafon import grafon_base

def make_parser(desc = ''):
  """
  Makes base parser for entailment
  """
  aparser = grafon_base.make_parser(desc)
  aparser.add_argument('-f', '--cclfile', dest='ccl',
                       required=True, type=str,
                       help=(
                         'Ccl file. For each sentences into file, '
                         'the graph will be build.'))
  aparser.add_argument('-r', '--relcclfile', dest='relccl',
                       required=True, type=str,
                       help=(
                         'Relccl file. Depends on the edge '
                         'builders relccl file may be used for build edges.'))
  aparser.add_argument('-o', '--out-res-file', dest = 'out_res_file',
                       required=False, help=(
                         'If options is given then '
                         'then the results will be stroed into given filename. '
                         'If not gvien then the result will be stored into '
                         'standard output.'))
  if argcomplete:
    argcomplete.autocomplete(aparser)
  return aparser

def read_pack(pack_dir):
  """
  Reads single pack from given directory
  
  @return Dictionary of dictionary of files into pack with sentences into its.
          The keys of the global dictionary is pack number, value is dictionary
          of files and sentences into it.
  @rtype: dict: {str} => {{str}=>[int, ..]}
  """
  docdict = {}
  for pfile in os.listdir(pack_dir):
    pfilepath = os.path.join(pack_dir, pfile)
    # print >> sys.stderr, "Checking file: {}".format(pfilepath)
    with open(pfilepath) as pf:
      docdict[pfile] = []
      lastkey = None
      for line in pf:
        line = line.strip()
        if line.find('Dokument: ') != -1:
          lastkey = line.split('Dokument: ')[1]
          docdict[pfile].append( (lastkey, [], []) )
          continue
        if not lastkey:
          print >> sys.stderr, "ERROR! {}".format(pfilepath)
          return -1
        if line.startswith('==='):
          continue
        spline = line.split()
        if len(spline) < 2: 
          continue
        docdict[pfile][-1][1].append(int(line.split()[0]))
        docdict[pfile][-1][2].append(' '.join(w for w in line.split()[1:]))
  return docdict

def read_gen_graphs(graph_dir):
  """
  Z zadanego katalogu wczytuje grafy opisujace relacje entailmentowe.
  
  Zwraca słownik, w którym klczuem jest nazwa relacji, a wartością 
  jest słownik, w którym kluczem jest okreslenie pochodzenia grafu (to, from)
  a wartością wczytany graf networkxowy
  
  dict: {
    nazwa relacji : dict{ 'from' : graf, 'to' : graf},
    ...
  }
  """
  reldict = {}
  for gfile in os.listdir(graph_dir):
    if not gfile.endswith('.pickle'):
      continue
    relname, fromto = gfile.split('#')
    relname = relname[:-1]
    if not reldict.has_key(relname):
      reldict[relname]={}
    if fromto.find('_to__') != -1:
      reldict[relname]['to'] = networkx.read_gpickle(
        os.path.join(graph_dir, gfile))
    else:
      reldict[relname]['from'] = networkx.read_gpickle(
        os.path.join(graph_dir, gfile))
  return reldict

def make_gpack(pack_dict, graph_dir):
  """
  Z zadanej paczki, ktorą otrzymywali lingwiści tworzy paczki grafów.

  Zwraca słownik, w ktorym kluczem jest numer paczki z wartością jest słownik
  w którym kluczem jest nazwa pliku, a wartością jest słownik, w którym kluczem
  jest numer zdania, a wartością jest para: <graf zbudowany ze zdania, ścieżka
  do pliku w którym graf jest zapisany>.
  
  packnumber : {
    filename : {
      sentencenumber : <graph, graphpath>,
      ...
    } , ...
  }, ...
  """
  gdct = {}
  # build graphs paths from pack
  for packnum, docdict in pack_dict.iteritems():
    if not gdct.has_key(packnum):
      gdct[packnum] = {}
    for docname, sentences, _ in docdict:
      if not gdct[packnum].has_key(docname):
        gdct[packnum][docname] = {}
      for s in sentences:
        gpath = os.path.join(graph_dir, '{}#s{}.pickle'.format(docname, s))
        logging.info('Reading pickled graph from file {}'.format(gpath))
        graph = networkx.read_gpickle(gpath)
        gdct[packnum][docname][s] = (graph, gpath)

  return gdct

def generate_pairs(lingpack, lingdec=None):
  """
  Z paczek generuje mozliwe pary zdan, pomiedzy ktorymi moze zajsc jakas 
  relacja.

  Jesli podany zostal lingdec generowanie par uwzglednia wcześniejsze
  oznaczenia lingwistów.
  """
  pospairs = {}
  first_docs = []
  pair_docs = []
  if lingdec:
    first_docs = _extract_first_documents(lingdec)
    pair_docs = _extract_pairs_of_documents(lingdec)

  for packnum, documents in lingpack.iteritems():
    if len(documents) != 3:
      logging.warning('Package {} size is equal {}. Should be 3!'.format(
        packnum, len(documents)))
      continue

    pospairs[packnum] = []
    if lingdec:
      if not documents[0][0] in first_docs:
        pospairs[packnum].extend(generate_sentence_pairs(
          docname1=documents[0][0], sentences1=documents[0][1], 
          docname2=documents[0][0], sentences2=documents[0][1],
          lingdec=lingdec))
      """
      if not (documents[0][0], documents[1][0]) in pair_docs:
        pospairs[packnum].extend(generate_sentence_pairs(
        docname1=documents[0][0], sentences1=documents[0][1], 
        docname2=documents[1][0], sentences2=documents[1][1],
        lingdec=lingdec))
      if not (documents[0][0], documents[2][0]) in pair_docs:
        pospairs[packnum].extend(generate_sentence_pairs(
          docname1=documents[0][0], sentences1=documents[0][1], 
          docname2=documents[2][0], sentences2=documents[2][1],
          lingdec=lingdec))
       """
    else:
      pospairs[packnum].extend(generate_sentence_pairs(
          docname1=documents[0][0], sentences1=documents[0][1], 
          docname2=documents[0][0], sentences2=documents[0][1],
          lingdec=lingdec))
      """
      pospairs[packnum].extend(generate_sentence_pairs(
        docname1=documents[0][0], sentences1=documents[0][1], 
        docname2=documents[1][0], sentences2=documents[1][1],
        lingdec=lingdec))
      pospairs[packnum].extend(generate_sentence_pairs(
          docname1=documents[0][0], sentences1=documents[0][1], 
          docname2=documents[2][0], sentences2=documents[2][1],
          lingdec=lingdec))
      """
  return pospairs

def generate_sentence_pairs(
  docname1, sentences1, docname2, sentences2, lingdec):
  """
  Dla zadanych dwoch dokumentow oraz numerow zdan generuje mozliwe
  pary zdan pomiedzy ktorymi beda wykrywane relacje semantyczne.
  
  Zwraca liste par par.
  [
    (
      (docname1, sentnum1), 
      (docname2, sentnum2)
    ), ...
  ]
  """
  pospairs = []
  
  for s1 in sentences1:
    s1 = int(s1)
    for s2 in sentences2:
      s2 = int(s2)
      # skip the same pair of docname, sentence
      if (docname1, s1) == (docname2, s2):
        continue
      if lingdec:
        if ((docname1, s1), (docname2, s2)) in lingdec:
          print 'Juz oceniano {}:{} {}:{}'.format(docname1, s1, docname2, s2)
          continue
      # skip sentences in reverse
      if ((docname2, s2), (docname1, s1)) in pospairs:
        continue
      pospairs.append(((docname1, s1), (docname2, s2)))
  return pospairs

def read_ling_decision(decpath):
  """
  Wczytuje decyzje lingwistow do wczesniej ocenianych relacji.
  
  Wymagany format pliku csv:
   w pierwszej i drugiej kolumnie znajdują się zdania pomiędzy którymi
   oceniana była relacja. Pojedyncze zdanie to właściwie nazwapliku z numerem
   zdania z paczki podzielone dwukropkiem. Przykładowo:
   
   nazwa_pliku:numer zdania\tnazwa_pliku:numer zdania
   
  Wpisy rozdzielone są tabulatorem. W każdej nowej linii powinna 
  być zapisana kolejna decyzja.
  
  
  Zwraca listę par par:
  [
    (
      (nazwa dokumentu, numer zdania),
      (nazwa dokumentu, numer zdania)
    ), ..
  ]
  """
  decling = []
  with open(decpath) as infile:
    creader = csv.reader(infile, delimiter='\t')
    for dec in creader:
      d1, s1 = dec[0].strip().replace('ner-iob-', '').split(':')
      d2, s2 = dec[1].strip().replace('ner-iob-', '').split(':')
      s1 = int(s1[1:])
      s2 = int(s2[1:])
      decling.append(((d1, s1), (d2, s2)))
  return decling

def _extract_first_documents(lingdec):
  """
  Z zadanej listy par par ekstrahuje te dokumenty, ktore wystapily
  jako glowny w przekazanej liscie (jako pierwszy)
  
  Zwraca liste nazw takich dokumentow
  """
  return [d1 for (d1, _), (_, _) in lingdec]

def _extract_pairs_of_documents(lingdec):
  """
  Z zadanej listy par par ekstrahuje listy par dokumentow, 
  ktore wystapily w przekazanej liscie
  
  Zwraca liste nazw takich dokumentow
  """
  return [(d1, d2) for (d1, _), (d2, _) in lingdec]