# Copyright 2026 NAVER Corp. and 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 HyperClovaX Vision model."""

import unittest

from transformers import AutoModelForImageTextToText, AutoProcessor, HyperCLOVAXVisionV2Config, is_torch_available
from transformers.testing_utils import (
    Expectations,
    require_torch,
    require_torch_accelerator,
    require_torch_accelerator_memory,
    slow,
    torch_device,
)

from ...test_memory_cleanup_mixin import MemoryCleanupMixin
from ...test_modeling_common import floats_tensor, ids_tensor
from ...test_processing_common import url_to_local_path
from ...vlm_tester import VLMModelTest, VLMModelTester


if is_torch_available():
    import torch

    from transformers import (
        HyperCLOVAXVisionV2Config,
        HyperCLOVAXVisionV2ForConditionalGeneration,
        HyperCLOVAXVisionV2Model,
    )


class HCXVisionV2VisionText2TextModelTester(VLMModelTester):
    base_model_class = HyperCLOVAXVisionV2Model
    config_class = HyperCLOVAXVisionV2Config

    conditional_generation_class = HyperCLOVAXVisionV2ForConditionalGeneration

    def __init__(self, parent, **kwargs):
        kwargs.setdefault("image_size", 28)
        kwargs.setdefault("patch_size", 14)
        kwargs.setdefault("num_heads", 4)
        kwargs.setdefault("depth", 2)
        kwargs.setdefault("spatial_merge_size", 1)
        kwargs.setdefault("out_hidden_size", 32)
        kwargs.setdefault("tokens_per_second", 1)
        kwargs.setdefault("temporal_patch_size", 2)
        kwargs.setdefault("video_token_id", 4)

        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
        return floats_tensor(
            [
                batch_size * (self.image_size**2) // (self.patch_size**2),
                self.num_channels * (self.patch_size**2) * self.temporal_patch_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
        image_grid_thw = torch.tensor([[1, 2, 2]] * batch_size, device=torch_device)
        return {"image_grid_thw": image_grid_thw}


@require_torch
class HyperCLOVAXVisionV2ModelTest(VLMModelTest, unittest.TestCase):
    model_tester_class = HCXVisionV2VisionText2TextModelTester

    test_cpu_offload = False
    model_split_percents = [0.5, 0.7, 0.8]

    def test_video_forward(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        B = self.model_tester.batch_size
        C = config.vision_config.in_channels
        T = config.vision_config.temporal_patch_size
        P = config.vision_config.patch_size
        seq_length = 15

        input_ids = ids_tensor([B, seq_length], self.model_tester.vocab_size)

        F = 4
        patch_H = self.model_tester.image_size // P
        patch_W = self.model_tester.image_size // P
        patch_T = F // T
        patches_per_video = patch_T * patch_H * patch_W
        pixel_values_videos = floats_tensor(
            [
                # first dim: batch_size * num_patches
                B * patches_per_video,
                # second dim: in_channels * temporal_patch_size * patch_size^2
                C * T * (P**2),
            ]
        )
        video_grid_thw = torch.tensor([[patch_T, patch_H, patch_W]] * B, device=torch_device)

        # sanity check
        self.assertEqual(pixel_values_videos.shape[0], video_grid_thw.prod(dim=1).sum().item())

        # Insert video token sequence
        input_ids[:, -1] = self.model_tester.pad_token_id
        input_ids[input_ids == self.model_tester.video_token_id] = self.model_tester.pad_token_id
        input_ids[input_ids == self.model_tester.image_token_id] = self.model_tester.pad_token_id
        input_ids[:, self.model_tester.num_image_tokens] = self.model_tester.video_token_id

        insertion_point = self.model_tester.num_image_tokens

        self.assertLessEqual(patches_per_video + insertion_point, seq_length)
        for b in range(B):
            input_ids[b, insertion_point : insertion_point + patches_per_video] = self.model_tester.video_token_id

        for model_class in self.all_model_classes:
            second_per_grid_ts = torch.tensor([1.0] * B, device=torch_device)
            model = model_class(config).to(torch_device)
            outputs = model(
                input_ids=input_ids,
                pixel_values_videos=pixel_values_videos,
                video_grid_thw=video_grid_thw,
                second_per_grid_ts=second_per_grid_ts,
            )
            self.assertIsNotNone(outputs)

    def test_reverse_loading_mapping(self, check_keys_were_modified=True):
        super().test_reverse_loading_mapping(check_keys_were_modified, skip_base_model=True)


@require_torch
@require_torch_accelerator_memory(memory=100)
class HyperCLOVAXVisionV2IntegrationTest(MemoryCleanupMixin, unittest.TestCase):
    model_id = "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B"

    def setUp(self):
        super().setUp()
        self.processor = AutoProcessor.from_pretrained(self.model_id)
        self.model = HyperCLOVAXVisionV2ForConditionalGeneration.from_pretrained(
            self.model_id, dtype=torch.bfloat16, device_map="auto"
        )

    @slow
    def test_text_generate(self):
        messages = [
            {"role": "user", "content": "What is the capital of South Korea?"},
        ]
        text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.processor(text=[text], return_tensors="pt").to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                (None, None): [
                    "user\nWhat is the capital of South Korea?\nassistant\n<think>\nOkay, so I need to figure out what the capital of South Korea is. Let me start by recalling any prior knowledge I have. I remember that"
                ],
                ("rocm", (9, 4)): [
                    "user\nWhat is the capital of South Korea?\nassistant\n<think>\nOkay, so I need to figure out what the capital of South Korea is. Let me start by recalling any basic geography I know. I remember that"
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)

    @slow
    def test_image_generate(self):
        image_url = (
            "https://huggingface.co/datasets/hf-internal-testing/fixtures-captioning/resolve/main/cow_beach_1.png"
        )
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "url": url_to_local_path(image_url)},
                    {"type": "text", "text": "What animal is in the image?"},
                ],
            }
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                (None, None): [
                    'user\n{"id": "image_00", "type": "image/jpeg", "filename": "a.jpg"}\n\nWhat animal is in the image?\nassistant\n<think>\nOkay, so I need to figure out what animal is in the image based on the details provided. Let me start by reading through the image carefully.\n\n'
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)

    @slow
    def test_video_generate(self):
        video_url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_videos/resolve/main/tennis.mp4"

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "video", "url": url_to_local_path(video_url)},
                    {"type": "text", "text": "What is shown in this video?"},
                ],
            }
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                (None, None): [
                    'user\n{"id": "video_00", "type": "video/mp4", "filename": "a.mp4"}\n<|video_aux_start|>다음 중 video_duration은 비디오 길이 정보입니다. 참고하여 답변하세요. {"video_duration": 6.07}<|video_aux_end|>\n\nWhat is shown in this video?\nassistant\n<think>\nOkay, so I need to figure out what\'s shown in this video based on the image provided. Let me start by breaking down the details given.\n\n'
                ],
                ("cuda", (8, 6)): [
                    'user\n{"id": "video_00", "type": "video/mp4", "filename": "a.mp4"}\n<|video_aux_start|>다음 중 video_duration은 비디오 길이 정보입니다. 참고하여 답변하세요. {"video_duration": 6.07}<|video_aux_end|>\n\nWhat is shown in this video?\nassistant\n<think>\nOkay, so I need to figure out what\'s shown in the video based on the image provided. Let me start by breaking down the details given.\n\n'
                ],
                ("cuda", (10, 0)): [
                    'user\n{"id": "video_00", "type": "video/mp4", "filename": "a.mp4"}\n<|video_aux_start|>다음 중 video_duration은 비디오 길이 정보입니다. 참고하여 답변하세요. {"video_duration": 6.07}<|video_aux_end|>\n\nWhat is shown in this video?\nassistant\n<think>\nOkay, so I need to figure out what\'s shown in this video based on the image provided. Let me start by breaking down the details given.\n\n'
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)


@require_torch
@require_torch_accelerator
class HyperCLOVAXVisionV2TinyModelIntegrationTest(MemoryCleanupMixin, unittest.TestCase):
    model_id = "hf-internal-testing/tiny-random-HyperCLOVAXVisionV2ForConditionalGeneration"

    def setUp(self):
        super().setUp()
        self.processor = AutoProcessor.from_pretrained(self.model_id)
        self.model = AutoModelForImageTextToText.from_pretrained(
            self.model_id, dtype=torch.bfloat16, device_map="auto"
        )

    @slow
    def test_text_generate(self):
        messages = [
            {"role": "user", "content": "What is the capital of South Korea?"},
        ]
        text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.processor(text=[text], return_tensors="pt").to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                ("cuda", (8, 6)): [
                    "user\nWhat is the capital of South Korea?\nassistant\n<think>\n globally,www locally,www locally,www locally,www locally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www"
                ],
                ("xpu", 5): [
                    "user\nWhat is the capital of South Korea?\nassistant\n<think>\n globally,www locally,www locally,www locally,www locally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www"
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)

    @slow
    def test_image_generate(self):
        image_url = (
            "https://huggingface.co/datasets/hf-internal-testing/fixtures-captioning/resolve/main/cow_beach_1.png"
        )
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "url": url_to_local_path(image_url)},
                    {"type": "text", "text": "What animal is in the image?"},
                ],
            }
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                ("cuda", (8, 6)): [
                    'user\n{"id": "image_00", "type": "image/jpeg", "filename": "a.jpg"}\n\nWhat animal is in the image?\nassistant\n<think>\n globally,www locally([@ globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www'
                ],
                ("xpu", 5): [
                    'user\n{"id": "image_00", "type": "image/jpeg", "filename": "a.jpg"}\n\nWhat animal is in the image?\nassistant\n<think>\n globally,www locally([@ globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www globally,www'
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)

    @slow
    def test_video_generate(self):
        video_url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_videos/resolve/main/tennis.mp4"

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "video", "url": url_to_local_path(video_url)},
                    {"type": "text", "text": "What is shown in this video?"},
                ],
            }
        ]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to(torch_device)

        output = self.model.generate(**inputs, max_new_tokens=30, do_sample=False)
        output_text = self.processor.batch_decode(output, skip_special_tokens=True)

        EXPECTED_TEXTS = Expectations(
            {
                ("cuda", ("cuda", (8, 6))): [
                    'user\n{"id": "video_00", "type": "video/mp4", "filename": "a.mp4"}\n<|video_aux_start|>다음 중 video_duration은 비디오 길이 정보입니다. 참고하여 답변하세요. {"video_duration": 6.07}<|video_aux_end|>\n\nWhat is shown in this video?\nassistant\n<think>\n globally신지차도 globally신지파괴의장 globally신지예원 sometimes globally신지신분 =================================================================================([@ globally신지접이 globally신지과이 globally신지 globally신지 globally신지 globally신지'
                ],
                ("xpu", 5): [
                    'user\n{"id": "video_00", "type": "video/mp4", "filename": "a.mp4"}\n<|video_aux_start|>다음 중 video_duration은 비디오 길이 정보입니다. 참고하여 답변하세요. {"video_duration": 6.07}<|video_aux_end|>\n\nWhat is shown in this video?\nassistant\n<think>\n globally신지차도 globally신지파괴의장 globally신지예원 sometimes globally신지신분 =================================================================================([@ globally신지과이 globally신지접이 globally신지 globally신지 globally신지 globally신지'
                ],
            }
        )
        expected = EXPECTED_TEXTS.get_expectation()

        self.assertEqual(output_text, expected)