from typing import Tuple
from typing import Sequence
import math


def _conv2d_output_size(
    input_dimensions: Tuple[int, int],
    output_channels: Sequence[int],
    kernel_sizes: Sequence[int],
    stride_sizes: Sequence[int],
) -> int:
    """Calculates the output size of convolutional layers based on square kernel sizes
    and square stride sizes. This does not take into account padding or dilation.

        H_out = \floor((H_in - kernel_size) / stride_size + 1)
        W_out = \floor((W_in - kernel_size) / stride_size + 1)

    Args:
        input_dimensions (Tuple[int, int]): The input dimension in the form of
            (height, width).
        output_channels (Sequence[int]): Number of output channels for each layer of
            convolutions.
        kernel_sizes (Sequence[int]): The kernel size for each layer of
            convolutions.
        stride_sizes (Sequence[int]): The stride size for each layer of
            convolutions.

    Returns:
        int: The output size of the convolutional layers.
    """
    assert len(output_channels) == len(kernel_sizes) == len(stride_sizes)
    assert len(input_dimensions) == 2

    input_height = input_dimensions[0]
    input_width = input_dimensions[1]

    for kernel_size, stride_size in zip(kernel_sizes, stride_sizes):
        input_height = math.floor((input_height - kernel_size) / stride_size + 1)
        input_width = math.floor((input_width - kernel_size) / stride_size + 1)

    return input_height * input_width * output_channels[-1]