import torch


def get_output_size(model, input_shape=(1, 3, 224, 224), device="cpu", dtype='float32'):
    """
    Returns the shape of the convolutional features in output to the model.

    Parameters
    ----------
    model pytorch model,
        neural network model.
    input_shape: tuple of int,
        shape of the images in input to the model in the form (batch_size, channels, height, width).
        Defaults to (1, 3, 224, 224).
    device: string,
        device for the gradient computation. Choose between "cpu" and "gpu:x", where x is the number of the GPU device.
    dtype: string,
        datatype for the model. Choose between 'float32' and 'float16'. Defaults to 'float32'.

    Return
    ------
    output_size : int,
        shape of the flattened convolutional features in output to the model.

    Note: It si not possible to do model(x) on CPU with f16. To avoid problems, the model is cast to float32 for this
    computation and then it is converted back to float16.
    """

    if dtype == "float16":
        model.float()

    dummy_input = torch.ones(input_shape).to(device)

    if model.name[0:12] == "efficientnet":
        output_size = model.extract_features(dummy_input).shape[1:].numel()
    else:
        output_size = model(dummy_input).shape[1:].numel()

    if dtype == "float16":
        model.half()

    return output_size