# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2016-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

from packaging.version import parse as version_parse
import numpy as np
from scipy.ndimage import fourier_shift

from scipy.fft import ifftshift, fftshift, rfftn, rfftfreq, fftfreq, fftn, ifftn, irfftn
from scipy.special import erfc
# from skimage.registration import phase_cross_correlation as phase_cross_correlation_orig
# from skimage import __version__ as skimage_version
from ._phase_cross_correlation import phase_cross_correlation as phase_cross_correlation_orig


def shift(img, v):
    """
    Shift an image 2D or 3D, if necessary using a subpixel (FFT-based) approach. Values are wrapped around the array borders.

    Alternative methods:
        for real data: see scipy.ndimage.interpolation.shift (uses spline interpolation, complex not supported)
        for complex or real data: scipy.ndimage.fourier_shift
    
    Args:
        img: the image (2D or 3D) to be shifted, which will be treated as complex-valued data.
        v: the shift values (2 or 3 values according to array dimensions)

    Returns:
        the shifted array
    """
    assert (img.ndim == len(v))
    if all([(x % 1 == 0) for x in v]):
        # integer shifting
        simg = np.roll(img, v[-1], axis=-1)
        for i in range(img.ndim - 1):
            simg = np.roll(simg, v[i], axis=i)
        return simg
    else:
        if np.isrealobj(img):
            # use r2c transform
            return irfftn(fourier_shift(rfftn(img), v, n=img.shape[-1]))
        return ifftn(fourier_shift(fftn(img), v))


def phase_cross_correlation(reference_image, moving_image, low_cutoff=None, low_width=0.03,
                            high_cutoff=None, high_width=0.03, **kwargs):
    """
    phase cross-correlation from scikit-image. This version automatically adds
    the normalization keyword for skimage>=0.19, and adds low and/or high-bandpass
    filters.

    :param reference_image: the reference image
    :param moving_image: the moving image
    :param low_cutoff: a 0<value<<0.5 can be given (typically it should be a few 0.01),
        an erfc filter with a cutoff at low_cutoff*N (where N is the size along each dimension)
        will be applied, after the images have been FT'd
    :param low_width: the width of the low cutoff filter, also as a percentage of the size
    :param high_cutoff: same as low_cutoff fot the high frequency filter, should be close below 0.5
    :param high_width: same as low_width
    :param kwargs: keyword arguments which will be passed to skimage's phase_cross_correlation
    :return:
    """
    # if version_parse(skimage_version) >= version_parse('0.19'):
    if 'normalization' not in kwargs:
        kwargs['normalization'] = None

    sh = reference_image.shape

    assert np.allclose(sh, moving_image.shape)

    if low_cutoff is not None or high_cutoff is not None:
        need_ft = True
        r2c = kwargs['r2c'] if 'r2c' in kwargs else False
        if 'space' in kwargs:
            if 'fourier' == kwargs['space']:
                need_ft = False
        if need_ft:
            if np.isrealobj(reference_image) and np.isrealobj(moving_image):
                reference_image = rfftn(reference_image)
                moving_image = rfftn(moving_image)
                r2c = True
                kwargs['r2c'] = True
            else:
                reference_image = fftn(reference_image)
                moving_image = fftn(moving_image)
            kwargs['space'] = 'fourier'
        r = 0
        for i in range(len(sh)):
            if r2c and i == len(sh) - 1:
                if need_ft:
                    tmp = rfftfreq(sh[i]).astype(np.float32)
                else:
                    # space == 'fourier' and r2c
                    tmp = rfftfreq(2 * (sh[i] - 1)).astype(np.float32)
            else:
                tmp = fftfreq(sh[i]).astype(np.float32)
            for ii in range(len(sh) - i - 1):
                tmp = np.expand_dims(tmp, axis=1)
            r = r + tmp ** 2
        r = np.sqrt(r)
        if low_cutoff:
            tmp = 1 - 0.5 * erfc((r - low_cutoff) / low_width)
            reference_image *= tmp
            moving_image *= tmp
        if high_cutoff:
            tmp = 0.5 * erfc((r - high_cutoff) / high_width)
            reference_image *= tmp
            moving_image *= tmp
    return phase_cross_correlation_orig(reference_image, moving_image, **kwargs)


def phase_cross_correlation_paraboloid(reference_image, moving_image, low_cutoff=None, low_width=0.03,
                                       high_cutoff=None, high_width=0.03):
    """
    Image registration of 2D images using phase cross-correlation from scikit-image, providing
    sub-pixel accuracy through the paraboloid fit of the cross-correlation map.
    If a stack of 2D images is provided (instead of just 1), the shifts are returned for
    all images.

    :param reference_image: the reference image
    :param moving_image: the moving image
    :param low_cutoff: a 0<value<<0.5 can be given (typically it should be a few 0.01),
        an erfc filter with a cutoff at low_cutoff*N (where N is the size along each dimension)
        will be applied, after the images have been FT'd
    :param low_width: the width of the low cutoff filter, also as a percentage of the size
    :param high_cutoff: same as low_cutoff fot the high frequency filter, should be close below 0.5
    :param high_width: same as low_width
    :return: a tuple(cy, cx) of shifts for 2D images, or two arrays (cy, cx) for the pixel
        shifts, each having the same shape as the original arrays, minus the last two dimensions.
    """

    assert np.allclose(reference_image.shape, moving_image.shape)
    # Make sure we have a stack of 2D images
    sh0 = reference_image.shape
    if len(sh0) > 2:
        nz = np.prod(sh0[:-2])
        ny, nx = sh0[-2:]
        sh = sh0[-2:]
    else:
        ny, nx = sh0
        sh = sh0
        nz = 1
    reference_image = reference_image.reshape((nz, ny, nx))
    moving_image = moving_image.reshape((nz, ny, nx))

    r2c = np.isrealobj(reference_image)
    if r2c:
        reference_image = rfftn(reference_image.astype(np.float32), axes=(-2, -1))
        moving_image = rfftn(moving_image.astype(np.float32), axes=(-2, -1))
    else:
        reference_image = fftn(reference_image.astype(np.complex64), axes=(-2, -1))
        moving_image = fftn(moving_image.astype(np.complex64), axes=(-2, -1))

    if low_cutoff is not None or high_cutoff is not None:
        r = 0
        for i in range(len(sh)):
            if r2c and i == len(sh) - 1:
                tmp = rfftfreq(sh[i]).astype(np.float32)
            else:
                tmp = fftfreq(sh[i]).astype(np.float32)
            for ii in range(len(sh) - i - 1):
                tmp = np.expand_dims(tmp, axis=1)
            r = r + tmp ** 2
        r = np.sqrt(r)
        if low_cutoff:
            tmp = 1 - 0.5 * erfc((r - low_cutoff) / low_width)
            reference_image *= tmp
            moving_image *= tmp
        if high_cutoff:
            tmp = 0.5 * erfc((r - high_cutoff) / high_width)
            reference_image *= tmp
            moving_image *= tmp
    shu = list(moving_image.shape[:-1])
    shu[-1] *= sh[-1]  # shape with the registration axes collapsed
    if r2c:
        cc = irfftn(reference_image * moving_image.conj(), axes=(-2, -1))
        maxima = np.unravel_index(np.argmax(cc.reshape(shu), axis=-1), sh)
    else:
        cc = ifftn(reference_image * moving_image.conj(), axes=(-2, -1))
        maxima = np.unravel_index(np.argmax(np.abs(cc).reshape(shu), axis=-1), sh)

    midpoints = np.array([np.fix(axis_size / 2) for axis_size in sh])
    iy, ix = np.fix(maxima).astype(np.int32)
    if np.isscalar(ix):
        if iy > midpoints[0]:
            iy -= sh[0]
        if ix > midpoints[1]:
            ix -= sh[1]
    else:
        iy[iy > midpoints[0]] -= sh[0]
        ix[ix > midpoints[1]] -= sh[1]

    # Paraboloid fit
    ixm = (ix - 1) % nx
    ixp = (ix + 1) % nx
    iym = (iy - 1) % ny
    iyp = (iy + 1) % ny

    # Paraboloid fit
    # CC(x,y)=Ax^2 + By^2 + Cxy + Dx + Ey + F

    # First reshape so the ccmax equations work with array of coordinates
    shz = list(cc.shape)
    if len(shz) == 2:
        shz = [1] + shz
    vz = list(range(shz[0]))
    ccz = cc.reshape(shz)

    vf = ccz[vz, iy, ix]
    va = (ccz[vz, iy, ixp] + ccz[vz, iy, ixm]) / 2 - vf
    vb = (ccz[vz, iyp, ix] + ccz[vz, iym, ix]) / 2 - vf
    vd = (ccz[vz, iy, ixp] - ccz[vz, iy, ixm]) / 2
    ve = (ccz[vz, iyp, ix] - ccz[vz, iym, ix]) / 2
    vc = (ccz[vz, iyp, ixp] - ccz[vz, iym, ixp] - ccz[vz, iyp, ixm] + ccz[vz, iym, ixm]) / 4

    dx = ix + np.maximum(np.minimum((2 * vb * vd - vc * ve) / (vc * vc - 4 * va * vb), 0.5), -0.5)
    dy = iy + np.maximum(np.minimum((2 * va * ve - vc * vd) / (vc * vc - 4 * va * vb), 0.5), -0.5)
    cx = dx - nx * (dx > (nx / 2)) + nx * (dx < (-nx / 2))
    cy = dy - ny * (dy > (ny / 2)) + ny * (dy < (-ny / 2))
    if nz == 1:
        return cy[0], cx[0]
    return cy.reshape(sh0[:-2]), cx.reshape(sh0[:-2])
