import os
import math
import numpy as np
import itk

# <-res float float> DRR Pixel spacing in isocenter plane in mm {im_sx, im_sy}
# <-size int int> Size of DRR in number of pixels [default: 512x512] {dx, dy}
# <-scd float> Source to isocenter (i.e., 3D image center) distance in mm [default: 1000mm]
# <-t float float float> CT volume translation in x, y, and z direction in mm {tx, ty, tz}
# <-rx float> CT volume rotation about x axis in degrees
# <-ry float> CT volume rotation about y axis in degrees
# <-rz float> CT volume rotation about z axis in degrees
# <-2dcx float float> Central axis position of DRR in continuous indices {o2Dx, o2Dy}
# <-iso float float float> Continous voxel indices of CT isocenter (center of rotation and projection center) {cx, cy, cz}
# <-rp float> Projection angle in degrees {rprojection}
# <-threshold float> CT intensity threshold, below which are ignored [default: 0]
def siddonJacobRayTracingDRR(input_name=None, output_name=None, rprojection=0., rx=0., ry=0., rz=0., tx=0., ty=0., tz=0., 
                             cx=0., cy=0., cz=0., threshold=0., scd=1000, im_sx=0.65104, im_sy=0.65104, 
                             dx=512, dy=512, o2Dx=0.0, o2Dy=0.0):
    Dimension = 3
    InputPixelType = itk.F
    OutputPixelType = itk.UC
    
    InputImageType = itk.Image[InputPixelType, Dimension]
    OutputImageType = itk.Image[OutputPixelType, Dimension]
    
    ReaderType = itk.ImageFileReader[InputImageType]
    reader = ReaderType.New()
    reader.SetFileName(input_name)
    reader.Update()
    inputImage = reader.GetOutput()
    
    # To simply Siddon-Jacob's fast ray-tracing algorithm, we force the origin of the CT image
    # to be (0,0,0). Because we align the CT isocenter with the central axis, the projection
    # geometry is fully defined. The origin of the CT image becomes irrelavent.
    ctOrigin = [0.0, 0.0, 0.0]
    inputImage.SetOrigin(ctOrigin)
    
    FilterType = itk.ResampleImageFilter[InputImageType, InputImageType]
    filter = FilterType.New()
    filter.SetInput(inputImage)
    filter.SetDefaultPixelValue(0)
    
    # An Euler transformation is defined to position the input volume.
    TransformType = itk.Euler3DTransform[itk.D]
    transform = TransformType.New()
    transform.SetComputeZYX(True)
    
    translation = [tx, ty, tz]
    
    # converting degrees into radians old val 1./180.*math.pi 
    dtr = (math.atan(1.0) * 4.0) / 180.0
    
    transform.SetTranslation(translation)
    transform.SetRotation(dtr * rx, dtr * ry, dtr * rz)
    
    imOrigin = inputImage.GetOrigin()
    imRes = inputImage.GetSpacing()
    imRegion = inputImage.GetBufferedRegion()
    imSize = imRegion.GetSize()
    
    isocenter = [imOrigin[i] + (imRes[i]*imSize[i]/2.0) for i in range(3)]  
    
    transform.SetCenter(isocenter)
    
    InterpolatorType = itk.SiddonJacobsRayCastInterpolateImageFunction[InputImageType, itk.D]
    interpolator = InterpolatorType.New()
    
    interpolator.SetProjectionAngle(dtr * rprojection) # Set angle between projection central axis and -z axis
    interpolator.SetFocalPointToIsocenterDistance(scd) # Set source to isocenter distance
    interpolator.SetThreshold(threshold) # Set intensity threshold, below which are ignored.
    interpolator.SetTransform(transform)
    interpolator.Initialize()
    
    filter.SetInterpolator(interpolator)
    
    # number of pixels of the 2D DRR image
    size = [dx, dy, 1]
    # pixel spacing of the 2D DRR image [mm]
    spacing = [im_sx, im_sy, 1]
    
    filter.SetSize(size)
    filter.SetOutputSpacing(spacing)
    
    # Use the image centers as the central axis position.
    o2Dx = (dx - 1.) / 2.
    o2Dy = (dy - 1.) / 2.
    
    # Compute the origin (in mm) of the 2D Image
    origin = [-im_sx * o2Dx, -im_sy * o2Dy, -scd]
    
    filter.SetOutputOrigin(origin)
    filter.Update()
    
    RescaleType = itk.RescaleIntensityImageFilter[InputImageType, OutputImageType]
    rescaler = RescaleType.New()
    rescaler.SetOutputMinimum(0)
    rescaler.SetOutputMaximum(255)
    rescaler.SetInput(filter.GetOutput())
    rescaler.Update()
    
    # Out of some reason, the computed projection is upsided-down.
    # Here we use a FilpImageFilter to flip the images in y direction.
    FlipFilterType = itk.FlipImageFilter[OutputImageType]
    flipFilter = FlipFilterType.New()
    
    FlipAxesArray = itk.FixedArray[itk.D, 3]()
    FlipAxesArray[0] = 0
    FlipAxesArray[1] = 1
    FlipAxesArray[2] = 0
    
    flipFilter.SetFlipAxes(FlipAxesArray)
    flipFilter.SetInput(rescaler.GetOutput())
    flipFilter.Update()
    
    WriteType = itk.ImageFileWriter[OutputImageType]
    writer = WriteType.New()
    writer.SetFileName(output_name)
    writer.SetInput(flipFilter.GetOutput())
    writer.Update()
    
    
def generateAndSaveDRR(dataFolder, outputFolder, numberOfProjections):
    firstAngle = 0
    angularArc = 360
    
    all_img_vol_names = [x for x in os.listdir(dataFolder) if x.endswith(".nii.gz")]
    
    for i, relative_path in enumerate(all_img_vol_names):
        img_vol_path = os.path.join(dataFolder, relative_path)
        
        drr_img_name = relative_path.rsplit('.', 2)[0]
        
        for noProj in range(numberOfProjections):
            angle = firstAngle + noProj * angularArc / numberOfProjections
            drr_img_rel_path = drr_img_name + '_' + str(int(angle))
            output_drr_Path = os.path.join(outputFolder, drr_img_rel_path + ".png")
            siddonJacobRayTracingDRR(input_name=img_vol_path, output_name=output_drr_Path, rprojection=angle)
            

dataFolder = 'F:/all_CTs_HU_corrected'
outputFolder = 'F:/ITK_DRRs/all_drrs'
generateAndSaveDRR(dataFolder, outputFolder, 10)