#!/usr/bin/python
# -*- coding: utf-8 -*-
# PYTHON_ARGCOMPLETE_OK

"""
This module consist a class for processing ccl files, and corpus2.Document
objects.

Have fun ;).
"""
__author__ = 'blaz'


from carrot import options, resource
from rb import wccl_op
from defender.phrase_reader import PhraseReader
from defender.deepened_chunker import SyntaxRelationGenerator
import corpus2
import sys


class NpSemrelCclProcessor:
    def __init__(self, fst_type_pred, ger_dict, snd_type_pred, wccl_config,
                 npsemrel_config, tagset='nkjp', only_continous=False,
                 silent=True):
        self.op = options.Options(npsemrel_config)
        self.op.parse_cfg()
        self.tagset = corpus2.get_named_tagset(tagset)
        [self.wcc_op_cfg_file, self.wccl_ops] = \
                            load_wccl_operators(wccl_config, tagset)
        self.res = resource.Resource(self.tagset)
        self.res.load_fst_type_predicate_dictionary(fst_type_pred)
        self.res.load_snd_type_predicate_dictionary(snd_type_pred)
        self.res.load_ger_dict(ger_dict)
        self.phrase_reader = PhraseReader()
        self.only_continous = only_continous
        self.silent = silent

    def process_document(self, document, update = False):
        """
        This method processes corpus2.Document object, adds new annotations,
        relations. If the update is set to true, it won't create the copy
        of the document, but only extend the input document.
        
        @type document: corpus2.Document
        @rtype: corpus2.Document
        """
        np, adjp = self.__read_all_phrases_from_document(document)
        ann_doc = \
            self.__append_annotations_and_relations_to_document(
                document,
                (np, adjp),
                update = update)
        return ann_doc

    def __append_annotations_and_relations_to_document(self, document,
            phrase_channels, update = False):
        """
        This method appends new annotations and relations generated by npsemrel
        to corpus2.Document object
        
        @type document: corpus2.Document
        @type phrase_channels: tuple
        @rtype: corpus2.Document
        """
        if update:
            rel_dict = self.__mix_phrase_channels(phrase_channels)
            for paragraph in document.paragraphs():
                for sentence in paragraph.sentences():
                    if sentence.id() in rel_dict:
                        asent = \
                            corpus2.AnnotatedSentence_wrap_sentence(sentence)
                        if 'np_sem_roles' not in asent.all_channels():
                            asent.create_channel('np_sem_roles')
                        channel = asent.get_channel('np_sem_roles')
                        rels = self.__update_chann_and_gener_rels(
                                channel,
                                rel_dict[sentence.id()],
                                sentence.id())
                        for rel in rels:
                            document.add_relation(rel.clone_shared())
            return document
        else:
            new_document = corpus2.Document()
            rel_dict = self.__mix_phrase_channels(phrase_channels)
            for paragraph in document.paragraphs():
                new_paragraph = corpus2.Chunk()
                for sentence in paragraph.sentences():
                    if sentence.id() in rel_dict:
                        asent = \
                            corpus2.AnnotatedSentence_wrap_sentence(sentence)
                        if 'np_sem_roles' not in asent.all_channels():
                            asent.create_channel('np_sem_roles')
                        channel = asent.get_channel('np_sem_roles')
                        rels = \
                            self.__update_chann_and_gener_rels(channel,
                                                        rel_dict[sentence.id()],
                                                        sentence.id())
                        new_sent = \
                            corpus2.AnnotatedSentence_cast_as_sentence(asent)
                        new_paragraph.append(new_sent)
                        for rel in rels:
                            new_document.add_relation(rel.clone_shared())

                    else:
                        new_paragraph.append(sentence)
                new_document.add_paragraph(new_paragraph.clone_shared())
            for rel in document.relations():
                new_document.add_relation(rel.clone_shared())
            return new_document

    def __update_chann_and_gener_rels(self, channel, relations_list,
                                      sentence_id):
        """
        Method updates 'deepened_chunker' annotation channel in one sentence
        and generates ccl relations.

        @type channel: corpus2.AnnotationChannel
        @type relations_list: list
        @type sentence_id: str
        @rtype: list
        """
        output_relations = list()
        for semantic_rel, syntax_rel in relations_list:
            segments = list()
            positions = [syntax_rel.tokens_positions()[0],
                         syntax_rel.tokens_positions()[-1]]
            for position in positions:
                new_segment = channel.get_new_segment_index()
                if channel.get_segment_at(position) == 0:
                    channel.set_segment_at(position, new_segment)
                    segments.append(new_segment)
                else:
                    segments.append(channel.get_segment_at(position))
            relation = self.__generate_relation(semantic_rel, sentence_id,
                                                'np_sem_roles', segments)
            output_relations.append(relation)
        return output_relations

    def __generate_relation(self, relation_name, sentence_id, channel_name,
                                   annotation_numbers):
        """
        Method generates binary relation from syntax relation.

        @type relation_name: str
        @type sentence_id: str
        @type channel_name: str
        @type annotation_numbers: list
        @rtype: corpus2.Relation
        """
        first_point = corpus2.DirectionPoint(sentence_id, channel_name,
                                             annotation_numbers[0])
        second_point = corpus2.DirectionPoint(sentence_id, channel_name,
                                            annotation_numbers[1])
        return corpus2.Relation(relation_name, first_point, second_point)

    def __mix_phrase_channels(self, channels):
        """
        This method mixes dictionaries from 'channels' tuple into one. It's used
        for further generation of new annotations and relations.

        @type channels: tuple
        """
        common_dict = dict()
        for channel in channels:
            for sent_id, rels_list in channel.items():
                if sent_id in common_dict:
                    common_dict[sent_id].extend(rels_list)
                else:
                    common_dict[sent_id] = rels_list
        return common_dict

    def __get_relations_for_document(self, phrase_type, phrase_head_type,
                                     document):
        """
        It generates and returns map of sentence id's as keys and relations
        inside the sentence as value.

        @type phrase_type: str
        @type phrase_head_type: str
        @type document: corpus2.Document
        @rtype: corpus2.Document
        """
        sentid_rels_map = dict()
        syntax_rel_generator = SyntaxRelationGenerator()
        for np_adjp_phrase in \
            self.phrase_reader. \
                read_np_adjp_phrases_from_corpus2_document(document,
                                                           self.tagset,
                                                           phrase_type,
                                                           phrase_head_type,
                                                           self.only_continous,
                                                           self.res,
                                                           self.silent):
            srels = []
            for rule_string, syntax_rel in \
                syntax_rel_generator.generate_sytax_relations(np_adjp_phrase,
                                                              self.res,
                                                              self.tagset):
                semantic_rel = \
                    self.wccl_ops[rule_string].run_wccl_operator(self.tagset,
                                                        syntax_rel.ann_sentence)
                if semantic_rel:
                    min_segm = min(np_adjp_phrase.segments())
                    syntax_rel.tok_positions = \
                        [x + min_segm for x in syntax_rel.tokens_positions()]
                    srels.append((semantic_rel, syntax_rel))
                    if np_adjp_phrase.sent_id() in sentid_rels_map:
                        xrels = sentid_rels_map[np_adjp_phrase.sent_id()]
                        xrels.extend(srels)
                        sentid_rels_map[np_adjp_phrase.sent_id()] = xrels
                    else:
                        sentid_rels_map[np_adjp_phrase.sent_id()] = srels
        return sentid_rels_map

    def __read_all_phrases_from_document(self, document):
        """
        Reads all phrase channels and generates the relation for single
        corpus2.Document object.
        
        @type document: corpus2.Document
        @rtype: tuple
        """
        np_phrases = self.__read_document_np_phrases_relations(document)
        adjp_phrases = self.__read_document_adjp_phrases_relations(document)
        return np_phrases, adjp_phrases

    def __read_document_np_phrases_relations(self, document):
        """
        Generates relations for np channel.

        @type document: corpus2.Document
        @rtype list
        """
        return self.__get_relations_for_document(
            self.phrase_reader.NP_ANNOTATION,
            self.phrase_reader.NP_HEAD_ANNOTATION,
            document)

    def __read_document_adjp_phrases_relations(self, document):
        """
        Generates relations for adjp channel.

        @type document: corpus2.Document
        @rtype list
        """
        return self.__get_relations_for_document(
            self.phrase_reader.ADJP_ANNOTATION,
            self.phrase_reader.ADJP_HEAD_ANNOTATION,
            document)


def load_wccl_operators(cfg_file, tagset):
    from npsemrel.rb import options

    op_cfg_file = options.Options(cfg_file)
    op_cfg_file.parse_wccl_cfg()
    op_dict = dict()
    for op_name, op_file in op_cfg_file.operators.iteritems():
        print >> sys.stderr, 'Loading operator:', op_name, \
                             'from file:', op_file
        wccl_ops = wccl_op.WcclOp()
        wccl_ops.load_wccl_operators(op_file, tagset)
        op_dict[op_name] = wccl_ops
    rules_dict = dict()
    for rul_nmb, op_name in op_cfg_file.rules.iteritems():
        print >> sys.stderr, 'Setting operator:', op_name, \
                             'to rule number:', rul_nmb
        rules_dict[rul_nmb] = op_dict[op_name]
    return [op_cfg_file, rules_dict]


def print_ccl_file(document, tagset):
    """
    IO method used to print corpus2.Document in ccl format into std.out.
    """
    printer = corpus2.TokenWriter_create_stdout_writer('ccl', tagset)
    for parag in document.paragraphs():
        printer.write_chunk(parag)
    printer.finish()
    del printer


def write_ccl_file(document, document_path, tagset):
    """
    IO method used to print corpus2.Document in ccl format into a file.
    """
    writer = corpus2.TokenWriter_create_path_writer('ccl', document_path,
                                                    tagset)
    for parag in document.paragraphs():
        writer.write_chunk(parag)
    writer.finish()
    del writer


def write_document_relations(document, output_path):
    """
    IO method witch writes relations from document object into a file.

    @type document: corpus2.Document
    @type output_path: str
    """
    rel_writer = corpus2.RelationWriter(output_path)
    rel_writer.write(document.relations())
    del rel_writer


def print_document_relations(document):
    """
    IO method witch writes relations from document object into a file.

    @type document: corpus2.Document
    @type output_path: str
    """
    import string
    import os
    import random

    #Get random string for temporary output
    ran = \
        ''.join([random.choice(string.ascii_letters + string.digits) for n in
                 xrange(32)])
    rel_writer = corpus2.RelationWriter(ran)
    rel_writer.write(document.relations())
    del rel_writer
    fi = open(ran, 'r')
    print fi.read()
    fi.close()
    os.remove(ran)


def read_ccl_file(file_path, tagset):
    """
    IO method used for ccl reading. It returns corpus2.Document.

    @rtype: corpus2.Document
    """
    doc_reader = corpus2.CclRelReader(tagset, str(file_path),
                                      str(file_path))
    doc_reader.set_option("disamb_only")
    document = doc_reader.read()
    del doc_reader
    return document


def read_ccl_stream(tagset):
    """
    IO method for ccl stream reading. It returns corpus2.Document

    @rtype: corpus2.Document
    """
    doc_reader = corpus2.TokenReader_create_stdin_reader('ccl', tagset)
    document = corpus2.Document()
    while True:
        chunk = doc_reader.get_next_chunk()
        if not chunk:
            break
        document.add_paragraph(chunk.clone_shared())
    del doc_reader
    return document


def get_arg_parser():
    import argparse

    try:
        import argcomplete
    except ImportError:
        argcomplete = None
    parser = argparse.ArgumentParser(description='It generates some rels')
    parser.add_argument('input_file', help='Input ccl file. For stream' +
                                           ' use "-"')
    parser.add_argument('-o', '--output', help='Output file')
    parser.add_argument('-t', '--tagset', help='Tagset name', default='nkjp')
    parser.add_argument('-r', '--relation_output')
    if argcomplete:
        argcomplete.autocomplete(parser)
    return parser.parse_args()


def main():
    args = get_arg_parser()
    npsemrel = NpSemrelCclProcessor(
                        'resources/dictionaries/dict-syn-mpar.lex',
                        'resources/dictionaries/dict-gerundium.lex',
                        'resources/dictionaries/dict-pred-snd-type.lex',
                        '../cfg/ops.ini',
                        '../cfg/rand.ini',
                        only_continous=True,
                        silent=True,
                        tagset=args.tagset)

    #Get tagset
    tagset = corpus2.get_named_tagset(args.tagset)

    #Read ccl document for following procedures
    if args.input_file == "-":
        document = read_ccl_stream(tagset)
    else:
        document = read_ccl_file(args.input_file, tagset)

    annotated_document = npsemrel.process_document(document)

    #Return result
    if args.output:
        write_ccl_file(annotated_document, args.output, tagset)
    else:
        print_ccl_file(annotated_document, tagset)
    if args.relation_output:
        write_document_relations(annotated_document, args.relation_output)
    else:
        print_document_relations(annotated_document)


if __name__ == '__main__':
    main()
