import unittest
import torch
from model_architectures import (
    ConvolutionalProcessingBlockBN,
    ConvolutionalDimensionalityReductionBlockBN,
    ConvolutionalProcessingBlockBNRC,
)


class TestBatchNormalizationBlocks(unittest.TestCase):
    def setUp(self):
        # Common parameters
        self.input_shape = (1, 3, 32, 32)  # Batch size 1, 3 channels, 32x32 input
        self.num_filters = 16
        self.kernel_size = 3
        self.padding = 1
        self.bias = False
        self.dilation = 1
        self.reduction_factor = 2

    def test_convolutional_processing_block(self):
        # Create a ConvolutionalProcessingBlockBN instance
        block = ConvolutionalProcessingBlockBN(
            input_shape=self.input_shape,
            num_filters=self.num_filters,
            kernel_size=self.kernel_size,
            padding=self.padding,
            bias=self.bias,
            dilation=self.dilation,
        )

        # Generate a random tensor matching the input shape
        input_tensor = torch.randn(self.input_shape)

        # Forward pass
        try:
            output = block(input_tensor)
            self.assertIsNotNone(output, "Output should not be None.")
        except Exception as e:
            self.fail(f"ConvolutionalProcessingBlock raised an error: {e}")

    def test_convolutional_processing_block_with_rc(self):
        # Create a ConvolutionalProcessingBlockBNRC instance
        block = ConvolutionalProcessingBlockBNRC(
            input_shape=self.input_shape,
            num_filters=self.num_filters,
            kernel_size=self.kernel_size,
            padding=self.padding,
            bias=self.bias,
            dilation=self.dilation,
        )

        # Generate a random tensor matching the input shape
        input_tensor = torch.randn(self.input_shape)

        # Forward pass
        try:
            output = block(input_tensor)
            self.assertIsNotNone(output, "Output should not be None.")
        except Exception as e:
            self.fail(f"ConvolutionalProcessingBlock raised an error: {e}")

    def test_convolutional_dimensionality_reduction_block(self):
        # Create a ConvolutionalDimensionalityReductionBlockBN instance
        block = ConvolutionalDimensionalityReductionBlockBN(
            input_shape=self.input_shape,
            num_filters=self.num_filters,
            kernel_size=self.kernel_size,
            padding=self.padding,
            bias=self.bias,
            dilation=self.dilation,
            reduction_factor=self.reduction_factor,
        )

        # Generate a random tensor matching the input shape
        input_tensor = torch.randn(self.input_shape)

        # Forward pass
        try:
            output = block(input_tensor)
            self.assertIsNotNone(output, "Output should not be None.")
        except Exception as e:
            self.fail(f"ConvolutionalDimensionalityReductionBlock raised an error: {e}")


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