# Copyright 2020 The HuggingFace 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.
import unittest

from transformers import AutoTokenizer, is_torch_available
from transformers.testing_utils import (
    require_torch,
    slow,
)

from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester


if is_torch_available():
    import torch

    from transformers import (
        ModernBertDecoderForCausalLM,
        ModernBertDecoderForSequenceClassification,
        ModernBertDecoderModel,
    )


class ModernBertDecoderModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = ModernBertDecoderModel


@require_torch
class ModernBertDecoderModelTest(CausalLMModelTest, unittest.TestCase):
    model_tester_class = ModernBertDecoderModelTester


@slow
@require_torch
class ModernBertDecoderIntegrationTest(unittest.TestCase):
    def test_inference_causal_lm(self):
        model = ModernBertDecoderForCausalLM.from_pretrained("blab-jhu/test-32m-dec", attn_implementation="eager")
        tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")

        inputs = tokenizer("Paris is the capital of", return_tensors="pt")
        with torch.no_grad():
            output = model(**inputs)[0]
        expected_shape = torch.Size((1, 7, model.config.vocab_size))
        self.assertEqual(output.shape, expected_shape)

        # compare the actual values for a slice.
        expected_slice = torch.tensor(
            [[[-8.0183, -7.1578, -0.4453], [-6.2909, -6.1557, 4.9063], [-6.7689, -5.8068, 6.1078]]]
        )
        torch.testing.assert_close(output[:, :3, :3], expected_slice, rtol=1e-4, atol=1e-4)

    def test_inference_no_head(self):
        model = ModernBertDecoderModel.from_pretrained("blab-jhu/test-32m-dec", attn_implementation="eager")
        tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")

        inputs = tokenizer("Paris is the capital of", return_tensors="pt")
        with torch.no_grad():
            output = model(**inputs)[0]
        expected_shape = torch.Size((1, 7, model.config.hidden_size))
        self.assertEqual(output.shape, expected_shape)

        # compare the actual values for a slice.
        expected_slice = torch.tensor(
            [[[-0.0306, -0.0115, 0.0007], [-0.2485, -0.1381, 0.0872], [0.3133, -0.1777, 0.1667]]]
        )
        torch.testing.assert_close(output[:, :3, :3], expected_slice, rtol=1e-4, atol=1e-4)

    def test_generation(self):
        model = ModernBertDecoderForCausalLM.from_pretrained("blab-jhu/test-32m-dec", attn_implementation="eager")
        tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")

        inputs = tokenizer("The weather today is", return_tensors="pt")
        outputs = model.generate(**inputs, max_new_tokens=10, do_sample=False)
        output_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        # Check that we got some reasonable output
        self.assertEqual(len(output_text), 1)
        self.assertTrue(len(output_text[0]) > len("The weather today is"))

    def test_sliding_window_long_context(self):
        """
        Test that ModernBertDecoder works with sliding window attention for longer sequences.
        """

        model = ModernBertDecoderForCausalLM.from_pretrained("blab-jhu/test-32m-dec", attn_implementation="eager")
        tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")

        # Create a longer input to test sliding window attention
        long_input = "This is a test. " * 50  # Repeat to make it longer
        inputs = tokenizer(long_input, return_tensors="pt", truncation=True, max_length=512)

        outputs = model.generate(**inputs, max_new_tokens=20, do_sample=False)

        # Check that generation worked with longer context
        self.assertEqual(outputs.shape[0], 1)
        self.assertGreater(outputs.shape[1], inputs["input_ids"].shape[1])

    def test_sequence_classification(self):
        """
        Test that ModernBertDecoderForSequenceClassification works correctly.
        """

        model = ModernBertDecoderForSequenceClassification.from_pretrained(
            "blab-jhu/test-32m-dec", num_labels=2, attn_implementation="eager"
        )
        tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")

        # Test with sample input
        inputs = tokenizer("This is a positive example.", return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)

        # Check output shape
        expected_shape = (1, 2)  # batch_size=1, num_labels=2
        self.assertEqual(outputs.logits.shape, expected_shape)

        # Test with labels
        labels = torch.tensor([1])
        outputs_with_loss = model(**inputs, labels=labels)

        # Check that loss is computed
        self.assertIsNotNone(outputs_with_loss.loss)
        self.assertTrue(isinstance(outputs_with_loss.loss.item(), float))