from parlai.core.opt import Opt
from parlai.utils.typing import TShared
from parlai.agents.transformer.transformer import TransformerGeneratorAgent

from .generation_methods import VocabTopKSampling, RerankedTopKSampling
from .generation_utils import Wordlist, Reranker, load_wordlist, cefr_to_int

class ControllableBlender(TransformerGeneratorAgent):
    def __init__(self, opt: Opt, shared: TShared = None):
        super().__init__(opt, shared)

        if opt.get("inference", None) == "vocab":
            wordlist_path = opt.get("wordlist_path", None)
            assert wordlist_path, "Please provide path to vocab list, in order to use inference method 'vocab'"

            allowed_words = load_wordlist(wordlist_path)
            self.wordlist = Wordlist(allowed_words, self.dict)

        elif opt.get("inference", None) == "rerank":
            cefr = opt.get("rerank_cefr", None)
            assert cefr, "Please provide CEFR level, in order to use inference method 'rerank'"

            rerank_tokenizer = opt.get("rerank_tokenizer", None)
            rerank_model = opt.get("rerank_model", None)
            assert rerank_model, "Please provide path to directory containing model weights, in order to use inference method 'rerank'"

            device = opt.get("complexity_model_device", None)
            penalty_stddev = opt.get("penalty_stddev", None)
            text_truncate = opt.get("text_truncate", None)

            word_filter = None
            filter_path = opt.get("filter_path", "")
            if filter_path:
                word_filter = load_wordlist(filter_path)

            exempt_tokens = [self.dict.tok2ind.get(self.dict.null_token),
                             self.dict.tok2ind.get(self.dict.start_token),
                             self.dict.tok2ind.get(self.dict.end_token),
                             self.dict.tok2ind.get(self.dict.unk_token)]

            if penalty_stddev < 0:
                exempt_tokens = "all"

            self.reranker = Reranker(cefr=cefr_to_int(cefr),
                                     model=rerank_model,
                                     tokenizer=rerank_tokenizer,
                                     device=device,
                                     text_truncate=text_truncate,
                                     exempt_tokens=exempt_tokens,
                                     penalty_stddev=penalty_stddev,
                                     vocab_size=len(self.dict),
                                     word_filter=word_filter)

        else:
            raise ValueError(f"Inference method {opt.get('inference', None)} does not exist. "
                             f"Please use 'vocab' or 'rerank'.")


    def _treesearch_factory(self, device, verbose=False):
        method = self.opt.get('inference', 'greedy')
        beam_size = self.opt.get('beam_size', 1)
        if method == 'vocab':
            return VocabTopKSampling(
                k=self.opt.get('topk', 40),
                wordlist=self.wordlist,
                beam_size=beam_size,
                min_length=self.beam_min_length,
                block_ngram=self.beam_block_ngram,
                context_block_ngram=self.beam_context_block_ngram,
                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                padding_token=self.NULL_IDX,
                bos_token=self.START_IDX,
                eos_token=self.END_IDX,
                device=device,
                verbose=verbose,
            )
        elif method == "rerank":
            return RerankedTopKSampling(
                k=self.opt.get('topk', 40),
                reranker=self.reranker,
                tokenids_to_text=self._v2t,
                beam_size=beam_size,
                min_length=self.beam_min_length,
                block_ngram=self.beam_block_ngram,
                context_block_ngram=self.beam_context_block_ngram,
                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                padding_token=self.NULL_IDX,
                bos_token=self.START_IDX,
                eos_token=self.END_IDX,
                device=device,
                verbose=verbose,
            )
        else:
            return super()._treesearch_factory(device, verbose=verbose)

    def share(self):
        """
        Share internal states between parent and child instances.
        """
        shared = super().share()
        if hasattr(self, 'wordlist'):
            shared['wordlist'] = self.wordlist
        if hasattr(self, 'reranker'):
            shared['reranker'] = self.reranker
        return shared