import torch


def generic_fftshift(x,axis=[-2,-1],inverse=False):
    """
    Fourier shift to center the low frequency components

    Parameters
    ----------
    x : torch Tensor
        Input array
    inverse : bool
        whether the shift is for fft or ifft

    Returns
    -------
    shifted array

    """
    if len(axis) > len(x.shape):
        raise ValueError('Not enough axis to shift around!')
    
    y = x
    for axe in axis:
        dim_size = x.shape[axe]
        shift = int(dim_size/2)
        if inverse:
            if not dim_size%2 == 0:
                shift += 1
        
        y = torch.roll(y,shift,axe)
    
    return y