import logging
from typing import List

import numpy as np
import tensorflow as tf
from transformers import TFAutoModelForMaskedLM, AutoTokenizer

from rhyme_with_ai.utils import pairwise
from rhyme_with_ai.token_weighter import TokenWeighter


class RhymeGenerator:
    def __init__(
        self,
        model: TFAutoModelForMaskedLM,
        tokenizer: AutoTokenizer,
        token_weighter: TokenWeighter = None,
    ):
        """Generate rhymes.

        Parameters
        ----------
        model : Model for masked language modelling
        tokenizer : Tokenizer for model
        token_weighter : Class that weighs tokens
        """

        self.model = model
        self.tokenizer = tokenizer
        if token_weighter is None:
            token_weighter = TokenWeighter(tokenizer)
        self.token_weighter = token_weighter
        self._logger = logging.getLogger(__name__)

        self.tokenized_rhymes_ = None
        self.position_probas_ = None

        # Easy access.
        self.comma_token_id = self.tokenizer.encode(",", add_special_tokens=False)[0]
        self.period_token_id = self.tokenizer.encode(".", add_special_tokens=False)[0]
        self.mask_token_id = self.tokenizer.mask_token_id

    def start(self, query: str, rhyme_words: List[str]) -> None:
        """Start the sentence generator.

        Parameters
        ----------
        query : Seed sentence
        rhyme_words : Rhyme words for next sentence
        """
        # TODO: What if no content?
        self._logger.info("Got sentence %s", query)
        tokenized_rhymes = [
            self._initialize_rhymes(query, rhyme_word) for rhyme_word in rhyme_words
        ]
        # Make same length.
        self.tokenized_rhymes_ = tf.keras.preprocessing.sequence.pad_sequences(
            tokenized_rhymes, padding="post", value=self.tokenizer.pad_token_id
        )
        p = self.tokenized_rhymes_ == self.tokenizer.mask_token_id
        self.position_probas_ = p / p.sum(1).reshape(-1, 1)

    def _initialize_rhymes(self, query: str, rhyme_word: str) -> List[int]:
        """Initialize the rhymes.

        * Tokenize input
        * Append a comma if the sentence does not end in it (might add better predictions as it
            shows the two sentence parts are related)
        * Make second line as long as the original
        * Add a period

        Parameters
        ----------
        query : First line
        rhyme_word : Last word for second line

        Returns
        -------
        Tokenized rhyme lines
        """

        query_token_ids = self.tokenizer.encode(query, add_special_tokens=False)
        rhyme_word_token_ids = self.tokenizer.encode(
            rhyme_word, add_special_tokens=False
        )

        if query_token_ids[-1] != self.comma_token_id:
            query_token_ids.append(self.comma_token_id)

        magic_correction = len(rhyme_word_token_ids) + 1  # 1 for comma
        return (
            query_token_ids
            + [self.tokenizer.mask_token_id] * (len(query_token_ids) - magic_correction)
            + rhyme_word_token_ids
            + [self.period_token_id]
        )

    def mutate(self):
        """Mutate the current rhymes.

        Returns
        -------
        Mutated rhymes
        """
        self.tokenized_rhymes_ = self._mutate(
            self.tokenized_rhymes_, self.position_probas_, self.token_weighter.proba
        )

        rhymes = []
        for i in range(len(self.tokenized_rhymes_)):
            rhymes.append(
                self.tokenizer.convert_tokens_to_string(
                    self.tokenizer.convert_ids_to_tokens(
                        self.tokenized_rhymes_[i], skip_special_tokens=True
                    )
                )
            )
        return rhymes

    def _mutate(
        self,
        tokenized_rhymes: np.ndarray,
        position_probas: np.ndarray,
        token_id_probas: np.ndarray,
    ) -> np.ndarray:

        replacements = []
        for i in range(tokenized_rhymes.shape[0]):
            mask_idx, masked_token_ids = self._mask_token(
                tokenized_rhymes[i], position_probas[i]
            )
            tokenized_rhymes[i] = masked_token_ids
            replacements.append(mask_idx)

        predictions = self._predict_masked_tokens(tokenized_rhymes)

        for i, token_ids in enumerate(tokenized_rhymes):
            replace_ix = replacements[i]
            token_ids[replace_ix] = self._draw_replacement(
                predictions[i], token_id_probas, replace_ix
            )
            tokenized_rhymes[i] = token_ids

        return tokenized_rhymes

    def _mask_token(self, token_ids, position_probas):
        """Mask line and return index to update."""
        token_ids = self._mask_repeats(token_ids, position_probas)
        ix = self._locate_mask(token_ids, position_probas)
        token_ids[ix] = self.mask_token_id
        return ix, token_ids

    def _locate_mask(self, token_ids, position_probas):
        """Update masks or a random token."""
        if self.mask_token_id in token_ids:
            # Already masks present, just return the last.
            # We used to return thee first but this returns worse predictions.
            return np.where(token_ids == self.tokenizer.mask_token_id)[0][-1]
        return np.random.choice(range(len(position_probas)), p=position_probas)

    def _mask_repeats(self, token_ids, position_probas):
        """Repeated tokens are generally of less quality."""
        repeats = [
            ii for ii, ids in enumerate(pairwise(token_ids[:-2])) if ids[0] == ids[1]
        ]
        for ii in repeats:
            if position_probas[ii] > 0:
                token_ids[ii] = self.mask_token_id
            if position_probas[ii + 1] > 0:
                token_ids[ii + 1] = self.mask_token_id
        return token_ids

    def _predict_masked_tokens(self, tokenized_rhymes):
        return self.model(tf.constant(tokenized_rhymes))[0]

    def _draw_replacement(self, predictions, token_probas, replace_ix):
        """Get probability, weigh and draw."""
        # TODO (HG): Can't we softmax when calling the model?
        probas = tf.nn.softmax(predictions[replace_ix]).numpy() * token_probas
        probas /= probas.sum()
        return np.random.choice(range(len(probas)), p=probas)