import re


def flops_layer(layer):
    """
    Calculate the number of flops for given a string information of layer.
    We extract only resonable numbers and use them.

    References:
        https://discuss.pytorch.org/t/calculating-flops-of-a-given-pytorch-model/3711

    Args:
        layer (str) : example
            Linear (512 -> 1000)
            Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
            BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
    """
    # print(layer)
    idx_type_end = layer.find('(')
    type_name = layer[:idx_type_end]

    params = re.findall('[^a-z](\d+)', layer)
    flops = 1

    if layer.find('Linear') >= 0:
        C1 = int(params[0])
        C2 = int(params[1])
        flops = C1 * C2

    elif layer.find('Conv2d') >= 0:
        C1 = int(params[0])
        C2 = int(params[1])
        K1 = int(params[2])
        K2 = int(params[3])

        # image size
        H = 32
        W = 32
        flops = C1 * C2 * K1 * K2 * H * W

    return flops