'''
file:       CorrectDrift.py
author:     Michael Entrup b. Epping (michael.entrup@wwu.de)
version:    20161017
info:       This module will correct the drift between two images.
'''

from __future__ import with_statement, division

import operator
#import pprint
import math

# pylint: disable-msg=E0401
from EFTEMj_pyLib import CrossCorrelation as cc
from EFTEMj_pyLib import PySIFT
from EFTEMj_pyLib import Tools as tools

from ij import IJ, ImagePlus
from ij.plugin import Duplicator
# pylint: enable-msg=E0401

def get_options():
    '''Returns a list of drift correction mode titles.'''
    return ['Scale-invariant feature transform', 'Normalized cross-correlation', 'None']

def get_modes():
    '''Return a list of modes that are available for drift correction.'''
    return ['SIFT', 'CC', 'None']

def correct_drift(img1, img2, display_cc=False):
    ''' Returns two ImagePlus objects that are drift corrected to each other.
    The first image is not changed.
    The second image is a new image that is shifted using ImageJ's Translate.
    :param img1: The reference image.
    :param img2: The image to be shifted.
    :param display_cc: Activate displaying the CrossCorrelation image (default False).
    '''
    img1_cc, img2_cc = cc.scale_to_power_of_two([img1, img2])
    result = cc.perform_correlation(img1_cc, img2_cc)
    x_off, y_off = cc.get_shift(result)
    # style after maximum detection
    if not display_cc:
        result.hide()
    else:
        result.copyScale(img1)
        cc.style_cc(result)
    if x_off == 0 and y_off == 0:
        print('No drift has been detected.')
        return img1, img2
    title = img2.getTitle()
    img2_dk = Duplicator().run(img2)