import torch
from transformers import AutoTokenizer, PreTrainedModel
from transformers.models.t5.configuration_t5 import T5Config

from .configuration_hyperstyler import HyperStylerConfig
from .hyperstyler import HyperStyler
# Imported here so that Transformers copies these files with the remote code.
from .modeling_t5_hyperstyler import T5ForConditionalGeneration_HyperStyler  # noqa: F401
from .config_t5 import T5Config as _T5Config  # noqa: F401
from .style_utils import load_style_model, text_to_style


class HyperStylerModel(PreTrainedModel):
    config_class = HyperStylerConfig
    base_model_prefix = "hyperstyler"
    # T5 shares these with hyperstyler.model.shared.weight.
    _tied_weights_keys = [
        "hyperstyler.model.encoder.embed_tokens.weight",
        "hyperstyler.model.decoder.embed_tokens.weight",
    ]

    def __init__(self, config):
        super().__init__(config)
        t5_config = T5Config.from_pretrained(config.base_model_name)
        self.hyperstyler = HyperStyler(
            config.base_model_name, config.styler_config, use_style=True, t5_config=t5_config
        )
        self._tokenizer = None
        self._style_model = None

    def _init_weights(self, module):
        pass

    @property
    def tokenizer(self):
        if self._tokenizer is None:
            self._tokenizer = AutoTokenizer.from_pretrained(self.config.base_model_name, legacy=True)
        return self._tokenizer

    def style_embeddings(self, texts):
        """AnnaWegmann/Style-Embedding vectors for a list of texts: [len(texts), 768]."""
        if self._style_model is None:
            model, tokenizer, _ = load_style_model(self.config.style_model_name)
            self._style_model = (model.to(self.device).eval(), tokenizer)
        model, tokenizer = self._style_model
        embeds = text_to_style(model=model, tokenizer=tokenizer, texts=texts, device=self.device, model_type="style")
        return torch.stack(embeds)

    def forward(self, input_ids, attention_mask, labels=None, style_embs=None, style_masks=None, **kwargs):
        return self.hyperstyler(input_ids=input_ids, attention_mask=attention_mask, labels=labels,
                                style_embs=style_embs, style_masks=style_masks, **kwargs)

    def generate(self, input_ids, attention_mask, style_embs, style_masks, **kwargs):
        return self.hyperstyler.generate(input_ids=input_ids, attention_mask=attention_mask,
                                         style_embs=style_embs, style_masks=style_masks, **kwargs)

    @torch.no_grad()
    def transfer(self, source_texts, target_texts, temperature=0.8, top_p=1.0, do_sample=True, **generate_kwargs):
        """Rewrite source text(s) in the style of the author of ``target_texts``.

        Args:
            source_texts: a string or a list of strings.
            target_texts: example texts by the target author.
        Returns:
            a string, or a list of strings if ``source_texts`` is a list.
        """
        single = isinstance(source_texts, str)
        sources = [source_texts] if single else list(source_texts)

        style = self.style_embeddings(list(target_texts)).to(self.device)        # [n, 768]
        style_embs = style.unsqueeze(0).expand(len(sources), -1, -1)             # [B, n, 768]
        style_masks = torch.ones(style_embs.shape[:2], device=self.device)

        inputs = self.tokenizer(sources, return_tensors="pt", padding="max_length",
                                max_length=self.config.max_encoder_len, truncation=True).to(self.device)
        generate_kwargs.setdefault("max_new_tokens", self.config.max_decoder_len)
        out = self.generate(inputs["input_ids"], inputs["attention_mask"], style_embs, style_masks,
                            do_sample=do_sample, temperature=temperature, top_p=top_p, **generate_kwargs)
        texts = self.tokenizer.batch_decode(out, skip_special_tokens=True)
        return texts[0] if single else texts