import logging

""" This module describes the used file formats.

The 'basic_srl' project uses different file types each with their own
file formats. Reading and writing is needed, therefore, this module
offers parsing and writing these formats.
"""

def fix_tid(src_tid):
    """ There are currently four ways to describe indices, which have
    to be unrolled into a tuple:

    'x:y' -> [x, x + 1, .. y], e.g. '2:5' -> [2,3,4,5]
    'x_y_z' -> [min, min + 1, .. , max], e.g. '6_7_9' -> [6,7,8,9]
    'x,y,z' -> [x,y,z], e.g. '1,2,3' -> [1,2,3]
    'x' -> [x], e.g. '42' -> [42]

    The unrolled span is sorted and returnd as a tuple.

    Args:
        src_tid (str): tid to fix
    Returns:
        tuple: The unrolled tid
    """

    if ':' in src_tid:
        start, _, end = src_tid.partition(':')
        span = list(range(int(start), int(end) + 1))
    elif '_' in src_tid:
        data = [int(x) for x in src_tid.split('_')]
        min_el, max_el = min(data), max(data)
        span = list(range(min_el, max_el + 1))
    elif ',' in src_tid:
        span = [int(x) for x in src_tid.split(',')]
    else:
        span = [int(src_tid)]

    span.sort()

    return tuple(span)

def read_corpus(sentence_file, frame_annotations_file):
    """ Reads in a sentence and frame annotation file and returns them grouped.

    The sentences and annotations are returned as two different lists. The sentence
    with index i can be found in sentences[i], its annotations can be found in 
    annotations[i].

    Args:
        sentence_file (str): path to the sentence file
        frame_annotations_file (str): path to the frame annotations file
    Returns:
        sentences (list): The parsed sentences
        annotations (list): The annotations for each sentence
    """
    sentences = []
    annotations = []

    with open(sentence_file, 'r') as sf:
        for line in sf:
            sentence_file_entry = SentenceFileEntry.from_line(line)
            sentences.append(sentence_file_entry)
            annotations.append([])

    with open(frame_annotations_file, 'r') as af:
        for line in af:
            frame_file_entry = FrameFileEntry.from_line(line)
            annotations[frame_file_entry.sentence_number].append(frame_file_entry)

    return sentences, annotations

def read_corpus_with_babelfied_synset(sentence_file, frame_annotations_file, babelfied_synset_annotation_file):
    """ Reads in a sentence and frame annotation file and returns them grouped.

    The sentences and annotations are returned as two different lists. The sentence
    with index i can be found in sentences[i], its annotations can be found in 
    annotations[i].

    Args:
        sentence_file (str): path to the sentence file
        frame_annotations_file (str): path to the frame annotations file
    Returns:
        sentences (list): The parsed sentences
        annotations (list): The annotations for each sentence
    """
    sentences = []
    annotations = []
    babelfied_synset_annotations = []

    with open(sentence_file, 'r') as sf:
        for line in sf:
            sentence_file_entry = SentenceFileEntry.from_line(line)
            sentences.append(sentence_file_entry)
            annotations.append([])

    with open(frame_annotations_file, 'r') as af:
        for line in af:
            frame_file_entry = FrameFileEntry.from_line(line)
            annotations[frame_file_entry.sentence_number].append(frame_file_entry)    

    with open(babelfied_synset_annotation_file, 'r') as nf:
        for line in nf:
            babelfied_synset_file_entry = BabelfiedSynsetFileEntry.from_line(line)
            babelfied_synset_annotations.append(babelfied_synset_file_entry)

    return sentences, annotations, babelfied_synset_annotations

def write_corpus(output_sentence_file, output_annotation_file, sentence_entries, annotation_entries):
    """ Writes sentence and frame annotation file and returns them grouped.

    Args:
        output_sentence_file (str): path to the sentence file
        output_annotation_file (str): path to the sentence file
        sentence_entries (list): List of SentenceFileEntry
        annotation_entries (list): The annotations for each sentence. annotations[idx] is a list of FrameFileEntry that belongs to sentence[idx]
    """
    sentence_count = 0
    annotations_count = 0
    with open(output_sentence_file, 'w') as sf, open(output_annotation_file, 'w') as af:
        for sentence_id, (sentence_entry, frame_entries) in enumerate(zip(sentence_entries, annotation_entries)):
            sentence_count += 1
            sf.write(sentence_entry.to_line())
            sf.write('\n')
            for frame_entry in frame_entries:
                frame_entry.sentence_number = sentence_id
                af.write(frame_entry.to_line())
                af.write('\n')
                annotations_count += 1

    logging.info('Written corpus to [%s] and [%s] with [%d] sentences and [%d] annotations', output_sentence_file, output_annotation_file, sentence_count, annotations_count)


def load_sentences(sentence_file):
    sentences = []

    with open(sentence_file, 'r') as f:
        for line in f:
            sentence_file_entry = SentenceFileEntry.from_line(line)
            sentences.append(sentence_file_entry)

    return sentences    

def load_annotations(annotations_file):
    annotations = []

    with open(annotations_file, 'r') as f:
        for line in f:
            frame_file_entry = FrameFileEntry.from_line(line)
            annotations.append(frame_file_entry)

    return annotations    

def build_frame_to_id_mapping(*corpora):
    """ Given a list of corpora, assign an id to each frame in the corpus """
    frame_names = set()

    for file_name in corpora:
        annotations = load_annotations(file_name)
        for annotation in annotations:
            frame_names.add(annotation.frame_name)

    frame_to_id_mapping = {}
    for i, frame_name in enumerate(frame_names):
        frame_to_id_mapping[frame_name] = i

    return frame_to_id_mapping    

class SentenceFileEntry(object):
    """ Represents one line in the sentence file format (lemma.tags).

    The file format itself is tab seperated. It first lists the number of tokens N in 
    the sentence. Then follow 6 groups of N tokens, each for a different type of annotation.
    The following listing illustrates the different annotations. 

    [# tokens]
    [tokens]
    [POS tags]
    [dependency labels]
    [dependency heads]
    [O (optional)]
    [lemmas]
    """

    def __init__(self, tokens, pos_tags, dependency_labels, dependency_heads, lemmas):
        self.number_of_tokens = len(tokens)

        assert len(pos_tags) == self.number_of_tokens
        assert len(dependency_labels) == self.number_of_tokens
        assert len(dependency_heads) == self.number_of_tokens
        assert len(lemmas) == self.number_of_tokens

        self.tokens = tokens
        self.pos_tags = pos_tags
        self.dependency_labels = dependency_labels
        self.dependency_heads = [int(x) for x in dependency_heads]
        self.lemmas = lemmas

    @staticmethod
    def from_line(line):
        """ Parses a line in the sentence file format and constructs a SentenceFileEntry

        Args:
            line (str): The line to parse

        Returns:
            SentenceFileEntry: The parsed entry

        Raises:
            ValueError: If line has invalid number of fields
        """
        fields = line.strip('\n').split('\t')
        n = int(fields[0]) 
        data = fields[1:]

        if len(data) != 6 * n:
            print(line, data)
            raise ValueError('Line has invalid length! Expected {0}, was: {1}'.format(6 * n, len(data)))

        tokens = data[:n]
        pos_tags = data[n:2 * n]
        dependency_labels = data[2*n:3*n]
        dependency_heads = data[3*n:4*n]
        lemmas = data[5*n:6*n]

        return SentenceFileEntry(tokens, pos_tags, dependency_labels, dependency_heads, lemmas)

    def to_line(self, sep='\t'):
        """ Formats this entry so that it can be saved in a sentence file.

        Args:
            sep (:obj:`str`, optional): Seperator between fields. Defaults to '\t'

        Returns:
            str: This entry formatted
        """        
        n = len(self.tokens)
        data = []
        data.append(n)
        data.extend(self.tokens)
        data.extend(self.pos_tags)
        data.extend(self.dependency_labels)
        data.extend(self.dependency_heads)
        data.extend([0] * n)
        data.extend(self.lemmas)

        return sep.join([str(x) for x in data])

    def __str__(self):
        line = self.to_line(' ')
        return line.encode('ascii', 'backslashreplace')

class FrameFileEntry(object):
    """ Represents one line in the frame annotation format.

    The file format itself is tab seperated. It has the following fields:

    [0]
    [0]
    [# of roles]
    [frame name]
    [lemma.pos]
    [position of the FEE in the sentence]
    [FEE string]
    [line# in the sentence file (incl. 0)]
    [role1]
    [position1]
    ...
    [roleX]
    [positionX]
    """

    def __init__(self, sentence_number, frame_name, lemma_pos, fee_position, fee):
        self.sentence_number = sentence_number
        self.frame_name = frame_name
        self.lemma_pos = lemma_pos
        self.fee_position = fee_position
        self.fee = fee
        self.roles = []
        self.positions = []        

    def add_role(self, role, position):
        """ Adds a role to this entry.

        Args:
            role (str): The role to add
            position (int): The position in the sentence of the role to add
        """
        self.roles.append(role)
        self.positions.append(position)

    @staticmethod
    def from_line(line):
        """ Parses a line in the frame annotation format and constructs a FrameFileEntry from it

        Args:
            line (str): The line to parse

        Returns:
            FrameFileEntry: The parsed entry

        Raises:
            ValueError: If line has invalid number of fields
        """        
        fields = line.strip().split('\t')
        n = int(fields[2])
        frame_name = fields[3]
        lemma_pos = fields[4]
        fee_position = fields[5]
        fee = fields[6]
        sentence_number = int(fields[7])

        role_positions = fields[8:]

        if len(role_positions) % 2 != 0:
            raise ValueError('Line has imbalanced number of roles and positions!')

        entry = FrameFileEntry(sentence_number, frame_name, lemma_pos, fee_position, fee)

        for i in range(0, len(role_positions), 2):
            role = role_positions[i]
            position = role_positions[i+1]
            entry.add_role(role, position)

        return entry

    def to_line(self):
        """ Formats this entry so that it can be saved in a frame annotation file.

        Returns:
            str: This entry formatted
        """            
        assert len(self.roles) == len(self.positions), 'Has to have as many roles as role positions!'


        data = ['0', '0', str(len(self.roles) + 1)]
        data.append(self.frame_name)
        data.append(self.lemma_pos)
        data.append(self.fee_position)
        data.append(self.fee)
        data.append(self.sentence_number)

        for pos, role in zip(self.positions, self.roles):
            data.append(role)
            data.append(pos)

        return '\t'.join([str(x) for x in data])

    def to_lexicon_entry(self):
        """ Formats this entry so that it can be saved in a lexicon file.

        Returns:
            str: This entry formatted as a lexicon entry
        """

        return '{0}\t{1}'.format(self.frame_name, self.lemma_pos)
    
class BabelfiedSynsetFileEntry(object):
    
    def __init__(self, sentence_number, babelfied_synset_ids):
        self.sentence_number = sentence_number
        self.babelfied_synset_ids = babelfied_synset_ids
        self.number_of_babelfied_synset_ids = len(babelfied_synset_ids)

    @staticmethod
    def from_line(line):
        """ Parses a line in the sentence file format and constructs a SentenceFileEntry

        Args:
            line (str): The line to parse

        Returns:
            SentenceFileEntry: The parsed entry

        Raises:
            ValueError: If line has invalid number of fields
        """
        fields = line.strip('\n').split('\t')
        sentence_number = int(fields[0]) 
        babelfied_synset_ids = fields[1:]

        return BabelfiedSynsetFileEntry(sentence_number, babelfied_synset_ids)    

class LexiconEntry(object):
    """ Represents one line in the frame annotation format. The format is  <frame_name>\t<lemma_pos> """

    def __init__(self, frame_name, lemma_pos):
        self.frame_name = frame_name
        self.lemma_pos = lemma_pos

    @staticmethod
    def from_line(line):
        """ Parses a line in the lexicon format and constructs a LexiconEntry from it

        Args:
            line (str): The line to parse

        Returns:
            LexiconEntry: The parsed entry
        """          
        frame_name, lemma_pos = line.strip().split('\t')
        return LexiconEntry(frame_name, lemma_pos)

    def to_line(self):
        """ Formats this entry so that it can be saved in a lexicon file.

        Returns:
            str: This entry formatted as a lexicon entry
        """        
        return '{0}\t{1}'.format(self.frame_name, self.lemma_pos)