import random


def time_mask(spec, T=40, num_masks=1, replace_with_zero=False):
    """Time masking

    :param torch.Tensor spec: input tensor with shape (T, dim)
    :param int T: maximum width of each mask
    :param int num_masks: number of masks
    :param bool replace_with_zero: if True, masked parts will be filled with 0, if False, filled with mean
    """
    cloned = spec.unsqueeze(0).clone()
    len_spectro = cloned.shape[1]

    for i in range(0, num_masks):
        t = random.randrange(0, T)
        t_zero = random.randrange(0, len_spectro - t)

        # avoids randrange error if values are equal and range is empty
        if (t_zero == t_zero + t):
            return cloned.squeeze(0)

        mask_end = random.randrange(t_zero, t_zero + t)
        if (replace_with_zero):
            cloned[0][t_zero:mask_end, :] = 0
        else:
            cloned[0][t_zero:mask_end, :] = cloned.mean()
    return cloned.squeeze(0)