#!/usr/bin/env python
# /*##########################################################################
#
# Copyright (c) 2016-2023 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/
"""Module for (filtered) backprojection on the GPU"""

__authors__ = ["A. Mirone, P. Paleo"]
__license__ = "MIT"
__date__ = "25/01/2019"

import logging
import numpy as np

from .common import pyopencl
from .processing import EventDescription, OpenclProcessing, BufferDescription
from .sinofilter import SinoFilter

if pyopencl:
    mf = pyopencl.mem_flags
    import pyopencl.array as parray
else:
    raise ImportError("Please install pyopencl in order to use opencl backprojection")
logger = logging.getLogger(__name__)


def _sizeof(Type):
    """
    return the size (in bytes) of a scalar type, like the C behavior
    """
    return np.dtype(Type).itemsize


def _idivup(a, b):
    """
    return the integer division, plus one if `a` is not a multiple of `b`
    """
    return (a + (b - 1)) // b


class Backprojection(OpenclProcessing):
    """A class for performing the backprojection using OpenCL"""

    kernel_files = ["backproj.cl", "array_utils.cl"]

    def __init__(
        self,
        sino_shape,
        slice_shape=None,
        axis_position=None,
        angles=None,
        filter_name=None,
        ctx=None,
        devicetype="all",
        platformid=None,
        deviceid=None,
        profile=False,
        extra_options=None,
    ):
        """Constructor of the OpenCL (filtered) backprojection

        :param sino_shape: shape of the sinogram. The sinogram is in the format
                           (n_b, n_a) where n_b is the number of detector bins
                           and n_a is the number of angles.
        :param slice_shape: Optional, shape of the reconstructed slice. By
                            default, it is a square slice where the dimension
                            is the "x dimension" of the sinogram (number of
                            bins).
        :param axis_position: Optional, axis position. Default is
                              `(shape[1]-1)/2.0`.
        :param angles: Optional, a list of custom angles in radian.
        :param filter_name: Optional, name of the filter for FBP. Default is
                            the Ram-Lak filter.
        :param ctx: actual working context, left to None for automatic
                    initialization from device type or platformid/deviceid
        :param devicetype: type of device, can be "CPU", "GPU", "ACC" or "ALL"
        :param platformid: integer with the platform_identifier, as given by
                           clinfo
        :param deviceid: Integer with the device identifier, as given by clinfo
        :param profile: switch on profiling to be able to profile at the kernel
                        level, store profiling elements (makes code slightly
                        slower)
        :param extra_options: Advanced extra options in the form of a dict.
            Current options are: cutoff, use_numpy_fft
        """
        # OS X enforces a workgroup size of 1 when the kernel has
        # synchronization barriers if sys.platform.startswith('darwin'):
        #  assuming no discrete GPU
        #    raise NotImplementedError("Backprojection is not implemented on CPU for OS X yet")

        OpenclProcessing.__init__(
            self,
            ctx=ctx,
            devicetype=devicetype,
            platformid=platformid,
            deviceid=deviceid,
            profile=profile,
        )

        self._init_geometry(
            sino_shape, slice_shape, angles, axis_position, extra_options
        )
        self._allocate_memory()
        self._compute_angles()
        self._init_kernels()
        self._init_filter(filter_name)

    def _init_geometry(
        self, sino_shape, slice_shape, angles, axis_position, extra_options
    ):
        """Geometry Initialization

        :param sino_shape: shape of the sinogram. The sinogram is in the format
                           (n_b, n_a) where n_b is the number of detector bins
                           and n_a is the number of angles.
        :param slice_shape: shape of the reconstructed slice. By
                            default, it is a square slice where the dimension
                            is the "x dimension" of the sinogram (number of
                            bins).
        :param angles: list of projection angles in radian.
        :param axis_position: axis position
        :param dict extra_options: Advanced extra options
        """
        self.shape = sino_shape
        self.num_bins = np.int32(sino_shape[1])
        self.num_projs = np.int32(sino_shape[0])
        self.angles = angles
        if slice_shape is None:
            self.slice_shape = (self.num_bins, self.num_bins)
        else:
            self.slice_shape = slice_shape
        self.dimrec_shape = (
            _idivup(self.slice_shape[0], 32) * 32,
            _idivup(self.slice_shape[1], 32) * 32,
        )
        if axis_position:
            self.axis_pos = np.float32(axis_position)
        else:
            self.axis_pos = np.float32((sino_shape[1] - 1.0) / 2)
        self.axis_array = None  # TODO: add axis correction front-end
        self._init_extra_options(extra_options)

    def _init_extra_options(self, extra_options):
        """Backprojection extra option initialization

        :param dict extra_options: Advanced extra options
        """
        self.extra_options = {
            "cutoff": 1.0,
            "use_numpy_fft": False,
            # It is  axis_pos - (num_bins-1)/2  in PyHST
            "gpu_offset_x": 0.0,  # self.axis_pos - (self.num_bins - 1) / 2.,
            "gpu_offset_y": 0.0,  # self.axis_pos - (self.num_bins - 1) / 2.
        }
        if extra_options is not None:
            self.extra_options.update(extra_options)

    def _allocate_memory(self):
        # Host memory
        self.slice = np.zeros(self.dimrec_shape, dtype=np.float32)
        self._use_textures = self.check_textures_availability()

        # Device memory
        self.buffers = [
            BufferDescription("_d_slice", self.dimrec_shape, np.float32, mf.READ_WRITE),
            BufferDescription(
                "d_sino", self.shape, np.float32, mf.READ_WRITE
            ),  # before transferring to texture (if available)
            BufferDescription("d_cos", (self.num_projs,), np.float32, mf.READ_ONLY),
            BufferDescription("d_sin", (self.num_projs,), np.float32, mf.READ_ONLY),
            BufferDescription("d_axes", (self.num_projs,), np.float32, mf.READ_ONLY),
        ]
        self.allocate_buffers(use_array=True)
        self.d_sino = self.cl_mem["d_sino"]  # shorthand

        # Texture memory (if relevant)
        if self._use_textures:
            self._allocate_textures()

        # Local memory
        self.local_mem = 256 * 3 * _sizeof(np.float32)  # constant for all image sizes

    def _compute_angles(self):
        if self.angles is None:
            self.angles = np.linspace(0, np.pi, self.num_projs, False)
        h_cos = np.cos(self.angles).astype(np.float32)
        h_sin = np.sin(self.angles).astype(np.float32)
        self.cl_mem["d_cos"][:] = h_cos[:]
        self.cl_mem["d_sin"][:] = h_sin[:]
        if self.axis_array:
            self.cl_mem["d_axes"][:] = self.axis_array.astype(np.float32)[:]
        else:
            self.cl_mem["d_axes"][:] = (
                np.ones(self.num_projs, dtype="f") * self.axis_pos
            )

    def _init_kernels(self):
        compile_options = None
        if not (self._use_textures):
            compile_options = "-DDONT_USE_TEXTURES"
        OpenclProcessing.compile_kernels(
            self, self.kernel_files, compile_options=compile_options
        )
        # check that workgroup can actually be (16, 16)
        self.compiletime_workgroup_size = self.kernels.max_workgroup_size(
            "backproj_cpu_kernel"
        )
        # Workgroup and ndrange sizes are always the same
        self.wg = (16, 16)
        self.ndrange = (
            _idivup(int(self.dimrec_shape[1]), 32) * self.wg[0],
            _idivup(int(self.dimrec_shape[0]), 32) * self.wg[1],
        )
        # Prepare arguments for the kernel call
        if not (self._use_textures):
            d_sino_ref = self.d_sino.data
        else:
            d_sino_ref = self.d_sino_tex
        self._backproj_kernel_args = (
            # num of projections (int32)
            self.num_projs,
            # num of bins (int32)
            self.num_bins,
            # axis position (float32)
            self.axis_pos,
            # d_slice (__global float32*)
            self.cl_mem["_d_slice"].data,
            # d_sino (__read_only image2d_t or float*)
            d_sino_ref,
            # gpu_offset_x (float32)
            np.float32(self.extra_options["gpu_offset_x"]),
            # gpu_offset_y (float32)
            np.float32(self.extra_options["gpu_offset_y"]),
            # d_cos (__global float32*)
            self.cl_mem["d_cos"].data,
            # d_sin (__global float32*)
            self.cl_mem["d_sin"].data,
            # d_axis  (__global float32*)
            self.cl_mem["d_axes"].data,
            # shared mem (__local float32*)
            self._get_local_mem(),
        )

    def _allocate_textures(self):
        """
        Allocate the texture for the sinogram.
        """
        self.d_sino_tex = self.allocate_texture(self.shape)

    def _init_filter(self, filter_name):
        """Filter initialization

        :param str filter_name: filter name
        """
        self.filter_name = filter_name or "ram-lak"
        self.sino_filter = SinoFilter(
            self.shape,
            ctx=self.ctx,
            filter_name=self.filter_name,
            extra_options=self.extra_options,
        )

    def _get_local_mem(self):
        return pyopencl.LocalMemory(self.local_mem)  # constant for all image sizes

    def _cpy2d_to_slice(self, dst):
        ndrange = (int(self.slice_shape[1]), int(self.slice_shape[0]))
        slice_shape_ocl = np.int32(ndrange)
        wg = None
        kernel_args = (
            dst.data,
            self.cl_mem["_d_slice"].data,
            np.int32(self.slice_shape[1]),
            np.int32(self.dimrec_shape[1]),
            np.int32((0, 0)),
            np.int32((0, 0)),
            slice_shape_ocl,
        )
        return self.kernels.cpy2d(self.queue, ndrange, wg, *kernel_args)

    def _transfer_to_texture(self, sino):
        if isinstance(sino, parray.Array):
            return self._transfer_device_to_texture(sino)
        sino2 = sino
        if not (sino.flags["C_CONTIGUOUS"] and sino.dtype == np.float32):
            sino2 = np.ascontiguousarray(sino, dtype=np.float32)
        if not (self._use_textures):
            ev = pyopencl.enqueue_copy(self.queue, self.d_sino.data, sino2)
            what = "transfer filtered sino H->D buffer"
            ev.wait()
        else:
            ev = pyopencl.enqueue_copy(
                self.queue,
                self.d_sino_tex,
                sino2,
                origin=(0, 0),
                region=self.shape[::-1],
            )
            what = "transfer filtered sino H->D texture"
        return EventDescription(what, ev)

    def _transfer_device_to_texture(self, d_sino):
        if not (self._use_textures):
            if id(self.d_sino) == id(d_sino):
                return
            ev = pyopencl.enqueue_copy(self.queue, self.d_sino.data, d_sino)
            what = "transfer filtered sino D->D buffer"
            ev.wait()
        else:
            ev = pyopencl.enqueue_copy(
                self.queue,
                self.d_sino_tex,
                d_sino.data,
                offset=0,
                origin=(0, 0),
                region=self.shape[::-1],
            )
            what = "transfer filtered sino D->D texture"
        return EventDescription(what, ev)

    def backprojection(self, sino, output=None):
        """Perform the backprojection on an input sinogram

        :param sino: sinogram.
        :param output: optional, output slice.
            If provided, the result will be written in this array.
        :return: backprojection of sinogram
        """
        events = []
        with self.sem:
            events.append(self._transfer_to_texture(sino))
            # Call the backprojection kernel
            if not (self._use_textures):
                kernel_to_call = self.kernels.backproj_cpu_kernel
            else:
                kernel_to_call = self.kernels.backproj_kernel
            kernel_to_call(
                self.queue, self.ndrange, self.wg, *self._backproj_kernel_args
            )
            # Return
            if output is None:
                res = self.cl_mem["_d_slice"].get()
                res = res[: self.slice_shape[0], : self.slice_shape[1]]
            else:
                res = output
                self._cpy2d_to_slice(output)

        # /with self.sem
        if self.profile:
            self.events += events

        return res

    def filtered_backprojection(self, sino, output=None):
        """
        Compute the filtered backprojection (FBP) on a sinogram.

        :param sino: sinogram (`np.ndarray` or `pyopencl.array.Array`)
            with the shape (n_projections, n_bins)
        :param output: output (`np.ndarray` or `pyopencl.array.Array`).
            If nothing is provided, a new numpy array is returned.
        """
        # Filter
        self.sino_filter(sino, output=self.d_sino)
        # Backproject
        res = self.backprojection(self.d_sino, output=output)
        return res

    __call__ = filtered_backprojection
