#!/usr/bin/env python
#
#    Project: Sift implementation in Python + OpenCL
#             https://github.com/silx-kit/silx
#
#    Copyright (C) 2013-2023  European Synchrotron Radiation Facility, Grenoble, France
#
# 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.

"""
This module provides a class for CBF byte offset compression/decompression.
"""

__authors__ = ["Jérôme Kieffer"]
__contact__ = "jerome.kieffer@esrf.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "11/10/2018"
__status__ = "production"


import functools
import os
import numpy
from ..common import ocl, pyopencl
from ..processing import BufferDescription, EventDescription, OpenclProcessing

import logging

logger = logging.getLogger(__name__)

if pyopencl:
    import pyopencl.version

    if pyopencl.version.VERSION < (2016, 0):
        from pyopencl.scan import GenericScanKernel, GenericDebugScanKernel
    else:
        from pyopencl.algorithm import GenericScanKernel
        from pyopencl.scan import GenericDebugScanKernel
else:
    logger.warning("No PyOpenCL, no byte-offset, please see fabio")


class ByteOffset(OpenclProcessing):
    """Perform the byte offset compression/decompression on the GPU

    See :class:`OpenclProcessing` for optional arguments description.

    :param int raw_size:
        Size of the raw stream for decompression.
        It can be (slightly) larger than the array.
    :param int dec_size:
        Size of the decompression output array
        (mandatory for decompression)
    """

    def __init__(
        self,
        raw_size=None,
        dec_size=None,
        ctx=None,
        devicetype="all",
        platformid=None,
        deviceid=None,
        block_size=None,
        profile=False,
    ):
        OpenclProcessing.__init__(
            self,
            ctx=ctx,
            devicetype=devicetype,
            platformid=platformid,
            deviceid=deviceid,
            block_size=block_size,
            profile=profile,
        )
        if self.block_size is None:
            self.block_size = self.device.max_work_group_size
        wg = self.block_size

        buffers = [BufferDescription("counter", 1, numpy.int32, None)]

        if raw_size is None:
            self.raw_size = -1
            self.padded_raw_size = -1
        else:
            self.raw_size = int(raw_size)
            self.padded_raw_size = int((self.raw_size + wg - 1) & ~(wg - 1))
            buffers += [
                BufferDescription("raw", self.padded_raw_size, numpy.int8, None),
                BufferDescription("mask", self.padded_raw_size, numpy.int32, None),
                BufferDescription("values", self.padded_raw_size, numpy.int32, None),
                BufferDescription(
                    "exceptions", self.padded_raw_size, numpy.int32, None
                ),
            ]

        if dec_size is None:
            self.dec_size = None
        else:
            self.dec_size = numpy.int32(dec_size)
            buffers += [
                BufferDescription("data_float", self.dec_size, numpy.float32, None),
                BufferDescription("data_int", self.dec_size, numpy.int32, None),
            ]

        self.allocate_buffers(buffers, use_array=True)

        self.compile_kernels([os.path.join("codec", "byte_offset")])
        self.kernels.__setattr__("scan", self._init_double_scan())
        self.kernels.__setattr__("compression_scan", self._init_compression_scan())

    def _init_double_scan(self):
        """generates a double scan on indexes and values in one operation"""
        arguments = "__global int *value", "__global int *index"
        int2 = pyopencl.tools.get_or_register_dtype("int2")
        input_expr = "index[i]>0 ? (int2)(0, 0) : (int2)(value[i], 1)"
        scan_expr = "a+b"
        neutral = "(int2)(0,0)"
        output_statement = "value[i] = item.s0; index[i+1] = item.s1;"

        if self.block_size > 256:
            knl = GenericScanKernel(
                self.ctx,
                dtype=int2,
                arguments=arguments,
                input_expr=input_expr,
                scan_expr=scan_expr,
                neutral=neutral,
                output_statement=output_statement,
            )
        else:  # MacOS on CPU
            knl = GenericDebugScanKernel(
                self.ctx,
                dtype=int2,
                arguments=arguments,
                input_expr=input_expr,
                scan_expr=scan_expr,
                neutral=neutral,
                output_statement=output_statement,
            )
        return knl

    def decode(self, raw, as_float=False, out=None):
        """This function actually performs the decompression by calling the kernels

        :param numpy.ndarray raw: The compressed data as a 1D numpy array of char.
        :param bool as_float: True to decompress as float32,
                              False (default) to decompress as int32
        :param pyopencl.array out: pyopencl array in which to place the result.
        :return: The decompressed image as an pyopencl array.
        :rtype: pyopencl.array
        """
        assert (
            self.dec_size is not None
        ), "dec_size is a mandatory ByteOffset init argument for decompression"

        events = []
        with self.sem:
            len_raw = numpy.int32(len(raw))
            if len_raw > self.padded_raw_size:
                wg = self.block_size
                self.raw_size = int(len(raw))
                self.padded_raw_size = (self.raw_size + wg - 1) & ~(wg - 1)
                logger.info("increase raw buffer size to %s", self.padded_raw_size)
                buffers = {
                    "raw": pyopencl.array.empty(
                        self.queue, self.padded_raw_size, dtype=numpy.int8
                    ),
                    "mask": pyopencl.array.empty(
                        self.queue, self.padded_raw_size, dtype=numpy.int32
                    ),
                    "exceptions": pyopencl.array.empty(
                        self.queue, self.padded_raw_size, dtype=numpy.int32
                    ),
                    "values": pyopencl.array.empty(
                        self.queue, self.padded_raw_size, dtype=numpy.int32
                    ),
                }
                self.cl_mem.update(buffers)
            else:
                wg = self.block_size

            evt = pyopencl.enqueue_copy(
                self.queue, self.cl_mem["raw"].data, raw, is_blocking=False
            )
            events.append(EventDescription("copy raw H -> D", evt))
            evt = self.kernels.fill_int_mem(
                self.queue,
                (self.padded_raw_size,),
                (wg,),
                self.cl_mem["mask"].data,
                numpy.int32(self.padded_raw_size),
                numpy.int32(0),
                numpy.int32(0),
            )
            events.append(EventDescription("memset mask", evt))
            evt = self.kernels.fill_int_mem(
                self.queue,
                (1,),
                (1,),
                self.cl_mem["counter"].data,
                numpy.int32(1),
                numpy.int32(0),
                numpy.int32(0),
            )
            events.append(EventDescription("memset counter", evt))
            evt = self.kernels.mark_exceptions(
                self.queue,
                (self.padded_raw_size,),
                (wg,),
                self.cl_mem["raw"].data,
                len_raw,
                numpy.int32(self.raw_size),
                self.cl_mem["mask"].data,
                self.cl_mem["values"].data,
                self.cl_mem["counter"].data,
                self.cl_mem["exceptions"].data,
            )
            events.append(EventDescription("mark exceptions", evt))
            nb_exceptions = numpy.empty(1, dtype=numpy.int32)
            evt = pyopencl.enqueue_copy(
                self.queue,
                nb_exceptions,
                self.cl_mem["counter"].data,
                is_blocking=False,
            )
            events.append(EventDescription("copy counter D -> H", evt))
            evt.wait()
            nbexc = int(nb_exceptions[0])
            if nbexc == 0:
                logger.info("nbexc %i", nbexc)
            else:
                evt = self.kernels.treat_exceptions(
                    self.queue,
                    (nbexc,),
                    (1,),
                    self.cl_mem["raw"].data,
                    len_raw,
                    self.cl_mem["mask"].data,
                    self.cl_mem["exceptions"].data,
                    self.cl_mem["values"].data,
                )
                events.append(EventDescription("treat_exceptions", evt))

            # self.cl_mem["copy_values"] = self.cl_mem["values"].copy()
            # self.cl_mem["copy_mask"] = self.cl_mem["mask"].copy()
            evt = self.kernels.scan(
                self.cl_mem["values"],
                self.cl_mem["mask"],
                queue=self.queue,
                size=int(len_raw),
                wait_for=(evt,),
            )
            events.append(EventDescription("double scan", evt))
            # evt.wait()
            if out is not None:
                if out.dtype == numpy.float32:
                    copy_results = self.kernels.copy_result_float
                else:
                    copy_results = self.kernels.copy_result_int
            else:
                if as_float:
                    out = self.cl_mem["data_float"]
                    copy_results = self.kernels.copy_result_float
                else:
                    out = self.cl_mem["data_int"]
                    copy_results = self.kernels.copy_result_int
            evt = copy_results(
                self.queue,
                (self.padded_raw_size,),
                (wg,),
                self.cl_mem["values"].data,
                self.cl_mem["mask"].data,
                len_raw,
                self.dec_size,
                out.data,
            )
            events.append(EventDescription("copy_results", evt))
            # evt.wait()
            if self.profile:
                self.events += events
        return out

    __call__ = decode

    def _init_compression_scan(self):
        """Initialize CBF compression scan kernels"""
        preamble = """
        int compressed_size(int diff) {
            int abs_diff = abs(diff);

            if (abs_diff < 128) {
                return 1;
            }
            else if (abs_diff < 32768) {
                return 3;
            }
            else {
                return 7;
            }
        }

        void write(const int index,
                   const int diff,
                   global char *output) {
            int abs_diff = abs(diff);

            if (abs_diff < 128) {
                output[index] = (char) diff;
            }
            else if (abs_diff < 32768) {
                output[index] = -128;
                output[index + 1] = (char) (diff >> 0);
                output[index + 2] = (char) (diff >> 8);
            }
            else {
                output[index] = -128;
                output[index + 1] = 0;
                output[index + 2] = -128;
                output[index + 3] = (char) (diff >> 0);
                output[index + 4] = (char) (diff >> 8);
                output[index + 5] = (char) (diff >> 16);
                output[index + 6] = (char) (diff >> 24);
            }
        }
        """
        arguments = (
            "__global const int *data, __global char *compressed, __global int *size"
        )
        input_expr = "compressed_size((i == 0) ? data[0] : (data[i] - data[i - 1]))"
        scan_expr = "a+b"
        neutral = "0"
        output_statement = """
        if (prev_item == 0) { // 1st thread store compressed data size
            size[0] = last_item;
        }
        write(prev_item, (i == 0) ? data[0] : (data[i] - data[i - 1]), compressed);
        """

        if self.block_size >= 64:
            knl = GenericScanKernel(
                self.ctx,
                dtype=numpy.int32,
                preamble=preamble,
                arguments=arguments,
                input_expr=input_expr,
                scan_expr=scan_expr,
                neutral=neutral,
                output_statement=output_statement,
            )
        else:  # MacOS on CPU
            knl = GenericDebugScanKernel(
                self.ctx,
                dtype=numpy.int32,
                preamble=preamble,
                arguments=arguments,
                input_expr=input_expr,
                scan_expr=scan_expr,
                neutral=neutral,
                output_statement=output_statement,
            )
        return knl

    def encode(self, data, out=None):
        """Compress data to CBF.

        :param data: The data to compress as a numpy array
                     (or a pyopencl Array) of int32.
        :type data: Union[numpy.ndarray, pyopencl.array.Array]
        :param pyopencl.array out:
            pyopencl array of int8 in which to store the result.
            The array should be large enough to store the compressed data.
        :return: The compressed data as a pyopencl array.
                 If out is provided, this array shares the backing buffer,
                 but has the exact size of the compressed data and the queue
                 of the ByteOffset instance.
        :rtype: pyopencl.array
        :raises ValueError: if out array is not large enough
        """

        events = []
        with self.sem:
            if isinstance(data, pyopencl.array.Array):
                d_data = data  # Uses provided array

            else:  # Copy data to device
                data = numpy.ascontiguousarray(data, dtype=numpy.int32).ravel()

                # Make sure data array exists and is large enough
                if (
                    "data_input" not in self.cl_mem
                    or self.cl_mem["data_input"].size < data.size
                ):
                    logger.info("increase data input buffer size to %s", data.size)
                    self.cl_mem.update(
                        {
                            "data_input": pyopencl.array.empty(
                                self.queue, data.size, dtype=numpy.int32
                            )
                        }
                    )
                d_data = self.cl_mem["data_input"]

                evt = pyopencl.enqueue_copy(
                    self.queue, d_data.data, data, is_blocking=False
                )
                events.append(EventDescription("copy data H -> D", evt))

            # Make sure compressed array exists and is large enough
            compressed_size = d_data.size * 7
            if (
                "compressed" not in self.cl_mem
                or self.cl_mem["compressed"].size < compressed_size
            ):
                logger.info("increase compressed buffer size to %s", compressed_size)
                self.cl_mem.update(
                    {
                        "compressed": pyopencl.array.empty(
                            self.queue, compressed_size, dtype=numpy.int8
                        )
                    }
                )
            d_compressed = self.cl_mem["compressed"]
            d_size = self.cl_mem["counter"]  # Shared with decompression

            evt = self.kernels.compression_scan(d_data, d_compressed, d_size)
            events.append(EventDescription("compression scan", evt))
            byte_count = int(d_size.get()[0])

            if out is None:
                # Create out array from a sub-region of the compressed buffer
                out = pyopencl.array.Array(
                    self.queue,
                    shape=(byte_count,),
                    dtype=numpy.int8,
                    allocator=functools.partial(
                        d_compressed.base_data.get_sub_region, d_compressed.offset
                    ),
                )

            elif out.size < byte_count:
                raise ValueError(
                    "Provided output buffer is not large enough: "
                    "requires %d bytes, got %d" % (byte_count, out.size)
                )

            else:  # out.size >= byte_count
                # Create an array with a sub-region of out and this class queue
                out = pyopencl.array.Array(
                    self.queue,
                    shape=(byte_count,),
                    dtype=numpy.int8,
                    allocator=functools.partial(
                        out.base_data.get_sub_region, out.offset
                    ),
                )

                evt = pyopencl.enqueue_copy(
                    self.queue, out.data, d_compressed.data, byte_count=byte_count
                )
                events.append(EventDescription("copy D -> D: internal -> out", evt))

            if self.profile:
                self.events += events

        return out

    def encode_to_bytes(self, data):
        """Compresses data to CBF and returns compressed data as bytes.

        Usage:

        Provided an image (`image`) stored as a numpy array of int32,
        first, create a byte offset compression/decompression object:

        >>> from silx.opencl.codec.byte_offset import ByteOffset
        >>> byte_offset_codec = ByteOffset()

        Then, compress an image into bytes:

        >>> compressed = byte_offset_codec.encode_to_bytes(image)

        :param data: The data to compress as a numpy array
                     (or a pyopencl Array) of int32.
        :type data: Union[numpy.ndarray, pyopencl.array.Array]
        :return: The compressed data as bytes.
        :rtype: bytes
        """
        compressed_array = self.encode(data)
        return compressed_array.get().tobytes()
