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


from phrase_reader import PhraseReader
from deepened_chunker import SyntaxRelationGenerator
from npsemrel.carrot import resource
import corpus2


class Defender():
    """
    Defender class wraps core functionality of defender. It's mainly usage is
    to process and create annotated ccl files enriched with defender relations.
    """

    def __init__(self, first_type_predicate, gerundium_dict,
                 second_type_predicate, only_continous=True,
                 silent=True, tagset='nkjp'):
        self.tagset = corpus2.get_named_tagset(tagset)
        self.phrase_reader = PhraseReader()
        self.res = resource.Resource(self.tagset)
        self.res.load_fst_type_predicate_dictionary(first_type_predicate)
        self.res.load_ger_dict(gerundium_dict)
        self.res.load_snd_type_predicate_dictionary(second_type_predicate)
        self.only_continous = only_continous
        self.silent = silent

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

        @type document: corpus2.Document
        """
        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
        """
        return self.__get_relations_for_document(
            self.phrase_reader.ADJP_ANNOTATION,
            self.phrase_reader.ADJP_HEAD_ANNOTATION,
            document)

    def read_document_agp_phrases_relations(self, document):
        """
        Generates relations for agp channel.

        @type document: corpus2.Document
        """
        return self.__get_relations_for_document(
            self.phrase_reader.AGP_ANNOTATION,
            self.phrase_reader.AGP_HEAD_ANNOTATION,
            document)

    def read_all_phrases_from_ccl(self, ccl_path):
        """
        Reads ccl file and generates all relations from the document.
        It returns a tuple for all 3 channels of phrases.

        @type ccl_path: str
        """
        tagset = corpus2.get_named_tagset(self.tagset)
        document = read_ccl_file(ccl_path, tagset)
        return self.read_all_phrases_from_document(document)

    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
        """
        np_phrases = self.read_document_np_phrases_relations(document)
        agp_phrases = self.read_document_agp_phrases_relations(document)
        adjp_phrases = self.read_document_adjp_phrases_relations(document)
        return np_phrases, agp_phrases, adjp_phrases

    def update_document_annotations_and_relations(self, document,
            phrase_channels, update = False):
        """
        Methods update annotations od corpus2.Document, and adds

        @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 'deepened_chunker' not in asent.all_channels():
                            asent.create_channel('deepened_chunker')
                        channel = asent.get_channel('deepened_chunker')
                        out_rels = self.__update_chann_and_gener_rels(
                                        channel,
                                        rel_dict[sentence.id()],
                                        sentence.id())
                        for rel in out_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 'deepened_chunker' not in asent.all_channels():
                            asent.create_channel('deepened_chunker')
                        channel = asent.get_channel('deepened_chunker')
                        out_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 out_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 relation in relations_list:
            segments = list()
            for position in relation.tokens_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))
            if len(relation.tokens_positions()) == 3:
                rel1, rel2 =  \
                    self.__generate_triple_relation(sentence_id,
                                                    'deepened_chunker',
                                                    segments)
                output_relations.extend([rel1, rel2])
            else:
                rel = self.__generate_binary_relation(sentence_id,
                                                      'deepened_chunker',
                                                      segments)
                output_relations.append(rel)
        return output_relations

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

        @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('2args_rel', first_point, second_point)

    def __generate_triple_relation(self, sentence_id, channel_name,
                                   annotation_numbers):
        """
        Method generates two relations that represents tri-element syntax
        relation.

        @type sentence_id: str
        @type channel_name: str
        @type annotation_numbers: list
        @rtype: tuple
        """
        first_point = corpus2.DirectionPoint(sentence_id, channel_name,
                                             annotation_numbers[0])
        second_point = corpus2.DirectionPoint(sentence_id, channel_name,
                                            annotation_numbers[1])
        third_point = corpus2.DirectionPoint(sentence_id, channel_name,
                                             annotation_numbers[2])
        first_relation = corpus2.Relation('3args_rel_1st_part',
                                         first_point,
                                         second_point)
        second_relation = corpus2.Relation('3args_rel_2nd_part',
                                         second_point,
                                         third_point)
        return first_relation, second_relation

    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.
        """
        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 _, syntax_rel in \
                syntax_rel_generator.generate_sytax_relations(np_adjp_phrase,
                                                              self.res,
                                                              self.tagset):
                min_segm = min(np_adjp_phrase.segments())
                syntax_rel.tok_positions = \
                        [x + min_segm for x in syntax_rel.tokens_positions()]
                srels.append(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 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='Shallow semantic, phrase ' +
                                                 'parser')
    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()
    defender = Defender('resources/dictionaries/dict-syn-mpar.lex',
                        'resources/dictionaries/dict-gerundium.lex',
                        'resources/dictionaries/dict-pred-snd-type.lex',
                        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)

    #Read relations for all channels
    np, agp, adjp = defender.read_all_phrases_from_document(document)

    #Create new annotations and add relations to document
    annotated_document = \
            defender.update_document_annotations_and_relations(document,
                                                               (np, adjp))

    #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()