from typing import Optional, Union

import numpy as np

from transformers import AutoTokenizer, DonutImageProcessor
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_utils import ImageInput, is_valid_image
from transformers.processing_utils import (
    MultiModalData,
    ProcessingKwargs,
    ProcessorMixin,
    TextKwargs,
    Unpack,
)
from transformers.tokenization_utils_base import (
    AddedToken,
    PreTokenizedInput,
    TextInput,
)
from transformers.utils import logging


logger = logging.get_logger(__name__)

IMAGE_TOKEN = "<image>"
EXTRA_TOKENS = [f"<loc{i:0>4}>" for i in range(1024)] + [
    f"<seg{i:0>3}>" for i in range(128)
]


# Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/processing_utils.py
class PaliGemmaTextKwargs(TextKwargs):
    suffix: Optional[
        Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]
    ]


class PaliGemmaProcessorKwargs(ProcessingKwargs, total=False):
    text_kwargs: PaliGemmaTextKwargs
    _defaults = {
        "text_kwargs": {
            "padding": False,
            "return_mm_token_type_ids": False,
        },
        "images_kwargs": {
            "data_format": "channels_first",
        },
    }


# Copied from transformers.models.idefics2.processing_idefics2.is_url
def is_url(val) -> bool:
    return isinstance(val, str) and val.startswith("http")


# Copied from transformers.models.idefics2.processing_idefics2.is_image_or_image_url
def is_image_or_image_url(elem):
    return is_url(elem) or is_valid_image(elem)


def _is_str_or_image(elem):
    return isinstance(elem, (str)) or is_image_or_image_url(elem)


def build_string_from_input(prompt, bos_token, image_seq_len, image_token, num_images):
    """
    Builds a string from the input prompt and image tokens.
    For example, for the call:
    build_string_from_input(
        prompt="Prefix str"
        bos_token="<s>",
        image_seq_len=3,
        image_token="<im>",
    )
    The output will be:
    "<im><im><im><s>Initial str"
    Args:
        prompt (`list[Union[str, ImageInput]]`): The input prompt.
        bos_token (`str`): The beginning of sentence token.
        image_seq_len (`int`): The length of the image sequence.
        image_token (`str`): The image token.
        num_images (`int`): Number of images in the prompt.
    """
    return f"{image_token * image_seq_len * num_images}{bos_token}{prompt}\n"


# Copied and adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/paligemma/processing_paligemma.py
class DIVEdocProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "DonutImageProcessor"  # change from the original SigLipImageProcessor to DonutImageProcessor
    tokenizer_class = "GemmaTokenizerFast"
    r"""
    Constructs a PaliGemma processor which wraps a PaliGemma image processor and a PaliGemma tokenizer into a single processor.

    [`PaliGemmaProcessor`] offers all the functionalities of [`SiglipImageProcessor`] and [`GemmaTokenizerFast`]. See the
    [`~PaliGemmaProcessor.__call__`] and [`~PaliGemmaProcessor.decode`] for more information.

    Args:
        image_processor ([`SiglipImageProcessor`], *optional*):
            The image processor is a required input.
        tokenizer ([`GemmaTokenizerFast`], *optional*):
            The tokenizer is a required input.
        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
            in a chat into a tokenizable string.
    """

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        chat_template=None,
        **kwargs,
    ):
        if not hasattr(image_processor, "image_seq_length"):
            raise ValueError(
                "Image processor is missing an `image_seq_length` attribute."
            )

        self.image_seq_length = image_processor.image_seq_length

        if not hasattr(tokenizer, "image_token"):
            image_token = AddedToken(IMAGE_TOKEN, normalized=False, special=True)
            tokens_to_add = {"additional_special_tokens": [image_token]}
            tokenizer.add_special_tokens(tokens_to_add)
            self.image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
            self.image_token = IMAGE_TOKEN
        else:
            self.image_token_id = tokenizer.image_token_id
            self.image_token = tokenizer.image_token

        tokenizer.add_tokens(EXTRA_TOKENS)
        tokenizer.add_bos_token = False
        tokenizer.add_eos_token = False

        super().__init__(image_processor, tokenizer, chat_template=chat_template)

    def __call__(
        self,
        images: Optional[ImageInput] = None,
        text: Union[
            TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]
        ] = None,
        **kwargs: Unpack[PaliGemmaProcessorKwargs],
    ) -> BatchFeature:
        """
        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
        and `kwargs` arguments to GemmaTokenizerFast's [`~GemmaTokenizerFast.__call__`] if `text` is not `None` to encode
        the text. To prepare the image(s), this method forwards the `images` and `kwargs` arguments to
        SiglipImageProcessor's [`~SiglipImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
        of the above two methods for more information.

        The usage for PaliGemma fine-tuning preparation is slightly different than usual. suffix passed are suffixes to
        the prompt in `text`, and will be placed after the prompt. This is because attention is handled differently for
        the prefix and the suffix. For instance,
        ```python
        image = PIL_cow_image
        prompt = "answer en Where is the cow standing?"
        suffix = "on the beach"
        inputs = processor(text=prompt, images=image, suffix=suffix)
        ```
        Here `inputs` will contain the `input_ids` and `token_type_ids` that follow
        ```python
        inputs["input_ids"][:, 256:]
        # tensor([[     2,   6006,    603,    573,  13910,   9980, 235336,    108,    477,   573,   8318]])
        inputs["token_type_ids"][:, 256:]
        tensor([[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1]])
        ```
        Meaning the last three tokens are of "label" ("suffix") type while the other ones are of "prefix" type.


        Args:
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
                number of channels, H and W are image height and width.
            text (`str`, `list[str]`, `list[list[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors of a particular framework. Acceptable values are:

                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
            suffix (`str`, `list[str]`, `list[list[str]]`):
                The suffixes or batch of suffixes to be encoded. Only necessary for finetuning. See https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/README.md
                for more information. If your prompt is "<image> What is on the image", the suffix corresponds to the expected prediction "a cow sitting on a bench".

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. If `suffix`
              is provided, the `input_ids` will also contain the suffix input ids.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **labels** -- Labels compatible with training if `suffix` is not None
        """

        output_kwargs = self._merge_kwargs(
            PaliGemmaProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )
        suffix = output_kwargs["text_kwargs"].pop("suffix", None)

        return_token_type_ids = True

        if images is None:
            raise ValueError(
                "`images` are expected as arguments to a `PaliGemmaProcessor` instance."
            )
        if text is None:
            logger.warning_once(
                "You are using PaliGemma without a text prefix. It will perform as a picture-captioning model."
            )
            text = ""

        if _is_str_or_image(text):
            text = [text]
        elif isinstance(text, list) and _is_str_or_image(text[0]):
            pass

        if text is not None and images is not None:
            if not any(IMAGE_TOKEN in sample for sample in text):
                logger.warning(
                    "You are passing both `text` and `images` to `PaliGemmaProcessor`. The processor expects special "
                    "image tokens in the text, as many tokens as there are images per each text. It is recommended to "
                    "add `<image>` tokens in the very beginning of your text. For this call, we will infer how many images "
                    "each text has and add special tokens."
                )

                if isinstance(text, list) and isinstance(images, list):
                    if len(images) != len(text):
                        raise ValueError(
                            f"Received {len(images)} images for {len(text)} prompts. Each prompt should be associated with an image or list of images."
                        )

                # make a nested list of lists to be able to iterate over the images and text below

                if is_valid_image(images):
                    images = [images]
                elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
                    images = [image for image in images]
                elif not (
                    isinstance(images, (list, tuple))
                    # and isinstance(images[0], (list, tuple))
                    and is_valid_image(images[0])
                ):
                    raise ValueError(
                        "images must be an image, list of images or list of list of images"
                    )

                input_strings = [
                    build_string_from_input(
                        prompt=prompt,
                        bos_token=self.tokenizer.bos_token,
                        image_seq_len=self.image_seq_length,
                        image_token=IMAGE_TOKEN,
                        num_images=len(image_list)
                        if isinstance(image_list, list)
                        else 1,
                    )
                    for prompt, image_list in zip(text, images)
                ]
            else:
                expanded_samples = []
                for sample in text:
                    expanded_sample = sample.replace(
                        IMAGE_TOKEN, IMAGE_TOKEN * self.image_seq_length
                    )
                    bos_rfind_index = expanded_sample.rfind(IMAGE_TOKEN)
                    bos_index = (
                        bos_rfind_index + len(IMAGE_TOKEN)
                        if bos_rfind_index != -1
                        else 0
                    )
                    expanded_sample = (
                        expanded_sample[:bos_index]
                        + self.tokenizer.bos_token
                        + expanded_sample[bos_index:]
                    )
                    expanded_samples.append(expanded_sample)
                input_strings = [f"{sample}\n" for sample in expanded_samples]

        if suffix is not None and _is_str_or_image(suffix):
            suffix = [suffix]
        if suffix is not None:
            suffix = [sfx + self.tokenizer.eos_token for sfx in suffix]
        pixel_values = self.image_processor(images, **output_kwargs["images_kwargs"])[
            "pixel_values"
        ]

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop(
            "return_mm_token_type_ids", None
        )
        inputs = self.tokenizer(
            input_strings,
            text_pair=suffix,
            return_token_type_ids=return_token_type_ids,
            **output_kwargs["text_kwargs"],
        )
        # self._check_special_mm_tokens(input_strings, inputs, modalities=["image"])

        return_data = {**inputs, "pixel_values": pixel_values}

        # TODO: ideally we would control label generation separately, now that we always return token_type_ids.
        if return_token_type_ids:
            labels = np.array(inputs["input_ids"])
            labels[np.array(inputs["token_type_ids"]) == 0] = -100
            return_data.update({"labels": labels})

        if return_mm_token_type_ids:
            array_ids = np.array(return_data["input_ids"])
            mm_token_type_ids = np.zeros_like(return_data["input_ids"])
            mm_token_type_ids[array_ids == self.image_token_id] = 1
            return_data["mm_token_type_ids"] = mm_token_type_ids.tolist()

        return BatchFeature(data=return_data, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.

        Args:
            image_sizes (list[list[str]], *optional*):
                The input sizes formatted as (height, width) per each image.
        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """
        vision_data = {}
        if image_sizes is not None:
            num_image_tokens = [self.image_seq_length] * len(image_sizes)
            num_image_patches = [1] * len(image_sizes)
            vision_data.update(
                {
                    "num_image_tokens": num_image_tokens,
                    "num_image_patches": num_image_patches,
                }
            )
        return MultiModalData(**vision_data)

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names + [
            "token_type_ids",
            "labels",
        ]
        image_processor_input_names = self.image_processor.model_input_names
        return list(tokenizer_input_names + image_processor_input_names)


def get_processor(hf_token, img_height, img_width, img_lm_input_seq_length):
    tokenizer = AutoTokenizer.from_pretrained(
        "google/paligemma-3b-ft-docvqa-896",
        token=hf_token,
        revision="acbe61b1b8507f7c7af03a0d42e9908e7b6d4d5d",
    )
    image_processor = DonutImageProcessor.from_pretrained(
        "naver-clova-ix/donut-base-finetuned-docvqa",
        revision="b19d2e332684b0e2d35d9144ce34047767335cf8",
    )
    image_processor.image_seq_length = img_lm_input_seq_length
    image_processor.size["height"], image_processor.size["width"] = (
        img_height,
        img_width,
    )
    processor = DIVEdocProcessor(tokenizer=tokenizer, image_processor=image_processor)
    return processor