u"""
The module contains th MaltRunner class which has the ability to add information
from the MaltParser to a ccl and rel ccl file. A new annotation channel, named
'malt' is created.

You may use the MaltRunner class explicitly or run the module as if it were 
a script ($python malt.py -h).

You can download the MaltParser from the website:
http://www.maltparser.org/download.html
"""

import traceback, codecs, subprocess, corpus2, os
from tempfile import NamedTemporaryFile

S_MALTPARSER = "maltparser"
O_MALT_PARSER_PATH = "malt_parser_path"
O_DATA_MODEL = "data_model"
O_DATA_MODEL_LOCALIZATION = "data_model_localization"

ENCODING = "UTF-8"
MALT_CHANNEL = 'malt'
MALT_SET = 'malt_relations'

class MaltRunner(object):
  """Run MaltParser in order to get extra information."""

  def __init__(self, config):
    u"""Initialize MaltRunner class.

    The config must have section 'maltparser' which contains: 
    'malt_parser_path', 'data_model', 'data_model_localization'.

    Args:
      config (ConfigParser.RawConfigParser): it contains MaltParser 
        configuration.
    """
    self.malt_parser_path = config.get(
      S_MALTPARSER, 
      O_MALT_PARSER_PATH)
    self.data_model = config.get(
      S_MALTPARSER, 
      O_DATA_MODEL)
    self.data_model_localization = config.get(
      S_MALTPARSER, 
      O_DATA_MODEL_LOCALIZATION)
    self.tagset = corpus2.get_named_tagset(config.get('saper','tagset')) \
        if config.has_section('saper') else None

  @staticmethod
  def convert2conll(input_file, output_file,
      input_format="ccl", converter="nkjp2skladnica"):
    u"""Run maca-convert in order to get the file in the conll format.

    Args:
      input_file (str): the path to the input file.
      output_file (file-like object): the output file object.
      input_format (str): the format of the input file, it must be one
        of the available maca-convert input formats.
      converter (str): the name of the converter used by maca-convert,
        it must be accepted by the maca-convert tool.
    """
    cmd = "maca-convert -i {:} -o conll -I {:} {:}".format(
      input_format, 
      input_file,
      converter)
    try:
      process = subprocess.Popen(cmd, shell = True, stdout = output_file) 
      process.wait()
    except Exception as exc:
      print "CMD:", cmd
      print "maca-convert exception,", exc
      print traceback.format_exc()

  @staticmethod
  def run_maltparser(input_file, output_file,
      malt_parser_path, data_model, data_model_localization):
    u"""Run MaltParser jar file in order to extend conll file.

    Args:
      input_file (str): the path to the input file.
      output_file (str): the path to the output file.
      malt_parser_path (str): the path to the MaltParser jar file.
      data_model (str): the name of the mco data model for the MaltParser
        without the '.mco' extension (i.e. my_data_model, 
        not my_data_model.mco).
      data_model_localization (str): the path to the directory where the
        data model is.
    """
    cmd = ('java -Xmx1024m -jar {:} -c {:} -v warn '
      '-i {:} -o {:} -m parse -w {:}').format(
      malt_parser_path, 
      data_model, 
      input_file, 
      output_file,
      data_model_localization)
    try:
      process = subprocess.Popen(cmd, shell = True) 
      process.wait()
    except Exception as exc:
      print "CMD:", cmd
      print "MaltParser exception,", exc
      print traceback.format_exc()

  @staticmethod
  def read_conll_file(conll_file):
    u"""Return the list of sentences from the conll file.

    Args:
      conll_file (str): the path to the conll file.

    Returns:
      (tuple): the tuple has 2 elements. The former is the list of sentences 
        (list of list of tuple). Each sentence is a list of triplets, 
        which looks like:  (from_token, to_token, relation_name),
        the numeration of tokens starts from 1 (not 0!). 
        The latter is the list of positions of tokens for each sentence 
        (list of list of int).
    """
    conll_sentences = list()
    tokens_numbers_by_sentences = list()
    with codecs.open(conll_file, 'r', encoding = ENCODING) as conll_file:
      sentence = list()
      tokens_numbers = list()
      for line in conll_file:
        columns = line.split('\t')
        # End of the current sentence, add it to coll_sentences
        # The sentence MAY be EMPTY! That's correct!!!
        if len(columns)<8: 
          conll_sentences.append(sentence)
          sentence = list()
          tokens_numbers_by_sentences.append(tokens_numbers)
          tokens_numbers = list()
        # Extract a tuple such as: (from_token, to_token, rel_name)
        # numerate tokens 1 as in conll file
        # However, don't take roots!
        elif len(columns)>=8 and int(columns[6]) != 0:
          # 0 position   - the current token's position
          # 6th position   - the position of the token the relation
          #          leads to
          # 7th position  - the name of the relation
          sentence.append(
            (int(columns[0]), 
            int(columns[6]), 
            columns[7].encode(ENCODING)))
        # Remember all tokens numbers
        if len(columns)>=8:
          tokens_numbers.append(int(columns[0]))
    return conll_sentences, tokens_numbers_by_sentences


  @staticmethod
  def add_conll2document(conll_file, document):
    u"""Add relations from the conll file to the corpus2.Document.

    Args:
      conll_file (str): the path to the conll file.
      document (corpus2.Document): the document to extend.
    """
    # Read the conll file
    conll_sentences, tokens_numbers_by_sentences = \
      MaltRunner.read_conll_file(conll_file)
    # Add malt channel to the document
    ccl_sentences = [sentence for paragraph in document.paragraphs() \
        for sentence in paragraph.sentences()]
    for conll_sentence, tokens_numbers, ccl_sentence in \
        zip(conll_sentences, tokens_numbers_by_sentences, ccl_sentences):
      ann_sentence = corpus2.AnnotatedSentence.wrap_sentence(ccl_sentence)
      # Add the channel only if it doesn't exist
      if not MALT_CHANNEL in ann_sentence.all_channels():
        ann_sentence.create_channel(MALT_CHANNEL)
        channel = ann_sentence.get_channel(MALT_CHANNEL)
      # Add channel's segments (ccl)
        for token_number in tokens_numbers:
          # Token number is actually segment's value.
          # The numeration of tokens in a ccl sentence starts from 0,
          # but the value 0 of a segment means that the channel is 
          # 'off'. Thus start the numeration of segments normally from
          # 1, but keep in mind it's meaning.
          channel.set_segment_at(token_number-1, token_number)
      # Add relations now (rel-ccl)
      for from_token, to_token, relation_name in conll_sentence:
        # Add the relation between the tokens
        from_point = corpus2.DirectionPoint(
          ccl_sentence.id(), 
          MALT_CHANNEL,
          from_token
          )
        to_point = corpus2.DirectionPoint(
          ccl_sentence.id(),
          MALT_CHANNEL,
          to_token
          )
        ccl_relation = corpus2.Relation(
          relation_name, 
          from_point, 
          to_point
          )
        # all malt relations are into the same set
        ccl_relation.set_set(MALT_SET)
        # Clone in order to get RelationPtr
        document.add_relation(ccl_relation.clone_shared())

  @staticmethod
  def save_document(document, tagset, path, document_format='ccl'):
    u"""Save corpus2.Document in the given localization.

    Args:
      document (corpus2.Document): the object to save.
      tagset (corpus2.Tagset): used tagset.
      path (str): the path to the file.
      document_format (str, optional): the format of the written file.
        Default to 'ccl'.
    """
    writer = corpus2.TokenWriter.create_path_writer(
      document_format, 
      path,
      tagset
      )
    for paragraph in document.paragraphs():
      writer.write_chunk(paragraph)
    writer.finish()
    del writer

  @staticmethod
  def save_relations(document, path):
    u"""Save corpus2.Document in the given localization as the relccl file.

    Args:
      document (corpus2.Document): the object to save.
      path (str): the path to the file.
    """
    writer = corpus2.RelationWriter(path)
    writer.write(document.relations())
    del writer

  def process(self, input_ccl, input_relccl=None):
    u"""Extend the input file by information from the MaltParser.

    Args:
      input_ccl (str): the path of the ccl file to extend.
      input_relccl (str): the path of the rel ccl file to extend.
    """
    # Create a tmp conll file
    tmp_conll = NamedTemporaryFile(delete = False)
    tmp_conll_out = NamedTemporaryFile(delete = False)
    # Convert the input ccl file to the conll file
    self.convert2conll(input_ccl, tmp_conll)
    # Extend the conll file
    self.run_maltparser(
      tmp_conll.name, 
      tmp_conll_out.name,
      self.malt_parser_path, 
      self.data_model, 
      self.data_model_localization
      )
    # Add information from the created conll file to the original ccl file
    # and the relccl file if it was given
    # Check whether the rel ccl file exists, if not create it
    if not self.tagset:
      self.tagset = corpus2.get_named_tagset('nkjp')
    document = corpus2.CclRelReader(
      self.tagset, 
      input_ccl, 
      input_relccl if input_relccl and os.path.exists(input_relccl) \
          else input_ccl
      ).read()
    self.add_conll2document(tmp_conll_out.name, document)
    # Save the updates
    self.save_document(document, self.tagset, input_ccl)
    if input_relccl:
      self.save_relations(document, input_relccl)
    # Delete the tmp conll file
    os.remove(tmp_conll.name)
    os.remove(tmp_conll_out.name)

  def process_document(self, document):
    u"""Add MaltParser information to the corpus2.Document.

    Args:
      document (corpus2.Document): the document to extend.

    Returns:
      (corpus2.Document): extended input document.
    """
    # Create a tmp files. The MaltParser must have a file to process.
    tmp_conll = NamedTemporaryFile(delete = False)
    tmp_conll_out = NamedTemporaryFile(delete = False)
    tmp_ccl = NamedTemporaryFile(delete = False)
    # Save the document as a ccl file
    self.save_document(document, self.tagset, tmp_ccl.name)
    # Convert the input ccl file to the conll file
    self.convert2conll(tmp_ccl.name, tmp_conll)
    # Extend the conll file
    self.run_maltparser(
      tmp_conll.name, 
      tmp_conll_out.name,
      self.malt_parser_path, 
      self.data_model, 
      self.data_model_localization
      )
    self.add_conll2document(tmp_conll_out.name, document)
    # Delete the tmp files
    os.remove(tmp_conll.name)
    os.remove(tmp_conll_out.name)
    os.remove(tmp_ccl.name)
    return document


def parse_args():
  u"""Extend the input channel by information from the MaltParser."""
  import argparse
  parser = argparse.ArgumentParser()
  parser.add_argument('input_ccl',
    help="The input ccl file.")
  parser.add_argument('--input_relccl',
    help="The input relccl file.")
  parser.add_argument('config', 
    help="The config file including information about MaltParser.")
  return parser.parse_args()

if __name__ == '__main__': 
  from ConfigParser import RawConfigParser
  cfg_parser = RawConfigParser(allow_no_value = True)
  args = parse_args()
  cfg_parser.readfp(open(args.config))
  malt_runner = MaltRunner(cfg_parser)
  malt_runner.process(
    args.input_ccl,
    args.input_relccl if 'input_relccl' in args else None
    )    
