import numpy as np
import cv2
from skimage.transform import resize 
from skimage.data import shepp_logan_phantom

# simulate 3D data
phantom = shepp_logan_phantom()
phantom = resize(phantom, (1024, 1024))
phantom_vec = np.zeros([432, 1024, 1024])  
for i in range(24,406):
    phantom_vec[i, :, :] = phantom 
    
#params
ViewNum = 720
det_row_count = 32
det_col_count = 864
detector_spacing_x = 1.06 
detector_spacing_y = 1.12
px_of_one_det = 32
sod = 630 
odd = 340
absolute_pitch = 12.0 # mm/circle
helical_pitch = absolute_pitch / (detector_spacing_y  * det_row_count  * sod / (sod + odd))
print(helical_pitch)

#Leap setting
from leapctype import *
leapct = tomographicModels()
volume_section_num_per_circle = 24
circle_num = 20

angles = leapct.setAngleArray(circle_num * ViewNum, 360.0*circle_num) 
leapct.set_conebeam(circle_num * ViewNum, det_row_count, det_col_count, detector_spacing_y, detector_spacing_x, 0.5*(det_row_count-1), 0.5*(det_col_count-1), angles, sod, sod + odd)
leapct.set_curvedDetector()

voxel_height = abs(absolute_pitch / volume_section_num_per_circle)
print(f'slice thick:{voxel_height}')
leapct.set_normalizedHelicalPitch(helical_pitch)

volume_section_num = volume_section_num_per_circle * (circle_num - 2)
leapct.set_volume(1024, 1024, volume_section_num, voxelWidth=0.5, voxelHeight=voxel_height)

#project
g = leapct.allocate_projections()
print(g.shape)
leapct.project(g,phantom_vec.astype(np.float32))

#back project
res = leapct.FBP(g)

import matplotlib.pyplot as plt
slice_height, slice_width = res.shape[0], res.shape[1]  
dpi = plt.gcf().dpi
fig = plt.figure(figsize=(slice_width / dpi, slice_height / dpi))
ax2 = fig.add_subplot(1, 1, 1)
ax2.imshow(res[:, :, 256], vmin=np.max(res[:, :, 256]) * 0.9, vmax=np.max(res[:, :, 256]), cmap='gray')
plt.show()
plt.close(fig)