import sys
import itk
from itk import RTK as rtk
import itk

# Loading 3D CT image
CT = itk.imread("training_001_ct.mha", pixel_type=itk.F)

# Defines the image type
Dimension_CT = 3
PixelType = itk.F
ImageType = itk.Image[PixelType, Dimension_CT]

# Create a stack of empty projection images
ConstantImageSourceType = rtk.ConstantImageSource[ImageType]
constantImageSource = ConstantImageSourceType.New()

# Define origin, sizeOutput and spacing (still need to change these)
origin = [ -127., -127., 0.]
numberOfProjections = 360
sizeOutput = [ 512, 512, numberOfProjections]
spacing = [ 0.653595, 2.0, 1.]

constantImageSource.SetOrigin( origin )
constantImageSource.SetSpacing( spacing )
constantImageSource.SetSize( sizeOutput )
constantImageSource.SetConstant(0.)

# Defines the RTK geometry object
geometry = rtk.ThreeDCircularProjectionGeometry.New()
firstAngle = 0.
angularArc = 360.
sid = 600 # source to isocenter distance
sdd = 1200 # source to detector distance
for x in range(0,numberOfProjections):
  angle = firstAngle + x * angularArc / numberOfProjections
  geometry.AddProjection(sid,sdd,angle)

# Writing the geometry to disk
xmlWriter = rtk.ThreeDCircularProjectionGeometryXMLFileWriter.New()
xmlWriter.SetFilename ( 'geo.xml' )
xmlWriter.SetObject ( geometry )
xmlWriter.WriteFile()

REIType = rtk.JosephForwardProjectionImageFilter[ImageType, ImageType]
rei = REIType.New()
semiprincipalaxis = [ 50, 50, 50]
center = [ 0, 0, 10]

rei.SetGeometry( geometry )
rei.SetInput(0, constantImageSource.GetOutput())
rei.SetInput(1, CT)
rei.Update()
itk.imwrite(rei.GetOutput(), 'stack.mha')