#import tifffile
from cil.framework import DataProcessor, ImageData, DataOrder
import gi

gi.require_version('Ufo', '0.0')
from gi.repository import Ufo
import numpy as np


class UfoStandardBackProjector(DataProcessor):

    def __init__(self,
                 volume_geometry=None,
                 sinogram_geometry=None):
        kwargs = {
            'volume_geometry': volume_geometry,
            'sinogram_geometry': sinogram_geometry}

        super(UfoStandardBackProjector, self).__init__(**kwargs)

        self.set_ImageGeometry(volume_geometry)
        self.set_AcquisitionGeometry(sinogram_geometry)

    def check_input(self, dataset):

        if self.sinogram_geometry.shape != dataset.geometry.shape:
            raise ValueError("Dataset not compatible with geometry used to create the projector")

        return True

    def set_ImageGeometry(self, volume_geometry):

        DataOrder.check_order_for_engine('astra', volume_geometry)

        if len(volume_geometry.dimension_labels) > 3:
            raise ValueError("Supports 2D and 3D data only, got {0}".format(volume_geometry.number_of_dimensions))

        self.volume_geometry = volume_geometry.copy()

    def set_AcquisitionGeometry(self, sinogram_geometry):

        DataOrder.check_order_for_engine('astra', sinogram_geometry)

        if len(sinogram_geometry.dimension_labels) > 3:
            raise ValueError("Supports 2D and 3D data only, got {0}".format(self.volume_geometry.number_of_dimensions))

        self.sinogram_geometry = sinogram_geometry.copy()

    def process(self, out=None):

        DATA = self.get_input()

        data_temp = DATA.as_array()

        arr_out = self.create_backprojection(data_temp, out)

        if out is None:
            out = ImageData(arr_out, deep_copy=False, geometry=self.volume_geometry.copy(), suppress_warning=True)
            return out
        else:
            out.fill(arr_out)

    def create_backprojection(self, sino, vol):

        pm = Ufo.PluginManager()
        graph = Ufo.TaskGraph()
        scheduler = Ufo.Scheduler()

        read = pm.get_task('memory-in')
        read.props.pointer = sino.__array_interface__['data'][0]
        read.props.width = sino.shape[1]
        read.props.height = sino.shape[0]
        if(len(sino.shape)==2):
            read.props.number = 1
        else:
            read.props.number = sino.shape[2]
        read.props.bitdepth = 32

        backproject = pm.get_task('backproject')
        backproject.props.axis_pos = sino.shape[1] / 2

        vol_arr = np.zeros(self.volume_geometry.shape, dtype=np.float32)

        write = pm.get_task('memory-out')
        write.props.pointer = vol_arr.__array_interface__['data'][0]
        write.props.max_size = vol_arr.nbytes

        graph.connect_nodes(read, backproject)
        graph.connect_nodes(backproject, write)

        scheduler.run(graph)

        return vol_arr