%matplotlib inline

import sys
# first you need to run:
# cd simpleitk && tar -zxvf SimpleITK-0.8.0-STRAW-branch-py2.7-linux-x86_64.egg_FILES.tar.gz
sys.path.insert(0,'simpleitk/SimpleITK-0.8.0-STRAW-branch-py2.7-linux-x86_64.egg_FILES')
sys.path.insert(0,'../src')

import numpy as np
import math as mt
import SimpleITK as sitk

import matplotlib.pyplot as plt

from imagedisplay import myshow
from IPython.display import display, clear_output

#load images
fixedInputFilename = "./data/BrainProtonDensitySliceBorder20.png"
fixedInput = sitk.ReadImage(fixedInputFilename)

#set fixed and moving to the same image
fixed = sitk.VectorIndexSelectionCast(fixedInput,0,sitk.sitkFloat32)
moving = fixed    

   
# the following defines a function that takes two images, an array with 1D translations,
# and a flag defining the (dis)similarity measure
def metric_calc_translation (fixed_img, moving_img, tx_array, metric='SSD'):
    
    # Define the registration object and properties 
    # (just to calculate the metric value - we are not going to perform optimization)
    R = sitk.ImageRegistrationMethod()    
    R.SetOptimizerAsRegularStepGradientDescent( maxStep=1e-10, 
                                                minStep=1e-11,
                                                numberOfIterations=1 )
    
    
    # Interpolation method
    # MAKE A SELECTION HERE
    
    # Linear interpolation
    R.SetInterpolator(sitk.sitkLinear)
    
    # Nearest neighbor interpolation
    #R.SetInterpolator(sitk.sitkNearestNeighbor)
    

    if metric == 'SSD':
        R.SetMetricAsMeanSquares()
        title_str = 'Sum of Squared Differences'
    elif metric == 'CC':
        R.SetMetricAsNormalizedCorrelation(subtractMean=True)
        title_str = 'Correlation Coefficient'
    elif metric == 'NMI':
        R.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
        title_str = 'Normalized Mutual Information'
    else:
        raise RuntimeError('undefined similarity metric')
        
    cost = []
    for tx in tx_array:
    
        # Create a transformation (translation in x-direction)
        t_input=sitk.Transform(fixed.GetDimension(), sitk.sitkTranslation)
        t_input.SetParameters( [tx,0] ) # t_input[0]=offset_x t_input[1]=offset_y

        # Run the cost function calculation
        R.SetTransform(t_input)
        R.Execute(fixed,moving)
        
        # Collect value
        cost.append(R.GetMetricValue())

    # Plot the metric values
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_title(title_str,fontsize=16)   

    ax.set_xlabel('Translation in x direction (in mm)',fontsize=16)
    ax.set_ylabel('Cost function',fontsize=16)
    ax.plot(tx_arr,cost,'r')
        
# CHANGE CODE HERE: Try different a different translation range
tx_arr = np.linspace(-15.0, 15.0, num=100, endpoint=True)

# Run the experiment
metric_calc_translation (fixed,moving,tx_arr,'SSD')
metric_calc_translation (fixed,moving,tx_arr,'CC')
metric_calc_translation (fixed,moving,tx_arr,'NMI')

tx_arr = np.linspace(-15.0, 15.0, num=100, endpoint=True)

# Apply Gaussian smoothing to the fixed and the moving image
moving = sitk.DiscreteGaussian(moving, 2.0)
fixed = sitk.DiscreteGaussian(fixed, 2.0)

# Re-run the little translation experiments
metric_calc_translation (fixed,moving,tx_arr,metric='SSD')
metric_calc_translation (fixed,moving,tx_arr,metric='CC')
metric_calc_translation (fixed,moving,tx_arr,metric='NMI')

# CHANGE CODE HERE: Try different values for the rotation
angle = 3.14/4


# Define display function
def cmyshow(img1,img2,*args, **kw):
    simg1 = sitk.Cast(sitk.RescaleIntensity(img1), sitk.sitkUInt8)
    simg2 = sitk.Cast(sitk.RescaleIntensity(img2), sitk.sitkUInt8)
    cimg = sitk.Compose(simg1, simg2, simg1/2.+simg2/2.)
    myshow(cimg,*args, **kw)

    
# Define the warping function
def warpmyimage(reference,img_source,transformation):
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(reference)
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetDefaultPixelValue(0.0)
    resampler.SetTransform(transformation)
    moving = resampler.Execute(img_source)
    return moving


# Read the image again
fixedInputFilename = "./data/BrainProtonDensitySliceBorder20.png"
fixedInput = sitk.ReadImage(fixedInputFilename)
fixed = sitk.VectorIndexSelectionCast(fixedInput,0,sitk.sitkFloat32)

# Smooth the input image
fixed = sitk.DiscreteGaussian(fixed, 1.0)

# Move the image center to origin - the rotation will be performed around the point (0,0)
spacing = fixed.GetSpacing();dimension = fixed.GetSize();  
fixed.SetOrigin([ -1*spacing[0]*dimension[0]/2,
                  -1*spacing[1]*dimension[1]/2 ])

# Transformation object for rotation: 
# http://www.itk.org/Doxygen/html/classitk_1_1Euler2DTransform.html
# Get the parameters that uniquely define the transform.
# These are typically used by optimizers. There are 3 parameters.
# The first one represents the angle or rotation in radians and the last two represents the translation. 
# The center of rotation is fixed.
t_input=sitk.Transform(fixed.GetDimension(), sitk.sitkEuler)
p = list(t_input.GetParameters()); p[0] = angle; t_input.SetParameters(p);

# Resample the moving image based on this transformation
moving = warpmyimage(fixed,fixed,t_input)
cmyshow(fixed,moving, title="Fixed Red, Moving Green", dpi = 50)

# Run the (dis)similarity evaluation for a range of translations
tx_arr = np.linspace(-15.0, 15.0, num=100, endpoint=True)
metric_calc_translation (fixed,moving,tx_arr,metric='SSD')
metric_calc_translation (fixed,moving,tx_arr,metric='CC')
metric_calc_translation (fixed,moving,tx_arr,metric='NMI')

# Setting up the registration
R = sitk.ImageRegistrationMethod()


# (Dis)similarity measure
# MAKE A SELECTION HERE

# SSD
R.SetMetricAsMeanSquares()

# CC
#R.SetMetricAsNormalizedCorrelation(subtractMean=True)

# NMI
#R.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)


R.SetOptimizerAsRegularStepGradientDescent(maxStep=2.0, minStep=.01, numberOfIterations=100 )
R.SetOptimizerScales([100.0, 1.0, 1.0])
R.SetTransform(sitk.Transform(fixed.GetDimension(), sitk.sitkEuler))
R.SetInterpolator(sitk.sitkLinear)

def command_iteration(method) :
    print("{0} = {1} : {2}".format( method.GetOptimizerIteration(),
                                    method.GetMetricValue(),
                                    method.GetOptimizerPosition() ) )
    sys.stdout.flush()

def command_end(R):
    print("-------")
    print("Optimizer stop condition: {0}".format(R.GetOptimizerStopConditionDescription()))
    print(" Iteration: {0}".format(R.GetOptimizerIteration()))
    print(" Metric value: {0}".format(R.GetMetricValue()))
    sys.stdout.flush()
    
R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda: command_iteration(R) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_end(R) )

t_output=R.Execute(fixed,moving)
warped_image = warpmyimage(fixed,moving,t_output)

cmyshow(fixed,moving,title="Before registration, Fixed Red, Moving Green", dpi = 50)
cmyshow(fixed,warped_image, title="After registration, Fixed Red, Moving Green", dpi = 50)

def command_plot_value_start(method,title):
    global x, ax, f
    x = []
    f, ax = plt.subplots()
    ax.set_title(title,fontsize=16)   
    ax.set_xlabel('Optimization Iteration Number',fontsize=16)
    ax.set_ylabel('Metric Value',fontsize=16)

def command_plot_value_end(method):
    global x, ax
    del x
    del ax
    plt.close()

def command_plot_value(method):
    global x, ax, f
    v = method.GetMetricValue()
    x.append(v)
    if( len(ax.lines) > 0 ):
        ax.lines.pop(0)
    ax.plot(x,'r')
    display(f)
    
titleSSD='SSD Registration -GradientDescent Optimizer'
R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda:  clear_output() )
R.AddCommand( sitk.sitkStartEvent, lambda: command_plot_value_start(R,titleSSD) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_plot_value_end(R) )
R.AddCommand( sitk.sitkIterationEvent, lambda: command_plot_value(R) )

R.SetTransform(sitk.Transform(fixed.GetDimension(), sitk.sitkEuler))

R.Execute(fixed,moving);

titleNelderMead='SSD Registration - Nelder-Mead Simplex'
R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda:  clear_output() )
R.AddCommand( sitk.sitkStartEvent, lambda: command_plot_value_start(R,titleNelderMead) )
R.AddCommand( sitk.sitkIterationEvent, lambda: command_plot_value(R) )

R.SetTransform(sitk.Transform(fixed.GetDimension(), sitk.sitkEuler))
R.SetOptimizerAsAmoeba(simplexDelta=3,parametersConvergenceTolerance=0.01,functionConvergenceTolerance=0.01); 

R.Execute(fixed,moving);