# In this script, a digitally reconstructed radiograph is created of a CT image from different angles
# For this, the itk wrapper is used. (so if simple itk wrapper is needed for something, the image should be transformed)
import itk
from itk import RTK as rtk
from downloaddata import fetch_data as fdata
import numpy as np

# Always write output to a separate directory, we don't want to pollute the source directory.
import os
OUTPUT_DIR = 'Output'

# Load CT image and change imagetype to float
filename = 'ctpreop.mha'

CT = itk.imread(filename, pixel_type=itk.F)

CT_size = itk.size(CT)
CT_spacing = CT.GetSpacing()
# DataCollectionDiameter = 500 #TAG (0018, 0090)
# ReconstructionDiamgeter = 425 #TAG (0018, 1100)
DistanceSourceToDetectorCT = 1040 # TAG (0018, 1110)
DistanceSourceToPatientCT = 570 # TAG (0018, 1111)
# ImagePosition(Patient) = -203.2024\153.705314226727\85.0300865766952 # TAG (0020,0032) (wordt in itk-snap origin genoemd)
# ImageOrientation(Patient) = 0.9960354\-0.08895722\1.83691E-16\1.884096E-16\4.464693E-17\-1 (0020, 0037)
# Slice Location = 0 (0020, 1041)

# Change values that air is 0
CT_np = itk.GetArrayFromImage(CT)
CT_np = CT_np + 1024
CT = itk.GetImageFromArray(CT_np)

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

numberOfProjections = 360

# 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] # Third coordinate is ignored since DRRs are 2D
sizeOutput = [CT_size[0], CT_size[1],  numberOfProjections ]
#spacing = [ CT_spacing[0], CT_spacing[1], 1. ]
spacing = [ 2.0, 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 = DistanceSourceToPatientCT # source to isocenter distance
sdd = DistanceSourceToDetectorCT # 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 (os.path.join(OUTPUT_DIR,'geometry.xml'))
xmlWriter.SetObject ( geometry )
xmlWriter.WriteFile()

# Create forward projection filter
REIType = rtk.JosephForwardProjectionImageFilter[ImageType, ImageType]
rei = REIType.New()
semiprincipalaxis = [ 50, 50, 50]
center = [-203.2024, 153.705314226727, 85.0300865766952]

# Constructing
rei.SetGeometry( geometry )
rei.SetInput(0, constantImageSource.GetOutput())
rei.SetInput(1, CT)
rei.Update()

# Write file
itk.imwrite(rei.GetOutput(), os.path.join(OUTPUT_DIR,'projection.mha'))