<gh_stars>1-10
# -*- coding: utf-8 -*-
import numpy

from ..utils import ChunksTagsTranslator
from .base import IO


class ConllIO(IO):
    """An IO interface of CoNLL-format files. 
    
    Parameters
    ----------
    sentence_sep_starts: List[str]
        * For OntoNotes (i.e., Conll2011, Conll2012), `sentence_sep_starts` should be `["#end", "pt/"]`
        If `document_level` is False, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`; 
        else, it will *skip* the current line. 
    document_sep_starts: List[str]
        * For Conll2003, `document_sep_starts` should be `["-DOCSTART-"]`
        * For OntoNotes (i.e., Conll2011, Conll2012), `document_sep_starts` should be `["#begin"]`
        Whether `document_level` is False or True, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`.
    """
    def __init__(self, 
                 text_col_id=0, 
                 tag_col_id=1, 
                 sep=None, 
                 scheme='BIO1', 
                 tag_sep='-',
                 breaking_for_types=True, 
                 additional_col_id2name=None, 
                 sentence_sep_starts=None, 
                 document_sep_starts=None,
                 document_level: bool=False,
                 encoding=None, 
                 verbose: bool=True, 
                 **token_kwargs):
        self.text_col_id = text_col_id
        self.tag_col_id = tag_col_id
        self.sep = sep
        
        self.tags_translator = ChunksTagsTranslator(scheme=scheme, sep=tag_sep, breaking_for_types=breaking_for_types)
        if additional_col_id2name is None:
            self.additional_col_id2name = {}
        else:
            assert all(isinstance(col_id, int) for col_id in additional_col_id2name.keys())
            assert all(isinstance(col_name, str) for col_name in additional_col_id2name.values())
            self.additional_col_id2name = additional_col_id2name
            
        assert sentence_sep_starts is None or all(isinstance(start, str) for start in sentence_sep_starts)
        self.sentence_sep_starts = [] if sentence_sep_starts is None else sentence_sep_starts
        assert document_sep_starts is None or all(isinstance(start, str) for start in document_sep_starts)
        self.document_sep_starts = [] if document_sep_starts is None else document_sep_starts
        self.document_level = document_level
        super().__init__(is_tokenized=True, encoding=encoding, verbose=verbose, **token_kwargs)
        
        
    def read(self, file_path):
        data = []
        with open(file_path, 'r', encoding=self.encoding) as f:
            text, tags = [], []
            additional = {col_id: [] for col_id in self.additional_col_id2name.keys()}
            
            _is_skipping_last = True
            for line in f:
                line = line.strip()
                
                if self._is_breaking(line):
                    if len(text) > 0:
                        additional_tags = {self.additional_col_id2name[col_id]: atags for col_id, atags in additional.items()}
                        tokens = self._build_tokens(text, additional_tags=additional_tags)
                        chunks = self.tags_translator.tags2chunks(tags)
                        data.append({'tokens': tokens, 'chunks': chunks})
                        
                        text, tags = [], []
                        additional = {col_id: [] for col_id in self.additional_col_id2name.keys()}
                    _is_skipping_last = True
                elif self._is_skipping(line):
                    _is_skipping_last = True
                else:
                    line_seperated = line.split(self.sep)
                    text.append(line_seperated[self.text_col_id])
                    tags.append(line_seperated[self.tag_col_id])
                    for col_id in self.additional_col_id2name.keys():
                        additional[col_id].append(line_seperated[col_id])
                    
                    # Fix for cases like ['I-ORG', '', 'I-ORG'], where the second is skipped. 
                    if (self.tags_translator.scheme == 'BIO1' 
                        and len(tags) >= 2 and tags[-1].startswith('I') and tags[-1][1:] == tags[-2][1:]
                        and _is_skipping_last):
                        tags[-1] = tags[-1].replace('I', 'B', 1)

                    _is_skipping_last = False
                        
            if len(text) > 0:
                additional_tags = {self.additional_col_id2name[col_id]: atags for col_id, atags in additional.items()}
                tokens = self._build_tokens(text, additional_tags=additional_tags)
                chunks = self.tags_translator.tags2chunks(tags)
                data.append({'tokens': tokens, 'chunks': chunks})
            
        return data
    
    
    def _is_sentence_seperator(self, line: str):
        if line.strip() == "":
            return True
        for start in self.sentence_sep_starts:
            if line.startswith(start):
                return True
        return False

    def _is_document_seperator(self, line: str):
        for start in self.document_sep_starts:
            if line.startswith(start):
                return True
        return False

    def _is_breaking(self, line: str):
        if self._is_document_seperator(line):
            return True
        if (not self.document_level) and self._is_sentence_seperator(line):
            return True
        return False

    def _is_skipping(self, line: str):
        return self._is_document_seperator(line) or self._is_sentence_seperator(line)


    def flatten_to_characters(self, data: list):
        additional_keys = [key for key in data[0]['tokens'][0].__dict__.keys() if key not in ('text', 'raw_text')]
        
        new_data = []
        for entry in data:
            tokenized_raw_text = entry['tokens'].raw_text
            char_seq_lens = [len(tok) for tok in tokenized_raw_text]
            cum_char_seq_lens = [0] + numpy.cumsum(char_seq_lens).tolist()
            
            flattened_tokenized_raw_text = [char for tok in tokenized_raw_text for char in tok]
            # Repeat additional-tags for every character in a token
            flattened_additional_tags = {key: [atag for atag, tok in zip(getattr(entry['tokens'], key), tokenized_raw_text) for char in tok] 
                                             for key in additional_keys}
            flattened_tokens = self._build_tokens(flattened_tokenized_raw_text, additional_tags=flattened_additional_tags)
            flattened_chunks = [(label, cum_char_seq_lens[start], cum_char_seq_lens[end]) for label, start, end in entry['chunks']]
            new_data.append({'tokens': flattened_tokens, 'chunks': flattened_chunks})
            
        return new_data

# -*- coding: utf-8 -*-
import numpy

from ..utils import ChunksTagsTranslator
from .base import IO


class ConllIO(IO):
    """An IO interface of CoNLL-format files. 
    
    Parameters
    ----------
    sentence_sep_starts: List[str]
        * For OntoNotes (i.e., Conll2011, Conll2012), `sentence_sep_starts` should be `["#end", "pt/"]`
        If `document_level` is False, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`; 
        else, it will *skip* the current line. 
    document_sep_starts: List[str]
        * For Conll2003, `document_sep_starts` should be `["-DOCSTART-"]`
        * For OntoNotes (i.e., Conll2011, Conll2012), `document_sep_starts` should be `["#begin"]`
        Whether `document_level` is False or True, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`.
    """
    def __init__(self, 
                 text_col_id=0, 
                 tag_col_id=1, 
                 sep=None, 
                 scheme='BIO1', 
                 tag_sep='-',
                 breaking_for_types=True, 
                 additional_col_id2name=None, 
                 sentence_sep_starts=None, 
                 document_sep_starts=None,
                 document_level: bool=False,
                 encoding=None, 
                 verbose: bool=True, 
                 **token_kwargs):
        self.text_col_id = text_col_id
        self.tag_col_id = tag_col_id
        self.sep = sep
        
        self.tags_translator = ChunksTagsTranslator(scheme=scheme, sep=tag_sep, breaking_for_types=breaking_for_types)
        if additional_col_id2name is None:
            self.additional_col_id2name = {}
        else:
            assert all(isinstance(col_id, int) for col_id in additional_col_id2name.keys())
            assert all(isinstance(col_name, str) for col_name in additional_col_id2name.values())
            self.additional_col_id2name = additional_col_id2name
            
        assert sentence_sep_starts is None or all(isinstance(start, str) for start in sentence_sep_starts)
        self.sentence_sep_starts = [] if sentence_sep_starts is None else sentence_sep_starts
        assert document_sep_starts is None or all(isinstance(start, str) for start in document_sep_starts)
        self.document_sep_starts = [] if document_sep_starts is None else document_sep_starts
        self.document_level = document_level
        super().__init__(is_tokenized=True, encoding=encoding, verbose=verbose, **token_kwargs)
        
        
    def read(self, file_path):
        data = []
        with open(file_path, 'r', encoding=self.encoding) as f:
            text, tags = [], []
            additional = {col_id: [] for col_id in self.additional_col_id2name.keys()}
            
            _is_skipping_last = True
            for line in f:
                line = line.strip()
                
                if self._is_breaking(line):
                    if len(text) > 0:
                        additional_tags = {self.additional_col_id2name[col_id]: atags for col_id, atags in additional.items()}
                        tokens = self._build_tokens(text, additional_tags=additional_tags)
                        chunks = self.tags_translator.tags2chunks(tags)
                        data.append({'tokens': tokens, 'chunks': chunks})
                        
                        text, tags = [], []
                        additional = {col_id: [] for col_id in self.additional_col_id2name.keys()}
                    _is_skipping_last = True
                elif self._is_skipping(line):
                    _is_skipping_last = True
                else:
                    line_seperated = line.split(self.sep)
                    text.append(line_seperated[self.text_col_id])
                    tags.append(line_seperated[self.tag_col_id])
                    for col_id in self.additional_col_id2name.keys():
                        additional[col_id].append(line_seperated[col_id])
                    
                    # Fix for cases like ['I-ORG', '', 'I-ORG'], where the second is skipped. 
                    if (self.tags_translator.scheme == 'BIO1' 
                        and len(tags) >= 2 and tags[-1].startswith('I') and tags[-1][1:] == tags[-2][1:]
                        and _is_skipping_last):
                        tags[-1] = tags[-1].replace('I', 'B', 1)

                    _is_skipping_last = False
                        
            if len(text) > 0:
                additional_tags = {self.additional_col_id2name[col_id]: atags for col_id, atags in additional.items()}
                tokens = self._build_tokens(text, additional_tags=additional_tags)
                chunks = self.tags_translator.tags2chunks(tags)
                data.append({'tokens': tokens, 'chunks': chunks})
            
        return data
    
    
    def _is_sentence_seperator(self, line: str):
        if line.strip() == "":
            return True
        for start in self.sentence_sep_starts:
            if line.startswith(start):
                return True
        return False

    def _is_document_seperator(self, line: str):
        for start in self.document_sep_starts:
            if line.startswith(start):
                return True
        return False

    def _is_breaking(self, line: str):
        if self._is_document_seperator(line):
            return True
        if (not self.document_level) and self._is_sentence_seperator(line):
            return True
        return False

    def _is_skipping(self, line: str):
        return self._is_document_seperator(line) or self._is_sentence_seperator(line)


    def flatten_to_characters(self, data: list):
        additional_keys = [key for key in data[0]['tokens'][0].__dict__.keys() if key not in ('text', 'raw_text')]
        
        new_data = []
        for entry in data:
            tokenized_raw_text = entry['tokens'].raw_text
            char_seq_lens = [len(tok) for tok in tokenized_raw_text]
            cum_char_seq_lens = [0] + numpy.cumsum(char_seq_lens).tolist()
            
            flattened_tokenized_raw_text = [char for tok in tokenized_raw_text for char in tok]
            # Repeat additional-tags for every character in a token
            flattened_additional_tags = {key: [atag for atag, tok in zip(getattr(entry['tokens'], key), tokenized_raw_text) for char in tok] 
                                             for key in additional_keys}
            flattened_tokens = self._build_tokens(flattened_tokenized_raw_text, additional_tags=flattened_additional_tags)
            flattened_chunks = [(label, cum_char_seq_lens[start], cum_char_seq_lens[end]) for label, start, end in entry['chunks']]
            new_data.append({'tokens': flattened_tokens, 'chunks': flattened_chunks})
            
        return new_data
# -*- coding: utf-8 -*- An IO interface of CoNLL-format files.           Parameters     ----------     sentence_sep_starts: List[str]         * For OntoNotes (i.e., Conll2011, Conll2012), `sentence_sep_starts` should be `["#end", "pt/"]`         If `document_level` is False, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`;          else, it will *skip* the current line.      document_sep_starts: List[str]         * For Conll2003, `document_sep_starts` should be `["-DOCSTART-"]`         * For OntoNotes (i.e., Conll2011, Conll2012), `document_sep_starts` should be `["#begin"]`         Whether `document_level` is False or True, it will *break* the proceeding sequence and *skip* the current line starting with any of `sentence_sep_starts`. # Fix for cases like ['I-ORG', '', 'I-ORG'], where the second is skipped. # Repeat additional-tags for every character in a token