from __future__ import print_function
import sys
sys.path.insert(0,'/scratch/blowekamp/build/SimpleITK/SimpleITK-build/lib')
sys.path.insert(0,'/scratch/blowekamp/build/SimpleITK/SimpleITK-build/Wrapping')

import sys
import os

import SimpleITK as sitk
print(sitk.Version())

from myshow import myshow

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)

help(sitk.Transform)

tx = sitk.Transform(2, sitk.sitkTranslation)
print(tx)

print(tx.GetParameters())
print(tx.GetFixedParameters())

help(sitk.ImageRegistrationMethod)

fixedInputFilename = "/scratch/blowekamp/build/SimpleITK/ITK//Examples/Data/BrainProtonDensitySliceBorder20.png"
movingInputFilename = "/scratch/blowekamp/build/SimpleITK/ITK//Examples/Data/BrainProtonDensitySliceShifted13x17y.png"

fixedInput = sitk.ReadImage(fixedInputFilename)
fixed = sitk.VectorIndexSelectionCast(fixedInput,0,sitk.sitkFloat32)
myshow(fixed, dpi = 50, title="Fixed Image")
    
movingInput = sitk.ReadImage(movingInputFilename)
moving = sitk.VectorIndexSelectionCast(movingInput,0,sitk.sitkFloat32)
myshow(moving, dpi = 50, title="Moving Image")

cmyshow(fixed,moving, title="Fixed Red, Moving Green", dpi = 50)

Tx=sitk.Transform(fixed.GetDimension(), sitk.sitkTranslation)
R = sitk.ImageRegistrationMethod()
R.SetMetricAsMeanSquares()
R.SetOptimizerAsRegularStepGradientDescent(maxStep=4.0, minStep=.01, numberOfIterations=200 )
R.SetOptimizerScales([1.0, 1.0])
R.SetTransform(Tx)
R.SetInterpolator(sitk.sitkLinear)

outTx = R.Execute(fixed,moving)
print(outTx)

resampler = sitk.ResampleImageFilter()
resampler.SetReferenceImage(fixed);
resampler.SetInterpolator(sitk.sitkLinear)
resampler.SetDefaultPixelValue(100)
resampler.SetTransform(outTx)

outImg = resampler.Execute(moving)
cmyshow(fixed,moving, title="Fixed Red, Moving Green", dpi = 50)
cmyshow(fixed,outImg, title="Fixed Red, Registered Moving Green", dpi = 50)

def command_iteration(method) :
    print("{0} = {1} : {2}".format(method.GetOptimizerIteration(),
                                   method.GetMetricValue(),
                                   method.GetOptimizerPosition()),
              end='\n');
    sys.stdout.flush();
    
def command_end(R):
    global Tx
    print("-------")
    print(Tx)
    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) )

Tx= sitk.Transform(2, sitk.sitkTranslation)
R.SetTransform(Tx)
outTx = R.Execute(fixed,moving)

from IPython.display import clear_output
import time

def command_plot_value_start(method):
    global x, ax, f
    x = []

    f, ax = plt.subplots()
    ax.set_yscale('log')
    ax.set_title("Metric Value")   

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)

R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda:  clear_output(stdout=True, stderr=False) )
R.AddCommand( sitk.sitkStartEvent, lambda: command_plot_value_start(R) )
R.AddCommand( sitk.sitkIterationEvent, lambda: command_plot_value(R) )

R.SetOptimizerAsAmoeba(simplexDelta=4,parametersConvergenceTolerance=0.01,functionConvergenceTolerance=0.01); 
outTx = R.Execute(fixed, moving)

def command_display(method, xf):
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed);
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetDefaultPixelValue(100)
    resampler.SetTransform(xf)

    outImg = resampler.Execute(moving)
    cmyshow(fixed,outImg, title="Iteration {0}".format(method.GetOptimizerIteration()), dpi = 50)

R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda: command_iteration(R) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_end(R) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_display(R,Tx) )

outTx = R.Execute(fixed, moving)

fixedInputFilename = "/scratch/blowekamp/build/SimpleITK/ITK//Examples/Data/BrainProtonDensitySliceBorder20.png"
movingInputFilename = "/scratch/blowekamp/build/SimpleITK/ITK//Examples/Data/BrainProtonDensitySliceR10X13Y17.png"

#movingInputFilename = "/scratch/blowekamp/build/SimpleITK/ITK//Examples/Data/BrainProtonDensitySliceR10X13Y17S12.png"

fixedInput = sitk.ReadImage(fixedInputFilename)
fixed = sitk.VectorIndexSelectionCast(fixedInput,0,sitk.sitkFloat32)
    
movingInput = sitk.ReadImage(movingInputFilename)
moving = sitk.VectorIndexSelectionCast(movingInput,0,sitk.sitkFloat32)

cmyshow(fixed,moving, title="Fixed Red, Moving Green", dpi = 50)

Tx = sitk.Transform(fixed.GetDimension(), sitk.sitkTranslation)
R = sitk.ImageRegistrationMethod()
R.SetMetricAsMeanSquares()
R.SetOptimizerAsRegularStepGradientDescent(maxStep=4.0, minStep=.01, numberOfIterations=200 )
R.SetOptimizerScales([1.0, 1.0])
R.SetTransform(Tx)
R.SetInterpolator(sitk.sitkLinear)

R.RemoveAllCommands()
R.AddCommand( sitk.sitkIterationEvent, lambda: command_iteration(R) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_end(R) )
R.AddCommand( sitk.sitkEndEvent, lambda: command_display(R,Tx) )

outTx = R.Execute(fixed, moving)

Tx = sitk.Transform(outTx)
#Tx.AddTransform(sitk.Transform(fixed.GetDimension(),sitk.sitkEuler))
#Tx.AddTransform(sitk.Transform(fixed.GetDimension(),sitk.sitkAffine))
Tx.AddTransform(sitk.Transform(fixed.GetDimension(),sitk.sitkSimilarity))
print(Tx)

print(Tx.GetParameters())
print(Tx.GetFixedParameters())

R.SetTransform(Tx)
R.SetMetricAsMeanSquares()
#R.SetOptimizerAsAmoeba(simplexDelta=1,parametersConvergenceTolerance=0.01,functionConvergenceTolerance=0.01,numberOfIterations=500)
R.SetOptimizerAsRegularStepGradientDescent(maxStep=1,minStep=.001,numberOfIterations=500)
#R.SetOptimizerScales([100.0, 1, 1])
#R.SetOptimizerScales([10]*6)
R.SetOptimizerScales([100,100,1,1])
R.SetInterpolator(sitk.sitkLinear)
R.DebugOn()
outTx2 = R.Execute(fixed, moving)

cIdx = [ idx//2 for idx in fixed.GetSize()]
cPt = fixed.TransformIndexToPhysicalPoint(cIdx)
print(cPt)

Tx.SetFixedParameters(cPt)