import unittest

import torch
from torch import nn

from src.litespecformer.model import LiteSpecFormerModel
from src.litespecformer.config import LiteSpecFormerConfig
from src.litespecformer.output import (
    LiteSpecFormerOutput,
    LiteSpecFormerEncoderOutput,
    LiteSpecFormerEncoderBlockOutput,
    AttentionOutput,
)


class TestLiteSpecFormer(unittest.TestCase):
    """Test the Forward of LiteSpecFormer and the Pipline in Spectrum Prediction."""

    config = LiteSpecFormerConfig()

    # Set the batch_size
    batch_size = 4

    def test_model_init(self) -> None:
        """Test the model initialization in loading."""

        # Load the model
        model = LiteSpecFormerModel(config=self.config)

        # Check the model instance
        self.assertIsInstance(
            obj=model,
            cls=LiteSpecFormerModel,
            msg="Model is not initialized correctly.",
        )
        self.assertIsInstance(
            obj=model.config,
            cls=LiteSpecFormerConfig,
            msg="Model config is not initialized correctly.",
        )
        self.assertIsInstance(
            obj=model.shared,
            cls=nn.Embedding,
            msg="Model shared is not initialized correctly.",
        )

    def test_model_forward(self) -> None:
        """Test the model loading and the forward pass."""

        # Load the model
        model = LiteSpecFormerModel(config=self.config)

        # Create the inputs and outputs
        for seq_len in [36, 128, 300, 512]:
            # Test various sequence lengths
            context = torch.randn(self.batch_size, seq_len)
            future_target = torch.randn(self.batch_size, 16)

            # Forward through the model
            outputs = model(context=context, future_target=future_target)
            loss = outputs.loss
            quantile_preds = outputs.quantile_preds

            # Check the output batch size
            self.assertEqual(quantile_preds.shape[0], self.batch_size)

            # Check the patch len of prediction
            self.assertEqual(quantile_preds.shape[2], 16)

            # Check the loss
            self.assertIsInstance(obj=loss, cls=torch.Tensor)
