"""
    Taken directly from http://scikit-image.org/docs/dev/auto_examples/plot_register_translation.html

"""

import numpy as np
import matplotlib.pyplot as plt

from skimage import data
from skimage.feature import register_translation
from skimage.feature.register_translation import _upsampled_dft
from scipy.ndimage import fourier_shift

def main():

    image = data.camera()
    shift = (-2.4, 1.32)
    # (-2.4, 1.32) pixel offset relative to reference coin
    offset_image = fourier_shift(np.fft.fftn(image), shift)
    offset_image = np.fft.ifftn(offset_image)
    print("Known offset (y, x):")
    print(shift)

    # pixel precision first
    shift, error, diffphase = register_translation(image, offset_image)

    fig = plt.figure(figsize=(8, 3))
    ax1 = plt.subplot(1, 3, 1, adjustable='box-forced')
    ax2 = plt.subplot(1, 3, 2, sharex=ax1, sharey=ax1, adjustable='box-forced')
    ax3 = plt.subplot(1, 3, 3)

    ax1.imshow(image)
    ax1.set_axis_off()
    ax1.set_title('Reference image')

    ax2.imshow(offset_image.real)
    ax2.set_axis_off()
    ax2.set_title('Offset image')

    # View the output of a cross-correlation to show what the algorithm is
    #    doing behind the scenes
    image_product = np.fft.fft2(image) * np.fft.fft2(offset_image).conj()
    cc_image = np.fft.fftshift(np.fft.ifft2(image_product))
    ax3.imshow(cc_image.real)
    ax3.set_axis_off()
    ax3.set_title("Cross-correlation")

    plt.show()

    print("Detected pixel offset (y, x):")
    print(shift)

    # subpixel precision
    shift, error, diffphase = register_translation(image, offset_image, 100)

    fig = plt.figure(figsize=(8, 3))
    ax1 = plt.subplot(1, 3, 1, adjustable='box-forced')
    ax2 = plt.subplot(1, 3, 2, sharex=ax1, sharey=ax1, adjustable='box-forced')
    ax3 = plt.subplot(1, 3, 3)

    ax1.imshow(image)
    ax1.set_axis_off()
    ax1.set_title('Reference image')

    ax2.imshow(offset_image.real)
    ax2.set_axis_off()
    ax2.set_title('Offset image')

    # Calculate the upsampled DFT, again to show what the algorithm is doing
    #    behind the scenes.  Constants correspond to calculated values in routine.
    #    See source code for details.
    cc_image = _upsampled_dft(image_product, 150, 100, (shift*100)+75).conj()
    ax3.imshow(cc_image.real)
    ax3.set_axis_off()
    ax3.set_title("Supersampled XC sub-area")


    plt.show()

    print("Detected subpixel offset (y, x):")
    print(shift)


if __name__ == '__main__':
    main()