from iobber.chunker import Chunker as IobberChunker
import corpus2


class Chunker():
    def __init__(self, config_path=None, model_dir=None):
        """
        Init Chunker processor

        @param config_path: filepath for Chunker's config
        @param model_dir: filepath to model
        """
        self.chunker = IobberChunker(config_path, model_dir, False)
        self.chunker.load_model()

    def chunking_sentence(self, sentences):
        """
        Do Chunker analysis

        @param sentences: list of corpus2.Sentence
        @return: list of chunked corpus2.Sentence
        """
        chunked_sentences = []
        for sent in sentences:
            new_asent = corpus2.AnnotatedSentence.wrap_sentence(sent.sentence)
            new_sent = corpus2.AnnotatedSentence.cast_as_sentence(new_asent)
            self.chunker.tag_sentence(new_sent)
            chunked_sentences.append(new_sent)
        return chunked_sentences

    def chunking_document(self, document, update = False):
        """
        Chunk the document and return it's copy or update the document only.


        @param document: corpus2.Document
        @return: corpus2.Document
        """
        if update:
            for chunk in document.paragraphs():
                for sentence in chunk.sentences():
                    self.chunker.tag_sentence(sentence)
            return document
        else:
            chunked_document = corpus2.Document(path=document.path())
            for chunk in document.paragraphs():
                chunked_chunk = corpus2.Chunk().clone_shared()
                chunked_chunk.set_attribute('id', chunk.get_attribute('id'))
                for sentence in self._chunked_sentences(chunk):
                    chunked_chunk.append(sentence)
                chunked_document.add_paragraph(chunked_chunk)

            return chunked_document

    def _chunked_sentences(self, chunk):
        """Yields chunked sentence

        @param chunk: corpus2.Chunk
        """
        for sentence in chunk.sentences():
            sentence_ptr = corpus2.AnnotatedSentence.wrap_sentence(
                sentence).clone_shared()
            self.chunker.tag_sentence(sentence_ptr)
            yield sentence_ptr