from typing import Tuple
from typing import Union


def multiple_cnn_size(
    input: Tuple[int, int],
    kernel: Union[int, Tuple[int, int]],
    padding: Union[int, Tuple[int, int]] = 0,
    stride: Union[int, Tuple[int, int]] = 1,
    n_convs: int = 1,
) -> Tuple[int, int]:
    """
    Return the size of the output of more than one convolutional layer. This function can replace cnn_size.
    Args:
        input: Size of the input image.
        kernel: Kernel size, it is assumed to be a square.
        padding: Padding size.
        stride: Stride.
        n_convs: Number of convolutions.
    Returns:
        The output size.
    """
    if isinstance(kernel, int):
        kernel = (kernel, kernel)

    if isinstance(padding, int):
        padding = (padding, padding)

    if isinstance(stride, int):
        stride = (stride, stride)

    # Function that computes our function.
    def f(i, k, p, s):
        count = 0
        for j in range(0, n_convs):
            s_j = s ** j
            count += -s_j * k + 2 * s_j * p + s_j * s
        count += i
        return count / (s ** n_convs)

    out_w = f(i=input[0], k=kernel[0], p=padding[0], s=stride[0])
    out_h = f(i=input[1], k=kernel[1], p=padding[1], s=stride[1])
    return int(out_w), int(out_h)