# -*- coding: utf-8 -*-
# Copyright 2007-2011 The Hyperspy developers
#
# This file is part of  Hyperspy.
#
#  Hyperspy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
#  Hyperspy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with  Hyperspy.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import scipy as sp
import scipy.signal
from scipy.fftpack import fftn, ifftn
import scipy.ndimage
import matplotlib.pyplot as plt


def shift_image(im, shift, fill_value=np.nan):
    fractional, integral = np.modf(shift)
    if fractional.any():
        order = 3
    else:
        # Disable interpolation
        order = 0
    im[:] = sp.ndimage.shift(im, shift, cval=fill_value, order=order)


def triu_indices_minus_diag(n):
    """Returns the indices for the upper-triangle of an (n, n) array 
    excluding its diagonal
    
    Parameters
    ----------
    n : int
        The length of the square array
        
    """
    ti = np.triu_indices(n)
    isnotdiag = ti[0] != ti[1]
    return (ti[0][isnotdiag], ti[1][isnotdiag])

def hanning2d(M, N):
    """
    A 2D hanning window created by outer product.
    """
    return np.outer(np.hanning(M),np.hanning(N))


def sobel_filter(im):
    sx = sp.ndimage.sobel(im, axis=0, mode='constant')
    sy = sp.ndimage.sobel(im, axis=1, mode='constant')
    sob = np.hypot(sx, sy)
    return sob
    
def fft_correlation(in1, in2, normalize=False):
    """Correlation of two N-dimensional arrays using FFT.
    
    Adapted from scipy's fftconvolve.
    
    Parameters
    ----------
    in1, in2 : array
    normalize: bool
        If True performs phase correlation

    """
    s1 = np.array(in1.shape)
    s2 = np.array(in2.shape)
    size = s1 + s2 - 1
    # Use 2**n-sized FFT
    fsize = 2 ** np.ceil(np.log2(size))
    IN1 = fftn(in1, fsize)
    IN1 *= fftn(in2, fsize).conjugate()
    fslice = tuple([slice(0, int(sz)) for sz in size])
    if normalize is True:
        ret = ifftn(np.nan_to_num(IN1/np.absolute(IN1))).real.copy()
    else:
        ret = ifftn(IN1).real.copy()
    del IN1
    return ret

def estimate_image_shift(ref, image, roi=None, sobel=True,
                         medfilter=True, hanning=True, plot=False,
                         dtype='float', normalize_corr=False,):
    """Estimate the shift in a image using phase correlation

    This method can only estimate the shift by comparing 
    bidimensional features that should not change the position 
    in the given axis. To decrease the memory usage, the time of 
    computation and the accuracy of the results it is convenient 
    to select a region of interest by setting the roi keyword.

    Parameters
    ----------
    
    roi : tuple of ints (top, bottom, left, right)
         Define the region of interest
    sobel : bool
        apply a sobel filter for edge enhancement 
    medfilter :  bool
        apply a median filter for noise reduction
    hanning : bool
        Apply a 2d hanning filter
    plot : bool
        If True plots the images after applying the filters and
        the phase correlation
    reference : \'current\' | \'cascade\'
        If \'current\' (default) the image at the current 
        coordinates is taken as reference. If \'cascade\' each image
        is aligned with the previous one.
    dtype : str or dtype
        Typecode or data-type in which the calculations must be
        performed.
        
    normalize_corr : bool
        If True use phase correlation instead of standard correlation
        
    Returns
    -------
    
    shifts: np.array
        containing the estimate shifts
    max_value : float
        The maximum value of the correlation
    
    """
    
    # Make a copy of the image so the don't get modified
    ref = ref.copy().astype(dtype)
    image = image.copy().astype(dtype)
    if roi is not None:
            top, bottom, left, right = roi
    else:
        top, bottom, left, right = [None,] * 4
        
    # Select region of interest
    ref = ref[top:bottom, left:right]
    image = image[top:bottom, left:right]
    
    # Apply filters
    for im in (ref,image):
        if hanning is True:
            im *= hanning2d(*im.shape) 
        if medfilter is True:
            im[:] = sp.signal.medfilt(im)
        if sobel is True:
            im[:] = sobel_filter(im)
    
    
    phase_correlation = fft_correlation(ref, image,
        normalize=normalize_corr)
    
    # Estimate the shift by getting the coordinates of the maximum
    argmax = np.unravel_index(np.argmax(phase_correlation),
                              phase_correlation.shape)
    threshold = (phase_correlation.shape[0]/2 - 1,
                phase_correlation.shape[1]/2 - 1)
    shift0 = argmax[0] if argmax[0] < threshold[0] else  \
        argmax[0] - phase_correlation.shape[0] 
    shift1 = argmax[1] if argmax[1] < threshold[1] else \
        argmax[1] - phase_correlation.shape[1]
    max_val = phase_correlation.max()
    
    # Plot on demand
    if plot is True:
        f, axarr = plt.subplots(1,3)
        axarr[0].imshow(ref)
        axarr[1].imshow(image)
        axarr[2].imshow(phase_correlation)
        axarr[0].set_title('Reference')
        axarr[1].set_title('Image')
        axarr[2].set_title('Phase correlation')
        plt.show()
    # Liberate the memory. It is specially necessary if it is a 
    # memory map
    del ref
    del image
    
    return -np.array((shift0, shift1)), max_val 
