import sys
import itk
from itk import RTK as rtk
import glob
import time

T0 = time.time()

# Defines the image type
GPUImageType = rtk.CudaImage[itk.F,3]
CPUImageType = rtk.Image[itk.F,3]


# Defines the RTK geometry object
geometry = rtk.ThreeDCircularProjectionGeometry.New()
numberOfProjections = 720
firstAngle = 0.
angularArc = 360.
sid = 200 # source to isocenter distance
sdd = 385 # 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 ( 'op.xml' )
xmlWriter.SetObject ( geometry )
xmlWriter.WriteFile()


# Read all tif images from /img folder
image_filenames = sorted(glob.glob('./data/gr/*.tif'))  

tiffio = itk.TIFFImageIO.New()

ProjectionsReaderType = rtk.ProjectionsReader[CPUImageType]
projectionsSource = ProjectionsReaderType.New()
projectionsSource.SetImageIO(tiffio)
projectionsSource.SetFileNames(image_filenames)
projOrigin = [ -1.0*(1536-1)/2, -1.0*(1536-1)/2, 0 ] #input images are 1536x1536 pixels
projSpacing = [ .44, .44, .44 ]
projectionsSource.SetOrigin( projOrigin  )
projectionsSource.SetSpacing( projSpacing )

print('Read TIFF Done.')


# REIType = rtk.RayEllipsoidIntersectionImageFilter[CPUImageType, CPUImageType]
# rei = REIType.New()
# semiprincipalaxis = [ 50, 50, 50]
# center = [ 0, 0, 10]
# # Set GrayScale value, axes, center...
# rei.SetDensity(2)
# rei.SetAngle(0)
# rei.SetCenter(center)
# rei.SetAxis(semiprincipalaxis)
# rei.SetGeometry( geometry )
# rei.SetInput(constantImageSource.GetOutput())

ConstantImageSourceType = rtk.ConstantImageSource[GPUImageType]
# Create reconstructed image
constantImageSource2 = ConstantImageSourceType.New()
sizeOutput = [ 1000, 1000, 1000 ]
origin = [ -500, -500, -500 ]
spacing = [ 1, 1, 1 ]
constantImageSource2.SetOrigin( origin  )
constantImageSource2.SetSpacing( spacing )
constantImageSource2.SetSize( sizeOutput )
constantImageSource2.SetConstant(0.)

# Graft the projections to an itk::CudaImage
projections = GPUImageType.New()
projectionsSource.Update()
projections.SetPixelContainer(projectionsSource.GetOutput().GetPixelContainer())
projections.CopyInformation(projectionsSource.GetOutput())
projections.SetBufferedRegion(projectionsSource.GetOutput().GetBufferedRegion())
projections.SetRequestedRegion(projectionsSource.GetOutput().GetRequestedRegion())

# FDK reconstruction
print("Reconstructing...")
FDKGPUType = rtk.CudaFDKConeBeamReconstructionFilter
feldkamp = FDKGPUType.New()
feldkamp.SetInput(0, constantImageSource2.GetOutput())
feldkamp.SetInput(1, projections)
feldkamp.SetGeometry(geometry)
feldkamp.GetRampFilter().SetTruncationCorrection(0.0)
feldkamp.GetRampFilter().SetHannCutFrequency(0.0)

# Field-of-view masking
FOVFilterType = rtk.FieldOfViewImageFilter[CPUImageType, CPUImageType]
fieldofview = FOVFilterType.New()
fieldofview.SetInput(0, feldkamp.GetOutput())
fieldofview.SetProjectionsStack(projectionsSource.GetOutput())
fieldofview.SetGeometry(geometry)

# Writer
print("Writing output image...")
WriterType = rtk.ImageFileWriter[CPUImageType]
writer = WriterType.New()
writer.SetFileName('cudaop.mhd')
writer.SetInput(feldkamp.GetOutput())
writer.Update()

writer2 = WriterType.New()
writer2.SetFileName('cudaop.tif')
writer2.SetInput(feldkamp.GetOutput())
writer2.Update()

print("Done!")

T1 = time.time()

print("Time cost: %ss" % (T1-T0) )