%matplotlib inline

import sys
sys.path.insert(0,'simpleitk/SimpleITK-0.8.0-STRAW-branch-py2.7-linux-x86_64.egg_FILES')
sys.path.insert(0,'../src')

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

import matplotlib.pyplot as plt

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

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

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)

# Multi Modal Image Registration
fixedInputFilename = "./data/ct.nii.gz"
movingInputFilename = "./data/mr.nii.gz"
fixed = sitk.ReadImage(fixedInputFilename)
moving = sitk.ReadImage(movingInputFilename)

# Define a transformation and warp the moving image (this time transformation is in 3D)
# Resample the moving image based on this transformation

# CHANGE CODE HERE: Try different initial transformation parameters
angleZ = -3.14*40/180; tx = -35.00;  ty = 41.00;

t_input=sitk.Transform(fixed.GetDimension(), sitk.sitkEuler)
t_input.SetParameters([0,0,angleZ,tx,ty,0])
moving = warpmyimage(fixed,moving,t_input)

global t_output
t_output = sitk.Transform(fixed.GetDimension(), sitk.sitkEuler)

# Multi-scale registration 
levels = [8,4]
for level in levels:

    fixed_s = sitk.DiscreteGaussian(fixed, 2.0 * level )
    moving_s = sitk.DiscreteGaussian(moving, 2.0 * level )
    
    fixed_s = sitk.Shrink(fixed_s, [1,1] * level )
    moving_s = sitk.Shrink(moving_s, [1,1] * level )
        
    # Create the registration object and set its properties
    R = sitk.ImageRegistrationMethod()
    R.SetOptimizerAsRegularStepGradientDescent(maxStep=4.0, minStep=.05, numberOfIterations=400 )
    R.SetOptimizerScales([1e10,1e10,100,1,1,1e10])
    R.SetInterpolator(sitk.sitkLinear)
    R.SetTransform(t_output)
    R.SetNumberOfThreads(8)
    
    # Define the (dis)similarity measure, sum of squared differences
    R.SetMetricAsMeanSquares()
    
    # To observe the cost function over iterations
    def command_iteration(method) :
        if ((method.GetOptimizerIteration() % 100) == 0):
            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) )

    # Run the registration and warp moving image
    t_output = R.Execute(fixed_s,moving_s)

warped_image = warpmyimage(fixed,moving,t_output)
cmyshow(fixed,moving, title="Fixed Red, Moving Green - Before Registration", dpi = 50)
cmyshow(fixed,warped_image, title="Fixed Red, Moving Green - After Registration", dpi = 50)