#!/usr/bin/env python
from __future__ import print_function
import SimpleRTK as srtk
import sys
import os
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import scipy.optimize

# Geometry parameters
slscale = 90.
xc = 0.
yc = -100.
xcr = xc
ycr = yc 
spacingVol = [ 2 ] * 3
RD = 400.
R = 700.
umin = -175.3
umax = 233.9
numberOfProjections = 720
I0=1e7
dH2O=0.01879 #mm^-1 at 75 keV

sizeOutput = [ 256, 256,  numberOfProjections ]
sizeOutputFDK = [ 2500, sizeOutput[1],  numberOfProjections ]
sizeOutputVol = [ 150 ] * 3

spacing = [ (umax-umin)/(sizeOutput[0]) ]*3
origin = [ umin+spacing[0]/-2., (sizeOutput[1]-1)/-2.*spacing[1], 0. ]
originVol = [ (sizeOutputVol[0]-1)/-2.*spacingVol[0] + xcr, (sizeOutputVol[1]-1)/-2.*spacingVol[1], (sizeOutputVol[2]-1)/-2.*spacingVol[2] + ycr ]
originSl = [ xcr / slscale, 0.25, ycr / slscale ]

# Source coordinates
beta = np.array(np.linspace(0, 2*np.pi,numberOfProjections))
xv = -R * np.sin(beta)
yv =  R * np.cos(beta)

# Fan angle of the FOV center
normvc = np.sqrt((xc-xv)**2+(yc-yv)**2)
alphac = np.arctan2((yc-yv)/normvc, (xc-xv)/normvc) - np.arctan2(-yv/R, -xv/R)
alphac = alphac-2*np.pi*np.round(alphac/(2*np.pi))

# Find alpha numericall
def alphafunc(alpha, alphac, xc, yc, RD, R, umin, umax):
    Dalpha = RD+R*np.cos(alpha);
    return alpha - alphac + 0.5*np.arctan((umin+umax-2*R*np.sin(alpha)) * Dalpha / (Dalpha**2-(umin-R*np.sin(alpha)) * (umax-R*np.sin(alpha))))
        
alpha, alphacov = scipy.optimize.leastsq(alphafunc, 0*beta, args=(alphac, xc, yc, RD, R, umin, umax))
print(180./np.pi*np.min(alpha))
print(180./np.pi*np.max(alpha))
sid = R*np.cos(alpha)
sdd = sid + RD
sox = R*np.sin(alpha)

# Defines the RTK geometry object
geometry = srtk.ThreeDCircularProjectionGeometry()
geometryFDK = srtk.ThreeDCircularProjectionGeometry()
for i in range(0, numberOfProjections):
  geometry.AddProjection(sid[i], sdd[i], -180/np.pi*(beta[i]+alpha[i]), 0, 0, 0, 0, sox[i], 0)
  geometryFDK.AddProjection(R, R, -180/np.pi*beta[i])

constantImageSource = srtk.ConstantImageSource()
constantImageSource.SetOrigin( originVol )
constantImageSource.SetSpacing( spacingVol )
constantImageSource.SetSize( sizeOutputVol )
constantImageSource.SetConstant(0.)
sourceVol = constantImageSource.Execute()

dpScale = 120.
cylinder = srtk.DrawCylinderImageFilter()
cylinder.SetAxis(dpScale*np.array([0.85, 0, 0.85]))
cylinder.SetCenter([xc,0,yc])
cylinder.SetDensity(1.)
vol = cylinder.Execute(sourceVol)

ellipsoid = srtk.DrawEllipsoidImageFilter()
ellipsoid.SetAxis(dpScale*np.array([0.6,0.03,0.6]))
for i in range(8):
    if i%2==0:
        ellipsoid.SetDensity(-0.7)
    else:
        ellipsoid.SetDensity(0.7)
    ellipsoid.SetCenter([xc,i*0.1*dpScale,yc])
    vol = ellipsoid.Execute(vol)

cube = srtk.DrawCubeImageFilter()
cube.SetAxis([40]*3)
cube.SetCenter([xc,-60,yc])
vol = cube.Execute ( vol )

writer = srtk.ImageFileWriter()
writer.SetFileName ( 'vol2.mha' )
writer.UseCompressionOn()
writer.Execute ( vol )

# Now create projections
constantImageSource.SetConstant(0.)
sourceProj = constantImageSource.Execute()

cylinder = srtk.RayCylinderIntersectionImageFilter()
cylinder.SetAxis(dpScale*np.array([0.85, 0, 0.85]))
cylinder.SetCenter([xc,0,yc])
cylinder.SetDensity(1.)
cylinder.SetGeometry( geometry )
proj = cylinder.Execute(sourceProj)

ellipsoid = srtk.ProjectEllipsoidImageFilter()
ellipsoid.SetAxis(dpScale*np.array([0.6,0.03,0.6]))
ellipsoid.SetGeometry( geometry )
for i in range(8):
    if i%2==0:
        ellipsoid.SetDensity(-0.7)
    else:
        ellipsoid.SetDensity(0.7)
    ellipsoid.SetCenter([xc,i*0.1*dpScale,yc])
    proj = ellipsoid.Execute(proj)

cube = srtk.ProjectCubeImageFilter()
cube.SetAxis([40]*3)
cube.SetCenter([xc,-60,yc])
cube.SetGeometry( geometry )
proj = cube.Execute ( proj )

# Reconstruct
constantImageSource = srtk.ConstantImageSource()
constantImageSource.SetOrigin( originVol )
constantImageSource.SetSpacing( spacingVol )
constantImageSource.SetSize( sizeOutputVol )
constantImageSource.SetConstant(0.)
sourceVol = constantImageSource.Execute()

fdkFilter = srtk.FDKConeBeamReconstructionFilter()
fdkFilter.SetGeometry( geometry )
fdk = fdkFilter.Execute( sourceVol, proj )
writer.SetFileName ( 'imagingring2.mha' )
writer.Execute ( fdk )
sys.exit()

# Simulate FDK type projections
spacing = [spacing[0]*R/(RD+R)]*3
print(spacing)
for i in range(0, 2):
    origin[i] = (sizeOutputFDK[i]-1)/-2.*spacing[i]
constantImageSource = srtk.ConstantImageSource()
constantImageSource.SetOrigin( origin )
constantImageSource.SetSpacing( spacing )
constantImageSource.SetSize( sizeOutputFDK )
constantImageSource.SetConstant(0.)
sourceProj = constantImageSource.Execute()

slFilter = srtk.SheppLoganPhantomFilter()
slFilter.SetGeometry( geometryFDK )
slFilter.SetOriginOffset( originSl )
slFilter.SetPhantomScale( slscale )
sl = slFilter.Execute(sourceProj)
if I0!=0:
    slarray = srtk.GetArrayFromImage(sl)
    slarray = I0*np.exp(-1.*dH2O*slarray)
    slarray = np.maximum(np.random.poisson(slarray), 1)
    slarray = np.log(I0/slarray)/dH2O
    slarray = srtk.GetImageFromArray(slarray.astype(np.float32))
    slarray.CopyInformation(sl)
    sl = slarray

writer = srtk.ImageFileWriter()
if I0==0:
    writer.SetFileName ( '/tmp/projFDK.mha' )
else:
    writer.SetFileName ( '/tmp/projFDK_poisson.mha' )
writer.Execute ( sl )

# Reconstruct
constantImageSource = srtk.ConstantImageSource()
constantImageSource.SetOrigin( originVol )
constantImageSource.SetSpacing( spacingVol )
constantImageSource.SetSize( sizeOutputVol )
constantImageSource.SetConstant(0.)
sourceVol = constantImageSource.Execute()

fdkFilter = srtk.FDKConeBeamReconstructionFilter()
fdkFilter.SetGeometry( geometryFDK )
fdk = fdkFilter.Execute( sourceVol, sl )

#fov = srtk.FieldOfView()
#srtk.SetGeometry(geometryFDK)
#fdk = fov.Execute(fdk, sl)

if I0==0:
    writer.SetFileName ( '/tmp/fdk.mha' )
else:
    writer.SetFileName ( '/tmp/fdk_poisson.mha' )
writer.Execute ( fdk )

sourceVol = constantImageSource.Execute()
slFilterRef = srtk.DrawSheppLoganFilter()
slFilterRef.SetPhantomScale( slscale )
slFilterRef.SetOriginOffset( originSl )
sl = slFilterRef.Execute(sourceVol)
writer.SetFileName ( '/tmp/ref.mha' )
writer.Execute ( sl )

# write geometry
geometryWriter = srtk.ThreeDCircularProjectionGeometryXMLFileWriter()
geometryWriter.SetFileName ( "geometryFDK.xml" )
geometryWriter.Execute ( geometryFDK );
if I0==0:
    os.system("rtkfdk -v --lowmem -p /tmp -r projFDK.mha -o /tmp/fdk_fullFDK.mha --dimension 600 --spacing 0.5 --origin -149.75,-149.75,-249.75 -g geometryFDK.xml --hardware cuda")
else:
    os.system("rtkfdk -v --lowmem -p /tmp -r projFDK_poisson.mha -o /tmp/fdk_fullFDK_poisson.mha --dimension 600 --spacing 0.5 --origin -149.75,-149.75,-249.75 -g geometryFDK.xml --hardware cuda")
os.system("rtkdrawshepploganphantom --offset 0,0.25,-1 -o /tmp/ref_full.mha --dimension 600 --spacing 0.5 --origin -149.75,-149.75,-249.75 --phantomscale 100 --hardware cuda")