import torch
from torch.fft import fftn


def roll_quadrants(data, backwards=False):
    """
    Shift low frequencies to the center of fourier transform, i.e. [-N/2, ..., +N/2] -> [0, ..., N-1]
    Args:
        data: fourier transform, (NxHxW)
        backwards: bool, if True shift high frequencies back to center

    Returns:
    Shifted fourier transform.
    """
    dim = data.ndim - 1

    if dim != 2:
        raise AttributeError(f'Data must be 2d but it is {dim}d.')
    if any(s % 2 == 0 for s in data.shape[1:]):
        raise RuntimeWarning('Roll quadrants for 2d input should only be used with uneven spatial sizes.')

    # for each dimension swap left and right half
    dims = tuple(range(1, dim+1))          # add one for batch dimension
    shifts = torch.tensor(data.shape[1:]) // 2 #.div(2, rounding_mode='floor')                # N/2 if N even, (N-1)/2 if N odd
    if backwards:
        shifts *= -1
    return data.roll(shifts.tolist(), dims=dims)


def batch_fft(data, normalize=False):
    """
    Compute fourier transform of batch.
    Args:
        data: input tensor, (NxHxW)

    Returns:
    Batch fourier transform of input data.
    """

    dim = data.ndim - 1     # subtract one for batch dimension
    if dim != 2:
        raise AttributeError(f'Data must be 2d but it is {dim}d.')
        
    dims = tuple(range(1, dim + 1))  # add one for batch dimension
    if normalize:
        norm = 'ortho'
    else:
        norm = 'backward'

    if not torch.is_complex(data):
        data = torch.complex(data, torch.zeros_like(data))
    freq = fftn(data, dim=dims, norm=norm)

    return freq


def azimuthal_average(image, center=None):
    # modified to tensor inputs from https://www.astrobetter.com/blog/2010/03/03/fourier-transforms-of-images-in-python/
    """
    Calculate the azimuthally averaged radial profile.
    Requires low frequencies to be at the center of the image.
    Args:
        image: Batch of 2D images, NxHxW
        center: The [x,y] pixel coordinates used as the center. The default is
             None, which then uses the center of the image (including
             fracitonal pixels).

    Returns:
    Azimuthal average over the image around the center
    """
    # Check input shapes
    assert center is None or (len(center) == 2), f'Center has to be None or len(center)=2 ' \
                                                 f'(but it is len(center)={len(center)}.'
    # Calculate the indices from the image
    H, W = image.shape[-2:]
    h, w = torch.meshgrid(torch.arange(0, H), torch.arange(0, W))

    if center is None:
        center = torch.tensor([(w.max() - w.min()) / 2.0, (h.max() - h.min()) / 2.0])

    # Compute radius for each pixel wrt center
    r = torch.stack([w-center[0], h-center[1]]).norm(2, 0)

    # Get sorted radii
    r_sorted, ind = r.flatten().sort()
    i_sorted = image.flatten(-2, -1)[..., ind]

    # Get the integer part of the radii (bin size = 1)
    r_int = r_sorted.long()             # attribute to the smaller integer

    # Find all pixels that fall within each radial bin.
    deltar = r_int[1:] - r_int[:-1]  # Assumes all radii represented, computes bin change between subsequent radii
    rind = torch.where(deltar)[0]  # location of changed radius

    # compute number of elements in each bin
    nind = rind + 1         # number of elements = idx + 1
    nind = torch.cat([torch.tensor([0]), nind, torch.tensor([H*W])])        # add borders
    nr = nind[1:] - nind[:-1]  # number of radius bin, i.e. counter for bins belonging to each radius

    # Cumulative sum to figure out sums for each radius bin
    if H % 2 == 0:
        raise NotImplementedError('Not sure if implementation correct, please check')
        rind = torch.cat([torch.tensor([0]), rind, torch.tensor([H * W - 1])])  # add borders
    else:
        rind = torch.cat([rind, torch.tensor([H * W - 1])])  # add borders
    csim = i_sorted.cumsum(-1, dtype=torch.float64)             # integrate over all values with smaller radius
    tbin = csim[..., rind[1:]] - csim[..., rind[:-1]]
    # add mean
    tbin = torch.cat([csim[:, 0:1], tbin], 1)

    radial_prof = tbin / nr.to(tbin.device)         # normalize by counted bins

    return radial_prof


def get_spectrum(data, normalize=False):
    dim = data.ndim - 1  # subtract one for batch dimension
    if dim != 2:
        raise AttributeError(f'Data must be 2d but it is {dim}d.')

    freq = batch_fft(data, normalize=normalize)
    power_spec = freq.real ** 2 + freq.imag ** 2
    N = data.shape[1]
    if N % 2 == 0:      # duplicate value for N/2 so it is put at the end of the spectrum
                        # and is not averaged with the mean value
        N_2 = N//2
        power_spec = torch.cat([power_spec[:, :N_2+1], power_spec[:, N_2:N_2+1], power_spec[:, N_2+1:]], dim=1)
        power_spec = torch.cat([power_spec[:, :, :N_2+1], power_spec[:, :, N_2:N_2+1], power_spec[:, :, N_2+1:]], dim=2)

    power_spec = roll_quadrants(power_spec)
    power_spec = azimuthal_average(power_spec)
    return power_spec


def plot_std(mean, std, x=None, ax=None, **kwargs):
    import matplotlib.pyplot as plt
    if ax is None:
        fig, ax = plt.subplots(1)

    # plot error margins in same color as line
    err_kwargs = {
        'alpha': 0.3
    }

    if 'c' in kwargs.keys():
        err_kwargs['color'] = kwargs['c']
    elif 'color' in kwargs.keys():
        err_kwargs['color'] = kwargs['color']

    if x is None:
        x = torch.linspace(0, 1, len(mean))     # use normalized x axis
    ax.plot(x, mean, **kwargs)
    ax.fill_between(x, mean-std, mean+std, **err_kwargs)

    return ax