import unittest
import jax  # type: ignore
import jax.numpy as jnp  # type: ignore
from jax import random  # type: ignore
from flax import linen as nn  # type: ignore

from capibara_model.core.config import CapibaraConfig
from capibara_model.core.model import CapibaraENT
from capibara_model.utils.data_processing import preprocess_text
from capibara_model.utils.generate_response import generate_response
from capibara_model.utils.language_utils import detect_language, translate_text


class TestCapibaraENT(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        """Set up the test environment for all test methods in the class."""
        cls.rng = random.PRNGKey(0)

    def setUp(self):
        """Sets up the test environment for each test method."""
        self.config = CapibaraConfig(
            d_model=512,
            d_state=256,
            d_conv=128,
            expand=2,
            base_model_name='gpt2',
            translation_model='facebook/m2m100_418M',
            get_active_layers=lambda: ['platonic', 'game_theory', 'ethics'],
            get_layer_config=lambda layer_name: {},
            personality={},
            context_window_size=10,
            max_length=50,
            vocab_size=1000
        )
        self.model = CapibaraENT(self.config)

    def test_model_creation(self):
        """Tests the creation of the CapibaraENT model."""
        self.assertIsInstance(self.model, CapibaraENT)

    def test_forward_pass(self):
        """Tests a forward pass through the model."""
        rng, input_rng = random.split(self.rng)
        input_ids = random.randint(
            input_rng, (1, self.config.max_length), 0, self.config.vocab_size)

        variables = self.model.init(rng, input_ids)
        output = self.model.apply(variables, input_ids)

        self.assertEqual(
            output['output'].shape, (1, self.config.max_length, self.config.vocab_size))

    def test_generate_response(self):
        """Tests the response generation."""
        rng, input_rng = random.split(self.rng)
        variables = self.model.init(rng, jnp.ones(
            (1, self.config.max_length), dtype=jnp.int32))

        user_input = "Hello, how are you?"
        conversation_history = ["User: Hi",
                                "AI: Hello! How can I assist you today?"]

        response = self.model.apply(
            variables, user_input, conversation_history, method=self.model.generate_response)

        self.assertIsInstance(response, str)
        self.assertGreater(len(response), 0)

    def test_gelu(self):
        """Tests the GELU activation function."""
        x = jnp.array([-2.0, -1.0, 0.0, 1.0, 2.0])
        output = gelu(x)
        expected = jnp.array(
            [-0.04550026, -0.15865529, 0.0, 0.8413447, 1.9544997])
        self.assertTrue(jnp.allclose(output, expected, atol=1e-6))

    def test_swish(self):
        """Tests the Swish activation function."""
        x = jnp.array([-2.0, -1.0, 0.0, 1.0, 2.0])
        output = swish(x)
        expected = jnp.array(
            [-0.23840584, -0.26894143, 0.0, 0.7310586, 1.7615942])
        self.assertTrue(jnp.allclose(output, expected, atol=1e-6))

    def test_positional_encoding(self):
        """Tests the positional encoding."""
        max_len, d_model = 100, 512
        pe = PositionalEncoding(d_model, max_len=max_len)
        encoding = pe(jnp.zeros((1, max_len, d_model)))
        self.assertEqual(encoding.shape, (1, max_len, d_model))
        self.assertTrue(jnp.all(encoding >= -1) and jnp.all(encoding <= 1))

    def test_create_masks(self):
        """Tests the creation of masks."""
        seq_len, batch_size = 10, 2
        src_mask, tgt_mask = create_masks(seq_len, batch_size)
        self.assertEqual(src_mask.shape, (batch_size, 1, seq_len))
        self.assertEqual(tgt_mask.shape, (batch_size, seq_len, seq_len))

    def test_detect_language(self):
        """Tests the language detection function."""
        text = "Hello, how are you?"
        lang = detect_language(text)
        self.assertEqual(lang, "en")

    def test_translate_text(self):
        """Tests the text translation function."""
        text = "Hello, how are you?"
        translated = translate_text(text, source_lang="en", target_lang="es")
        self.assertIsInstance(translated, str)
        self.assertNotEqual(text, translated)


if __name__ == '__main__':
    unittest.main()
