# Copyright 2026 IBM 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 GraniteMoeSWA model."""

import unittest

from transformers import is_torch_available
from transformers.testing_utils import (
    Expectations,
    require_torch,
    require_torch_accelerator,
    slow,
    torch_device,
)

from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester


if is_torch_available():
    import torch

    from transformers import AutoTokenizer, GraniteMoeSWAForCausalLM, GraniteMoeSWAModel


class GraniteMoeSWAModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = GraniteMoeSWAModel
        # With the default `num_hidden_layers=2`, `layer_types` resolves to
        # ["full_attention", "sliding_attention"], so both attention paths are exercised. The default
        # `sliding_window` stays larger than the short test sequences (matching gemma2/gpt_oss). Shared
        # experts stay disabled (`shared_intermediate_size=0`), matching the model's default.


@require_torch
class GraniteMoeSWAModelTest(CausalLMModelTest, unittest.TestCase):
    model_tester_class = GraniteMoeSWAModelTester


@slow
@require_torch_accelerator
class GraniteMoeSWAIntegrationTest(unittest.TestCase):
    model_id = "ibm-granite/granite-swash-3b-a600m"
    input_text = "The capital of France is"

    def test_model_logits_bf16(self):
        model = GraniteMoeSWAForCausalLM.from_pretrained(
            self.model_id, device_map="auto", dtype=torch.bfloat16, attn_implementation="eager"
        )
        tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        input_ids = tokenizer(self.input_text, return_tensors="pt").input_ids.to(torch_device)

        with torch.no_grad():
            out = model(input_ids)

        # fmt: off
        EXPECTED_MEANS = Expectations(
            {
                ("cuda", 8): torch.tensor([[-1.3672, -2.0156, -1.3359, -1.4531, -2.5156]]),
                ("cuda", (8, 6)): torch.tensor([[-1.3672, -2.0312, -1.3516, -1.4922, -2.5156]]),
                ("cuda", 9): torch.tensor([[-1.3594, -2.0156, -1.3594, -1.4922, -2.5156]]),
                ("xpu", None): torch.tensor([[-1.3672, -2.0312, -1.3516, -1.4766, -2.5156]]),
            }
        )
        EXPECTED_SLICES = Expectations(
            {
                ("cuda", 8): torch.tensor([0.4883, 4.3438, -0.0464, -0.2812, 1.8750, 1.3438, 2.3438, -0.7227, 3.5938, 1.9844, 1.4922, 3.7031, 1.7734, 3.5938, 2.8438]),
                ("cuda", (8, 6)): torch.tensor([ 0.4902,  4.4375, -0.0206, -0.2363,  1.8984,  1.3828,  2.4062, -0.7031, 3.7031,  2.0312,  1.5156,  3.7500,  1.8125,  3.6562,  2.8438]),
                ("cuda", 9): torch.tensor([0.4883, 4.4062, -0.0430, -0.2637, 1.8750, 1.3438, 2.3438, -0.7266, 3.6250, 2.0000, 1.5078, 3.7188, 1.8125, 3.5938, 2.8281]),
                ("xpu", None): torch.tensor([0.5352, 4.4375, -0.0038, -0.2559, 1.9219, 1.3594, 2.4375, -0.6758, 3.6719, 2.0312, 1.4844, 3.7969, 1.8359, 3.7188, 2.9062]),
            }
        )
        # fmt: on
        torch.testing.assert_close(
            EXPECTED_MEANS.get_expectation().to(torch_device), out.logits.mean(-1).float(), rtol=1e-2, atol=1e-2
        )
        torch.testing.assert_close(
            EXPECTED_SLICES.get_expectation().to(torch_device), out.logits[0, 0, :15].float(), rtol=1e-3, atol=1e-3
        )

    def test_model_generation(self):
        model = GraniteMoeSWAForCausalLM.from_pretrained(
            self.model_id, device_map="auto", dtype=torch.bfloat16, attn_implementation="eager"
        )
        tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        inputs = tokenizer(self.input_text, return_tensors="pt").to(torch_device)

        generated_ids = model.generate(**inputs, max_new_tokens=20, do_sample=False)
        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        EXPECTED_TEXTS = Expectations(
            {
                ("cuda", (8, 6)): (
                    "The capital of France is Paris.\nThe capital of France is Paris.\nThe capital of France is Paris.\nThe capital of France"
                ),
                ("xpu", None): (
                    "The capital of France is Paris.\nThe capital of France is also known as the City of Light.\nThe capital of France is"
                ),
            }
        )
        self.assertEqual(generated_text, EXPECTED_TEXTS.get_expectation())