# cil imports
from cil.framework import ImageData, ImageGeometry
from cil.framework import AcquisitionGeometry, AcquisitionData
from cil.optimisation.algorithms import CGLS, SIRT

import src.operators.ProjectionOperator as UfoProjectionOperator
import cil.plugins.astra.operators.ProjectionOperator as AstraProjectionOperator
import cil.plugins.astra.processors.FBP


# External imports
import tifffile
import os
import numpy as np
import matplotlib.pyplot as plt
import logging
import time

logging.basicConfig(level = logging.INFO)

# set up default colour map for visualisation
#cmap = "gray"

# set the backend for FBP and the ProjectionOperator
#device = 'gpu'

def reconstruct():
 # number of pixels
 n_pixels = 4096
 z_size = 1024
 #slices=2
 #rounding='int8'
 #rounding='half'
 #rounding='single'
 #n_pixels = 1024

 # Angles
 angles = np.linspace(0, 180, n_pixels, endpoint=False, dtype=np.float32)


 # Setup acquisition geometry
 # with sufficient number of projections
 ag = AcquisitionGeometry.create_Parallel3D()\
                            .set_angles(angles)\
                            .set_panel([n_pixels,z_size], pixel_size=1)
 #ag = AcquisitionGeometry.create_Parallel2D()\
 #                        .set_angles(angles)\
 #                        .set_panel(n_pixels,pixel_size=1)
 
 ag.set_labels(['vertical', 'angle', 'horizontal'])

 # Setup image geometry
 ig = ImageGeometry(voxel_num_x=n_pixels,
                   voxel_num_y=n_pixels,
                   voxel_num_z=z_size,
                   voxel_size_x=1,
                   voxel_size_y=1,
                   voxel_size_z=1)


 # Get phantom
 #home='/ccpi/data/new/phantom/'
 home='/ccpi/data/new/1024/phantom/'
 #home='/ccpi/data/test/1024/slices/'
# phantom = []
# for i in sorted(os.listdir(home)):
# for i in range(0, z_size):
  #file = os.path.join(home,i)
  #img = tifffile.imread(file) 
#  img = np.random.random_sample((n_pixels,n_pixels)).astype(np.float32)
#  if len(img.shape)==3:
#     img = np.squeeze(img, axis=0)
#  phantom.append(img)
# phantom = np.asarray(phantom)
 phantom = np.random.random_sample((z_size,n_pixels,n_pixels)).astype(np.float32)

 print(phantom.shape)
 phantom_reshaped = np.transpose(phantom, (0, 1, 2))
 print(phantom_reshaped.shape)
 phantom = ImageData(phantom_reshaped, deep_copy=False, geometry=ig)
 phantom_arr = phantom.as_array()
 tifffile.imsave('phantom_3d.tif', phantom_arr[:,:,0])

 # Create projection operator using Astra-Toolbox.
 out_file = "ufo_single.tif"
 sino_file= "sino_single.tif"

# Ufo-harish
 A = UfoProjectionOperator(ig, ag, True, rounding, slices)
# Ufo-farago
# A = UfoProjectionOperator(ig, ag, False, rounding, slices)
# Astra
# A = AstraProjectionOperator(ig, ag, 'gpu')
  
 #A.dot_test(A)

 # Create an acqusition data (numerically)
 sino = A.direct(phantom)
 sino_arr = sino.as_array()
 tifffile.imsave(sino_file,sino_arr[0,:,:])
 
 # back_projection
 back_projection = A.adjoint(sino)

 # initial estimate - zero array in this case
 initial = ig.allocate(0)

 # setup CGLS
 f = open("single.txt", "a")
 for i in range(20,21):
    total_time=0
    for j in range(0,1):
       cgls = CGLS(initial=initial,
                  operator=A,
                  data=sino,
                  tolerance=1e-16,
                  max_iteration = 50,
                  update_objective_interval = 1 )

       # run N interations
       start = time.time()
       cgls.run(i, verbose=True)
       end = time.time()
       if(j>0):
          total_time += (end-start)
    f.write('Iteration {} Time {} \n'.format(i,str(total_time/3)))
 f.close()
 print("finished reconstruction in single precision")

if __name__ == '__main__':
 reconstruct()