from typing import List, Tuple, Set, Optional

from pymorphy2 import MorphAnalyzer
from razdel import tokenize as regex_tokenize
from b_labs_models import Tokenizer, POSTagger
from russian_tagsets import converters


opencorpora_to_ud_convert = converters.converter('opencorpora-int', 'ud14')


class BaseTokenizer:
    '''
    Base tokenizer interface
    '''

    def split(self, text: str, lemmatize: bool = True) -> List[
        Tuple[str, Optional[str]]
    ]:
        '''
        Each subclass of BaseTokenizer must inplement this method
        Returns list with tuples like that one:
        [('приехать', 'VERB'), ('дом', 'NOUN')]
        '''
        raise NotImplementedError


class RegexTokenizer(BaseTokenizer):
    '''
    Tokenizer based on one of most accurate and fast tokenizers for russian
    language text - razdel. This tokenizer doesn't performs POS tagging
    and lemmatization.
    '''

    def split(self, text: str, lemmatize: bool = False) -> List[
        Tuple[str, Optional[str]]
    ]:
        return [
            (token.text.lower(), None) for token in regex_tokenize(text)
        ]


class BaselineTokenizer(BaseTokenizer):
    '''
    Simple tokenizer that uses pre-trained CRF model from B-Labs and
    takes first form of each word from pymorphy2 output
    '''

    POS_WITHOUT_NORMAL_FORM: Set[str] = {
        'SYMB',
        'NUMB',
        'PNCT',
        'LATN'
    }

    def __init__(self):
        self.tokenizer = Tokenizer()
        self.tagger = POSTagger()
        self.morph = MorphAnalyzer()

    def get_token_text(self, token: str, pos: str, lemmatize: bool = False):

        if pos in self.POS_WITHOUT_NORMAL_FORM or not lemmatize:
            return token.lower()

        text: str = token
        forms: List[Tuple[str, str, str]] = [
            (form.word, form.normal_form, form.tag)
            for form in self.morph.parse(token)
        ]
        for (_, normal_form, tag) in forms:
            if pos in tag:
                text = normal_form
                break
        return text.lower()

    def split(self, text: str, lemmatize: bool = False):
        tokens: List[str] = list(self.tokenizer.split(text))
        parts_of_speech: List[str] = list(self.tagger.tag(tokens))
        return [
            (self.get_token_text(token, pos, lemmatize), pos)
            for token, pos in zip(tokens, parts_of_speech)
        ]


class UDBaselineTokenizer(BaselineTokenizer):
    '''
    Tokenizer that works similar to baseline tokenizer, but also
    converts OpenCorpora part of speech tags to Universal Dependencies
    format, because some word vectors (notably, RusVectores) was trained on UD tagset
    '''

    UNKNOWN_POS_TAG: str = 'X'

    def split(self, text: str, lemmatize: bool = False):
        tokens: List[Tuple[str, str]] = (
            super(UDBaselineTokenizer, self)
            .split(text=text, lemmatize=lemmatize)
        )
        results: List[Tuple[str, str]] = [
            (token, opencorpora_to_ud_convert(pos).split(' ')[0])
            if pos else (token, self.UNKNOWN_POS_TAG) for token, pos in tokens
        ]
        return results
