import itk
from sys import argv
import numpy as np

if len(argv) != 4:
    print("Usage: " + argv[0] + " <fixed2DImage> " + " <moving3DVolume> " + " <outputImage> " )
    exit(1)


# ---------- User-defined parameters ----------------
movInitTranslx = 0.
movInitTransly = 0.
movInitTranslz = 0.

# parameters for the Ray cast interpolator (same ones used to generate the FixedImage)
focal_length = 1000.
threshold = 0 # threshold above which the volume's intensities will be integrated.

# parameters for the resampled DRR (same ones used to generate the FixedImage)
dx = 1000 # size (# pixels) of the output image
dy = 1000 # size (# pixels) of the output image
sx = 0.5 # spacing x [mm] of the output image
sy = 0.5 # spacing y [mm] of the output image
o2Dx = 0 # principal points coordinate x [mm] (usually the 2D projection normal is assumed to be perpendicular )
o2Dy = 0 # principal points coordinate y [mm]

# ---------- Initializations ----------------
fixedImageFile = argv[1]
movingImageFile = argv[2]

PixelType = itk.F
Dimension = 3
ScalarType = itk.D

FixedImageType = itk.Image[PixelType, Dimension]
MovingImageType = itk.Image[PixelType, Dimension]

fixedImageReader  = itk.ImageFileReader[FixedImageType].New()
movingImageReader = itk.ImageFileReader[MovingImageType].New()

OptimizerType = itk.RegularStepGradientDescentOptimizerv4[ScalarType]

RegistrationType = itk.ImageRegistrationMethodv4[FixedImageType,
                                                 MovingImageType]

TransformType = itk.TranslationTransform[ScalarType, Dimension]

MetricType = itk.MeanSquaresImageToImageMetricv4[FixedImageType,
                                                 MovingImageType]

FixedInterpolatorType = itk.LinearInterpolateImageFunction[FixedImageType, ScalarType]
MovingInterpolatorType = itk.RayCastInterpolateImageFunction[MovingImageType, ScalarType]

#WriterTypeOutput = itk.ImageFileWriter[OutputImageType]


# ---------- Instantiate the classes for the registration framework ----------------
registration = RegistrationType.New()
initialTransform = TransformType.New()
FixedInitialTransform = TransformType.New()
imageMetric  = MetricType.New()
optimizer    = OptimizerType.New()
fixedInterpolator = FixedInterpolatorType.New()
movingInterpolator = MovingInterpolatorType.New()


# ---------- Read fixed and moving images ----------------
fixedImageReader.SetFileName( fixedImageFile )
movingImageReader.SetFileName( movingImageFile )

fixedImageReader.Update()
movingImageReader.Update()

fixedImage = fixedImageReader.GetOutput()
movingImage = movingImageReader.GetOutput()

movSpacing = movingImage.GetSpacing()
movOrigin = movingImage.GetOrigin()
movSize = movingImage.GetBufferedRegion().GetSize()
movCenter = np.asarray(movOrigin) + np.multiply(movSpacing, np.divide(movSize,2.)) - np.divide(movSpacing,2)

fixSpacing = fixedImage.GetSpacing()
fixSize = fixedImage.GetBufferedRegion().GetSize()

# ---------- Set Transform ----------------
initial_parameters = initialTransform.GetParameters()
initial_parameters[0] = movInitTranslx
initial_parameters[1] = movInitTransly
initial_parameters[2] = movInitTranslz
initialTransform.SetParameters(initial_parameters)

registration.SetInitialTransform( initialTransform )

# ---------- Set Ray Cast Interpolator ----------------
#movingInterpolator.SetTransform(initialTransform)
movingInterpolator.SetThreshold(threshold) # We can then specify a threshold above which the volume's intensities will be integrated.

focal_point = np.zeros(Dimension)
focal_point[0] = movCenter[0]
focal_point[1] = movCenter[1]
focal_point[2] = movCenter[2] - focal_length/2.
movingInterpolator.SetFocalPoint(focal_point)

# ---------- Set metric ----------------
imageMetric.SetFixedInterpolator(fixedInterpolator)
imageMetric.SetMovingInterpolator(movingInterpolator)

registration.SetMetric(imageMetric)

# ---------- Set Optimizer ----------------
# optimizer scale
translationScale = 1.0 # / 100.0

optimizerScales = itk.OptimizerParameters[ScalarType](initialTransform.GetNumberOfParameters())
optimizerScales.SetElement(0, translationScale)
optimizerScales.SetElement(1, translationScale)
optimizerScales.SetElement(2, translationScale)

optimizer.SetScales( optimizerScales )

# Set other parameters
optimizer.SetNumberOfIterations( 200 )
optimizer.SetLearningRate( 0.2 )
optimizer.SetMinimumStepLength( 0.001 )
optimizer.SetReturnBestParametersAndValue(True)

registration.SetOptimizer(optimizer)


# ---------- Iteration Observer ----------------
def iterationUpdate():
    currentParameter = registration.GetTransform().GetParameters()
    print("M: %f   P: %f %f" % ( optimizer.GetValue(),
                                 currentParameter.GetElement(0),
                                 currentParameter.GetElement(1)))

iterationCommand = itk.PyCommand.New()
iterationCommand.SetCommandCallable( iterationUpdate )
registration.AddObserver( itk.IterationEvent(), iterationCommand )



# ---------- Registration ----------------
# Set image plane
fixOrigin = [0]*Dimension
fixOrigin[0] = movCenter[0] + o2Dx - fixSpacing[0]*(fixSize[0] - 1.)/2.
fixOrigin[1] = movCenter[1] + o2Dy - fixSpacing[1]*( fixSize[1] - 1.)/2.
fixOrigin[2] = movCenter[2] + focal_length/2.

fixedImageReader.GetOutput().SetOrigin(fixOrigin)

# Set fixed and moving images
registration.SetFixedImage(fixedImageReader.GetOutput())
registration.SetMovingImage(movingImageReader.GetOutput())

# One level registration process without shrinking and smoothing.
registration.SetNumberOfLevels(1)
registration.SetSmoothingSigmasPerLevel([0])
registration.SetShrinkFactorsPerLevel([1])

print("Starting registration")

# Start the registration process
registration.Update()