# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Testing suite for the PyTorch Llava-NeXT model."""

import unittest

import pytest

from transformers import (
    AutoProcessor,
    BitsAndBytesConfig,
    CLIPVisionConfig,
    LlamaConfig,
    LlavaNextConfig,
    LlavaNextForConditionalGeneration,
    LlavaNextModel,
    is_torch_available,
    is_vision_available,
)
from transformers.testing_utils import (
    Expectations,
    cleanup,
    require_bitsandbytes,
    require_torch,
    slow,
    torch_device,
)

from ...test_image_processing_common import load_test_image
from ...test_modeling_common import floats_tensor
from ...vlm_tester import VLMModelTest, VLMModelTester


if is_torch_available():
    import torch

    from transformers.models.llava_next.modeling_llava_next import image_size_to_num_patches


if is_vision_available():
    pass


class LlavaNextVisionText2TextModelTester(VLMModelTester):
    base_model_class = LlavaNextModel
    config_class = LlavaNextConfig
    conditional_generation_class = LlavaNextForConditionalGeneration
    text_config_class = LlamaConfig
    vision_config_class = CLIPVisionConfig

    def __init__(self, parent, **kwargs):
        kwargs.setdefault("num_patches_per_image", 2)
        # Compute num_image_tokens from LlavaNext's pack_image_features logic
        image_size = kwargs.get("image_size", 8)
        patch_size = kwargs.get("patch_size", 4)
        tokens_per_patch = (image_size // patch_size) ** 2
        height = width = image_size // patch_size
        grid_tokens = height * (width + 1)
        kwargs.setdefault("num_image_tokens", tokens_per_patch + grid_tokens)
        kwargs.setdefault("image_token_index", kwargs.get("image_token_id", 3))
        super().__init__(parent, **kwargs)

    def create_pixel_values(self, batch_size: int | None = None):
        # Override to 5D for patch-based models
        batch_size = batch_size if batch_size is not None else self.batch_size
        """LlavaNext expects 5D pixel_values: (batch_size, num_patches, channels, height, width)"""
        return floats_tensor(
            [
                batch_size,
                self.num_patches_per_image,
                self.num_channels,
                self.image_size,
                self.image_size,
            ]
        )

    def get_additional_inputs(self, config, input_ids, pixel_values, batch_size: int | None = None):
        batch_size = batch_size if batch_size is not None else self.batch_size
        return {
            "image_sizes": torch.tensor([[self.image_size, self.image_size]] * batch_size),
        }

    def get_config(self):
        config = super().get_config()
        # Set grid pinpoints compatible with our small test image size
        config.image_grid_pinpoints = [[self.image_size, self.image_size]]
        return config


@require_torch
class LlavaNextForConditionalGenerationModelTest(VLMModelTest, unittest.TestCase):
    """
    Model tester for `LlavaNextForConditionalGeneration`.
    """

    model_tester_class = LlavaNextVisionText2TextModelTester
    skip_test_image_features_output_shape = True

    @pytest.mark.xfail(reason="This architecture seems to not compute gradients for some layer.")
    def test_training_gradient_checkpointing(self):
        super().test_training_gradient_checkpointing()

    @pytest.mark.xfail(reason="This architecture seems to not compute gradients for some layer.")
    def test_training_gradient_checkpointing_use_reentrant_false(self):
        super().test_training_gradient_checkpointing_use_reentrant_false()

    @pytest.mark.xfail(reason="This architecture seems to not compute gradients for some layer.")
    def test_training_gradient_checkpointing_use_reentrant_true(self):
        super().test_training_gradient_checkpointing_use_reentrant_true()

    @unittest.skip(
        "VLMs need lots of steps to prepare images/mask correctly to get pad-free inputs. Can be tested as part of LLM test"
    )
    def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self):
        pass


@require_torch
class LlavaNextForConditionalGenerationIntegrationTest(unittest.TestCase):
    def setUp(self):
        self.processor = AutoProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
        url = "https://huggingface.co/datasets/hf-internal-testing/transformers-synthetic-assets/resolve/main/images/llava_v1_5_radar.jpg"
        self.image = load_test_image(url)

        self.prompt = "[INST] <image>\nWhat is shown in this image? [/INST]"

    def tearDown(self):
        cleanup(torch_device, gc_collect=True)

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test(self):
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
        )

        inputs = self.processor(images=self.image, text=self.prompt, return_tensors="pt").to(torch_device)

        # The input_ids/pixel_values dumps this used to compare against (`nielsr/test-image`) were recorded from
        # the original third-party radar image, so they cannot be reused once the image is a synthetic stand-in.
        # Re-recording them from this very processor would only compare it against itself, so the check is dropped
        # rather than repointed; the generation assertion below is what still exercises the pipeline end to end.

        # verify generation
        output = model.generate(**inputs, max_new_tokens=100)
        EXPECTED_DECODED_TEXT = "[INST]  \nWhat is shown in this image? [/INST] The image shows a stylized graphic with a hexagonal shape at its center, outlined in yellow. The hexagon is set against a dark background with a grid-like pattern. The overall design has a modern and abstract aesthetic. "

        self.assertEqual(
            self.processor.decode(output[0], skip_special_tokens=True),
            EXPECTED_DECODED_TEXT,
        )

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test_batch(self):
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf", quantization_config=BitsAndBytesConfig(load_in_4bit=True)
        )
        url = "https://huggingface.co/datasets/hf-internal-testing/fixtures-coco/resolve/main/val2017/000000039769.jpg"
        cats_image = load_test_image(url)

        inputs = self.processor(
            images=[self.image, cats_image],
            text=[self.prompt, self.prompt],
            return_tensors="pt",
            padding=True,
        ).to(torch_device)

        # it should not matter whether two images are the same size or not
        output = model.generate(**inputs, max_new_tokens=20)

        EXPECTED_DECODED_TEXT = ['[INST]  \nWhat is shown in this image? [/INST] The image shows a stylized graphic with a hexagonal shape at its center, outlined in yellow', '[INST]  \nWhat is shown in this image? [/INST] The image shows two cats lying on a pink surface, which appears to be a couch or a cush']  # fmt: skip
        self.assertEqual(
            self.processor.batch_decode(output, skip_special_tokens=True),
            EXPECTED_DECODED_TEXT,
        )

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test_unk_token(self):
        # related to (#29835)
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
        )

        prompt_with_unk = "[INST] <image>\nWhat is shown in this <unk> image? [/INST]"
        inputs = self.processor(images=self.image, text=prompt_with_unk, return_tensors="pt")

        # verify single forward pass
        inputs = inputs.to(torch_device)
        with torch.no_grad():
            output = model(**inputs)

        # verify generation
        output = model.generate(**inputs, max_new_tokens=40)
        EXPECTED_DECODED_TEXT = '[INST]  \nWhat is shown in this   image? [/INST] The image shows a stylized graphic with a hexagonal shape at its center, outlined in yellow. The hexagon is set against a dark background with a grid-like pattern. The overall design'  # fmt: skip

        self.assertEqual(
            self.processor.decode(output[0], skip_special_tokens=True),
            EXPECTED_DECODED_TEXT,
        )

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test_batch_different_resolutions(self):
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
        )

        url = "https://huggingface.co/datasets/hf-internal-testing/fixtures-coco/resolve/main/val2017/000000039769.jpg"
        lowres_url = "https://4.img-dpreview.com/files/p/TS560x560~forums/56876524/03975b28741443319e9a94615e35667e"
        cats_image = load_test_image(url)
        lowres_img = load_test_image(lowres_url)

        inputs = self.processor(
            images=[lowres_img, cats_image], text=[self.prompt, self.prompt], return_tensors="pt", padding=True
        ).to(torch_device)
        pixel_values = inputs["pixel_values"]

        # verify pixel values are padded correctly with 0 when one image has more num_patches than the other
        image_num_patches = [
            image_size_to_num_patches(
                image_size=imsize,
                grid_pinpoints=model.config.image_grid_pinpoints,
                patch_size=model.config.vision_config.image_size,
            )
            for imsize in inputs["image_sizes"]
        ]
        for pix_val, num_patch in zip(pixel_values, image_num_patches):
            self.assertTrue(torch.all(pix_val[num_patch:] == 0))  # pad on the right
            for i in range(num_patch):
                self.assertFalse(torch.all(pix_val[i : i + 1] == 0))  # no padding expected in any of patches

        # verify generation
        output = model.generate(**inputs, max_new_tokens=50)
        # fmt: off
        EXPECTED_DECODED_TEXT = Expectations(
            {
                (None, None): '[INST]  \nWhat is shown in this image? [/INST] The image shows two deer, likely fawns, in a grassy area with trees in the background. The setting appears to be a forest or woodland, and the photo is taken during what seems to be either dawn or dusk, given',
                ("xpu", 5): '[INST]  \nWhat is shown in this image? [/INST] The image shows two deer, likely fawns, in a grassy area with trees in the background. The setting appears to be a forest or woodland, and the time of day seems to be either dawn or dusk, given the soft',
            }
        )
        # fmt: on
        self.assertEqual(
            self.processor.decode(output[0], skip_special_tokens=True),
            EXPECTED_DECODED_TEXT.get_expectation(),
        )

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test_batch_matches_single(self):
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
        )

        url = "https://huggingface.co/datasets/hf-internal-testing/fixtures-coco/resolve/main/val2017/000000039769.jpg"
        lowres_url = "https://4.img-dpreview.com/files/p/TS560x560~forums/56876524/03975b28741443319e9a94615e35667e"
        cats_image = load_test_image(url)
        lowres_img = load_test_image(lowres_url)

        inputs_batched = self.processor(
            images=[lowres_img, cats_image], text=[self.prompt, self.prompt], return_tensors="pt", padding=True
        ).to(torch_device)

        inputs_single = self.processor(images=lowres_img, text=self.prompt, return_tensors="pt", padding=True).to(
            torch_device
        )

        # verify generation
        output_batched = model.generate(**inputs_batched, max_new_tokens=50)
        output_single = model.generate(**inputs_single, max_new_tokens=50)
        self.assertEqual(
            self.processor.decode(output_batched[0], skip_special_tokens=True),
            self.processor.decode(output_single[0], skip_special_tokens=True),
        )

    @slow
    @require_bitsandbytes
    def test_small_model_integration_test_full_vision_state_selection(self):
        model = LlavaNextForConditionalGeneration.from_pretrained(
            "llava-hf/llava-v1.6-mistral-7b-hf",
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
        )
        # test that changing `strategy` won't error out
        model.vision_feature_select_strategy = "full"

        inputs = self.processor(text=self.prompt, images=self.image, return_tensors="pt").to(model.device)

        # verify generation
        output = model.generate(**inputs, max_new_tokens=30)
        EXPECTED_DECODED_TEXT = '[INST]  \nWhat is shown in this image? [/INST] The image shows a stylized graphic with a hexagonal shape at its center, outlined in yellow. The hexagon is set against a dark background'  # fmt: skip

        self.assertEqual(
            self.processor.decode(output[0], skip_special_tokens=True),
            EXPECTED_DECODED_TEXT,
        )

    @slow
    def test_granite_vision(self):
        """
        Check the expected output of a granite vision model, which leverages
        multiple vision feature layers and a visual encoder with no CLS (siglip).
        """
        granite_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
        model = LlavaNextForConditionalGeneration.from_pretrained(granite_model_path)
        self.processor = AutoProcessor.from_pretrained(granite_model_path)
        prompt = "<|user|>\n<image>\nWhat is shown in this image?\n<|assistant|>\n"
        inputs = self.processor(text=prompt, images=self.image, return_tensors="pt").to(model.device)

        # verify generation
        output = model.generate(**inputs, max_new_tokens=30)
        EXPECTED_DECODED_TEXT = '<|user|>\n\nWhat is shown in this image?\n<|assistant|>\nA pentagon'  # fmt: skip
        self.assertEqual(
            self.processor.decode(output[0], skip_special_tokens=True),
            EXPECTED_DECODED_TEXT,
        )