from typing import List
from text.pipeline.base_pipeline import BasePipeline
from text.pipeline_components.input_parsing.parser import ArpabetSequence, GraphemeSequence, ParsedSequence, parse_grapheme_and_arpabet_sequence
from text.pipeline_components.normalization.espeak_normalization import normalize
from text.pipeline_components.phonetization.espeak_phonemization import EspeakPhonetization
from text.pipeline_components.sequencing.espeak_symbols import Symbols

class EspeakPipeline(BasePipeline):
    """
    This is a new text pipeline for FakeYou that will be introduced in July 2022.
    It introduces Arpabet and individually upgradable components.
    The code is not 1:1 with our notebooks, but we'll go to efforts to make sure that 
    the model input expectations on both sides match one another.
    """

    def __init__(self, lang='en'):
        self.lang = lang
        symbols = Symbols(lang)
        self.symbols_to_sequence = symbols.symbols_to_sequence
        self.arpabet_to_sequence = symbols.arpabet_to_sequence
        self.sequence_to_text = symbols.sequence_to_text
        self.SPACE_INTEGER_ENCODING = symbols.SPACE_INTEGER_ENCODING

        self.phonetization = EspeakPhonetization(lang)

    def user_input_to_sequence(self, input_text: str) -> List[int]:
        """Convert raw user input into a numerical sequence for ML inference."""

        parsed_sequence = parse_grapheme_and_arpabet_sequence(input_text)
        self._normalize_graphemes(parsed_sequence)

        parsed_sequence = self._convert_phonemes(parsed_sequence)
        
        return self._convert_to_integer_sequence(parsed_sequence)

    def sequence_to_text(self, input_sequence: List[int]) -> str:
        return self.sequence_to_text(input_sequence)

    def _normalize_graphemes(self, parsed_sequence: ParsedSequence):
        for i in range(len(parsed_sequence)):
            subsequence = parsed_sequence[i]
            if isinstance(subsequence, GraphemeSequence):
                normalized = normalize(subsequence.sequence, self.lang)
                parsed_sequence[i] = GraphemeSequence(normalized)


    def _convert_phonemes(self, parsed_sequence: ParsedSequence) -> ParsedSequence:
        # NB: We will likely break grapheme runs into multiple arpabet sequences, extending the 
        # output list length.
        updated_sequence = []

        for subsequence in parsed_sequence:
            if isinstance(subsequence, ArpabetSequence):
                updated_sequence.append(subsequence)

            elif isinstance(subsequence, GraphemeSequence):
                converted_sequence = self.phonetization.try_convert_to_phonemes(subsequence.sequence)
                updated_sequence.extend(converted_sequence)

        return updated_sequence

    def _convert_to_integer_sequence(self, parsed_sequence: ParsedSequence) -> List[int]:
        integer_sequence = []
        last_index = len(parsed_sequence) - 1

        for i, subsequence in enumerate(parsed_sequence):
            integer_subsequence = []

            if isinstance(subsequence, ArpabetSequence):
                integer_subsequence = self.arpabet_to_sequence(subsequence.sequence)
            elif isinstance(subsequence, GraphemeSequence):
                integer_subsequence = self.symbols_to_sequence(subsequence.sequence)
                
            if integer_subsequence:
                integer_sequence.extend(integer_subsequence)

                if i < last_index and integer_sequence[-1] != self.SPACE_INTEGER_ENCODING:
                    # Separate words with single spaces if not already present.
                    integer_sequence.append(self.SPACE_INTEGER_ENCODING)

        return integer_sequence