import unittest

import torch

from catasta.models import (
    FeedforwardRegressor,
    CNNImageClassifier,
    TransformerRegressor,
    TransformerFFTRegressor,
    MambaRegressor,
    MambaFFTRegressor,
)


class ModelContractTests(unittest.TestCase):
    def test_feedforward_regressor_no_hidden_layers_respects_output_dims(self):
        x = torch.randn(5, 4)

        single = FeedforwardRegressor(n_inputs=4, n_outputs=1, dropout=0.0, hidden_dims=[])
        multi = FeedforwardRegressor(n_inputs=4, n_outputs=3, dropout=0.0, hidden_dims=[])

        self.assertEqual(single(x).shape, (5,))
        self.assertEqual(multi(x).shape, (5, 3))

    def test_cnn_image_classifier_supports_empty_feedforward_dims(self):
        model = CNNImageClassifier(
            input_shape=(8, 8, 3),
            n_classes=4,
            conv_out_channels=[4],
            feedforward_dims=[],
        )
        x = torch.randn(2, 8, 8, 3)
        y = model(x)
        self.assertEqual(y.shape, (2, 4))

    def test_patch_regressor_guards(self):
        regressor_builders = [
            lambda: TransformerRegressor(
                n_inputs=8,
                n_outputs=1,
                n_patches=0,
                d_model=4,
                n_layers=1,
                n_heads=1,
                feedforward_dim=8,
                head_dim=4,
            ),
            lambda: TransformerFFTRegressor(
                n_inputs=8,
                n_outputs=1,
                n_patches=9,
                d_model=4,
                n_layers=1,
                n_heads=1,
                feedforward_dim=8,
                head_dim=4,
            ),
            lambda: MambaRegressor(
                n_inputs=8,
                n_outputs=1,
                n_patches=3,
                d_model=4,
                d_state=2,
                d_conv=2,
                expand=1,
                n_layers=1,
            ),
            lambda: MambaFFTRegressor(
                n_inputs=8,
                n_outputs=1,
                n_patches=3,
                d_model=4,
                d_state=2,
                d_conv=2,
                expand=1,
                n_layers=1,
            ),
        ]

        for builder in regressor_builders:
            with self.assertRaises(ValueError):
                builder()

    def test_regressor_squeeze_behavior_for_multi_output(self):
        x = torch.randn(3, 8)

        models = [
            TransformerRegressor(
                n_inputs=8,
                n_outputs=2,
                n_patches=2,
                d_model=4,
                n_layers=1,
                n_heads=1,
                feedforward_dim=8,
                head_dim=4,
            ),
            TransformerFFTRegressor(
                n_inputs=8,
                n_outputs=2,
                n_patches=2,
                d_model=4,
                n_layers=1,
                n_heads=1,
                feedforward_dim=8,
                head_dim=4,
            ),
            MambaRegressor(
                n_inputs=8,
                n_outputs=2,
                n_patches=2,
                d_model=4,
                d_state=2,
                d_conv=2,
                expand=1,
                n_layers=1,
            ),
            MambaFFTRegressor(
                n_inputs=8,
                n_outputs=2,
                n_patches=2,
                d_model=4,
                d_state=2,
                d_conv=2,
                expand=1,
                n_layers=1,
            ),
        ]

        for model in models:
            self.assertEqual(model(x).shape, (3, 2))


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