from __future__ import annotations

import numpy as np
import torch

from ....constants import Backends
from ....registry import register_model
from ....utils import is_backend_available
from ...model import Model
from ...model_outputs import Image2TextOutput
from .beit_roberta_image2text_config import BeitRobertaImage2TextConfig


if is_backend_available(Backends.TRANSFORMERS):
    from transformers import (
        BeitConfig,
        BeitModel,
        GenerationConfig,
        RobertaConfig,
        RobertaForCausalLM,
        VisionEncoderDecoderModel,
    )

if is_backend_available(Backends.PILLOW):
    from PIL import Image

_required_backends = [Backends.TRANSFORMERS, Backends.TOKENIZERS, Backends.PILLOW]


@register_model("beit_roberta_image2text", config_class=BeitRobertaImage2TextConfig)
class BeitRobertaImage2Text(Model):
    """
    BEiT + RoBERTa for image to text
    """

    is_generative = True
    required_backends = _required_backends
    loss_func_name = "cross_entropy"

    def __init__(self, config: BeitRobertaImage2TextConfig, **kwargs):
        super().__init__(config, **kwargs)
        encoder = BeitModel(config=BeitConfig(**self.config.encoder))
        decoder = RobertaForCausalLM(config=RobertaConfig(**self.config.decoder))
        self.beit_roberta = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder)

    def forward(
        self,
        pixel_values,
        decoder_input_ids=None,
        decoder_attention_mask=None,
        encoder_outputs=None,
        past_key_values=None,
        decoder_inputs_embeds=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        **kwargs,
    ):
        outputs = self.beit_roberta(
            pixel_values=pixel_values,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            encoder_outputs=encoder_outputs,
            past_key_values=past_key_values,
            decoder_inputs_embeds=decoder_inputs_embeds,
            labels=None,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        return outputs

    def compute_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        loss = self.loss_func(logits.reshape(-1, self.beit_roberta.decoder.config.vocab_size), labels.reshape(-1))
        return loss

    def generate(self, pixel_values, generation_config=None, **kwargs):
        if generation_config is None:
            generation_config = self.config.to_dict()["generation"]
        generation_config = GenerationConfig(**generation_config)
        outputs = self.beit_roberta.generate(inputs=pixel_values, generation_config=generation_config, **kwargs)

        return outputs

    def preprocess(self, inputs: list[str] | list[np.ndarray] | list[Image.Image] | list[torch.Tensor], **kwargs):
        image_processor = self.preprocessor.image_processor
        processed_outputs = image_processor(inputs, **kwargs)
        return processed_outputs

    def post_process(self, model_outputs, **kwargs):
        tokenizer = self.preprocessor.tokenizer
        decoded_outputs = tokenizer.decode(model_outputs.cpu().numpy().tolist())
        outputs = [Image2TextOutput(text=text) for text in decoded_outputs]
        return outputs
