# -*- coding: utf-8 -*-

# PyVkFFT
#   (c) 2021- : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import ctypes
import os
from os import name as os_name

import numpy as np
from vulkan import ffi

# from pyvkfft.tune import tune_vkfft
##import platform
##import sysconfig


def getVulkanPtr(obj):
    return int(ffi.cast("unsigned long long", obj))


# Type to pass array size, omit and batch as int arrays
# vkfft_path = os.path.join(os.path.dirname(__file__), '..', 'src', 'vulkan')
vkfft_path = os.path.join(os.path.dirname(__file__), "bin")
if os_name == "nt":  # windows
    vkfft_path = os.path.join(vkfft_path, "vkfft_vulkan.dll")
    vkfft_long_type = np.int32
else:  # unix
    vkfft_path = os.path.join(vkfft_path, "vkfft_vulkan.so")
    vkfft_long_type = np.int64

# Regrettably, VkFFT was compiled using "long" types, which have different sizes on Windows/Linux; we patch this here.
# TODO: recompile VkFFT replacing "long" with "int" (always 32 bit) or "long long' (always 64 bit).
ctype_int_size_p = np.ctypeslib.ndpointer(
    dtype=vkfft_long_type, ndim=1, flags="C_CONTIGUOUS"
)


from radis.gpu.vulkan.pyvkfft_base import (
    VkFFTApp as VkFFTAppBase,  # , check_vkfft_result
)

# from pyvkfft_base import VkFFTApp as VkFFTAppBase  # , check_vkfft_result


##def load_library(basename):
##    if platform.system() == "Windows":
##        # We patched build_ext so the module is a .so and not a dll
##        ext = ".so"
##    else:
##        ext = sysconfig.get_config_var("EXT_SUFFIX")
##    return ctypes.cdll.LoadLibrary(
##        os.path.join(os.path.dirname(__file__) or os.path.curdir, basename + ext)
##    )

_vkfft_vulkan = ctypes.cdll.LoadLibrary(vkfft_path)


##stdout = os.dup(1)
##silent = os.open(os.devnull, os.O_WRONLY)


def prepare_fft(
    buffer, name="", norm=1, compute_app=None, indirectOffset=0, exclusivePlan=0
):

    return VkFFTApp(
        buffer=buffer,
        indirectOffset=indirectOffset,
        indirectBuffer=compute_app.indirect_d,
        indirectHost=compute_app._indirect_h,
        physical_device=compute_app._physicalDevice,
        device=compute_app._device,
        queue=compute_app._queue,
        command_pool=compute_app._commandPool,
        fence=compute_app._fence,
        ndim=1,
        inplace=True,
        norm=norm,
        r2c=True,
        tune_config=None,
        name=name,
        exclusivePlan=exclusivePlan,
    )


class _types:
    """Aliases"""

    VkFFTConfiguration = ctypes.c_void_p
    VkBuffer = ctypes.c_void_p
    VkPhysicalDevice = ctypes.c_void_p
    VkDevice = ctypes.c_void_p
    VkQueue = ctypes.c_void_p
    VkCommandPool = ctypes.c_void_p
    VkFence = ctypes.c_void_p
    VkCommandBuffer = ctypes.c_void_p
    vkfft_app = ctypes.c_void_p


_vkfft_vulkan.get_dev_props.restype = ctypes.c_int
_vkfft_vulkan.get_dev_props.argtypes = [ctypes.c_void_p, ctypes.c_char_p]

_vkfft_vulkan.get_dev_props2.restype = ctypes.c_int
_vkfft_vulkan.get_dev_props2.argtypes = [_types.VkFFTConfiguration, ctypes.c_char_p]

_vkfft_vulkan.sync_app.restype = ctypes.c_int
_vkfft_vulkan.sync_app.argtypes = [_types.vkfft_app]


_vkfft_vulkan.make_config.restype = _types.VkFFTConfiguration
_vkfft_vulkan.make_config.argtypes = [
    ctype_int_size_p,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_size_t,
    _types.VkBuffer,
    _types.VkBuffer,
    # ctypes.c_int,
    # _types.VkBuffer,
    # ctypes.c_int,
    ctypes.c_int,
    _types.VkBuffer,
    ctypes.c_int,
    ctypes.c_void_p,
    ctypes.POINTER(_types.VkPhysicalDevice),
    ctypes.POINTER(_types.VkDevice),
    ctypes.POINTER(_types.VkQueue),
    ctypes.POINTER(_types.VkCommandPool),
    ctypes.POINTER(_types.VkFence),
    ctypes.c_ulonglong,
    ctypes.c_int,
    ctypes.c_size_t,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_size_t,
    ctype_int_size_p,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_int,
    ctype_int_size_p,
    ctypes.c_char_p,
    ctypes.c_int,
]

_vkfft_vulkan.init_app.restype = _types.vkfft_app
_vkfft_vulkan.init_app.argtypes = [
    _types.VkFFTConfiguration,
    ctypes.POINTER(ctypes.c_int),
]

_vkfft_vulkan.fft.restype = ctypes.c_int
_vkfft_vulkan.fft.argtypes = [
    _types.vkfft_app,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
]

_vkfft_vulkan.ffto.restype = ctypes.c_int
_vkfft_vulkan.ffto.argtypes = [
    _types.vkfft_app,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_int,
    ctypes.c_int,
]

_vkfft_vulkan.ifft.restype = ctypes.c_int
_vkfft_vulkan.ifft.argtypes = [
    _types.vkfft_app,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
]

_vkfft_vulkan.iffto.restype = ctypes.c_int
_vkfft_vulkan.iffto.argtypes = [
    _types.vkfft_app,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_int,
    ctypes.c_int,
]

_vkfft_vulkan.free_app.restype = None
_vkfft_vulkan.free_app.argtypes = [_types.vkfft_app]

_vkfft_vulkan.free_config.restype = None
_vkfft_vulkan.free_config.argtypes = [_types.VkFFTConfiguration]

_vkfft_vulkan.vkfft_max_fft_dimensions.restype = ctypes.c_uint32
_vkfft_vulkan.vkfft_max_fft_dimensions.argtypes = None


class VkFFTApp(VkFFTAppBase):
    """
    VkFFT application interface, similar to a cuFFT plan.
    """

    def __init__(
        self,
        buffer,
        # shape_in,
        # dtype: type,
        # buffer_size,
        # buffer_src,
        # buffer_dst,
        # currentBatchUBO,
        # currentBatchUBOOffset,
        indirectOffset,
        indirectBuffer,
        indirectHost,
        physical_device,
        device,
        queue,
        command_pool,
        fence,
        ndim=None,
        inplace=True,
        norm=1,
        r2c=False,
        dct=False,
        axes=None,
        # strides=None,
        tune_config=None,
        name="",
        exclusivePlan=0,
        **kwargs,
    ):

        shape_in = (
            (buffer._fftSize,)
            if buffer._batchSize == 1
            else (buffer._batchSize, buffer._fftSize)
        )
        dtype = buffer._dtype
        super().__init__(
            shape_in,
            dtype,
            ndim=ndim,
            inplace=inplace,
            norm=norm,
            r2c=r2c,
            dct=dct,
            axes=axes,
            **kwargs,
        )

        self.bufferSize = buffer._bufferSize
        self.bufferSrc = _types.VkBuffer(getVulkanPtr(buffer._buffer))
        self.bufferDest = _types.VkBuffer(getVulkanPtr(buffer._buffer))
        self.indirectOffset = indirectOffset
        self.indirectBuffer = _types.VkBuffer(getVulkanPtr(indirectBuffer._buffer))
        self.indirectHost = indirectBuffer._structPtr.contents
        self.physicalDevice = _types.VkPhysicalDevice(getVulkanPtr(physical_device))
        self.device = _types.VkDevice(getVulkanPtr(device))
        self.queue = _types.VkQueue(getVulkanPtr(queue))
        self.commandPool = _types.VkCommandPool(getVulkanPtr(command_pool))
        self.fence = _types.VkFence(getVulkanPtr(fence))
        self.name = name.encode()
        self.exclusivePlan = exclusivePlan
        # buf = ctypes.create_string_buffer(256)

        self.config = self._make_config()

        if self.config is None:
            raise RuntimeError(
                "Error creating VkFFTConfiguration. Was the Vulkan context properly initialised ?"
            )
        res = ctypes.c_int(0)

        self.app = _vkfft_vulkan.init_app(self.config, ctypes.byref(res))

        if self.app is None:
            raise RuntimeError(
                "Error {:d}  creating VkFFTApplication. Was the Vulkan driver initialised ?".format(
                    res.value
                )
            )

    def __del__(self):
        """Takes care of deleting allocated memory in the underlying
        VkFFTApplication and VkFFTConfiguration.
        """

        # print('VkFFT.app.__del__... ')
        # print('free_app @',hex(self.app), end='... ')
        if self.app is not None:
            _vkfft_vulkan.free_app(self.app)
            self.app = None
        # print('Done!')

        # print('free_config @',hex(self.config), end='... ')
        if self.config is not None:
            _vkfft_vulkan.free_config(self.config)
            self.config = None
        # print('Done!')

        # print('VkFFT.app.__del__ done!')

    def _make_config(self):
        """Create a vkfft configuration for a FFT transform"""
        if len(self.shape) > vkfft_max_fft_dimensions():
            raise RuntimeError(
                f"Too many FFT dimensions after collapsing non-transform axes: "
                f"{len(self.shape)}>{vkfft_max_fft_dimensions()}"
            )

        shape = np.ones(vkfft_max_fft_dimensions(), dtype=vkfft_long_type)
        skip = np.zeros(vkfft_max_fft_dimensions(), dtype=vkfft_long_type)
        grouped_batch = -np.ones(vkfft_max_fft_dimensions(), dtype=vkfft_long_type)

        shape[0] = self.shape[-1]
        # skip[1 : len(self.shape)] = 1
        # FFTdim = 1
        n_batch = 1 if len(self.shape) == 1 else self.shape[-2]

        indirectDispatch = 1 if self.exclusivePlan == 1 else 0

        return _vkfft_vulkan.make_config(
            shape,
            self.bufferSize,
            0,  # self.bufferSize,
            1,  # FFTdim,
            self.bufferSrc,
            None,  # self.bufferDest,
            indirectDispatch,
            self.indirectBuffer,
            self.indirectOffset,
            ctypes.byref(self.indirectHost),
            ctypes.byref(self.physicalDevice),
            ctypes.byref(self.device),
            ctypes.byref(self.queue),
            ctypes.byref(self.commandPool),
            ctypes.byref(self.fence),
            0,  # isCompilerInitialized
            1,  # norm,
            self.precision,
            1,  # int(self.r2c),
            0,  # int(self.dct),
            0,  # int(self.disableReorderFourStep),
            int(self.registerBoost),
            int(self.use_lut),
            int(self.keepShaderCode),
            # min(6,n_batch),
            n_batch,
            skip,
            int(self.coalescedMemory),
            int(self.numSharedBanks),
            int(self.aimThreads),
            int(self.performBandwidthBoost),
            int(self.registerBoostNonPow2),
            int(self.registerBoost4Step),
            int(self.warpSize),
            0,  # 1,#int(1), #specify offset at launch
            grouped_batch,
            self.name,
            self.exclusivePlan,
            0,  # debugEnable
        )

    def getBufSize(self, src):
        self.bufferSrc = _types.VkBuffer(getVulkanPtr(src))
        m = _vkfft_vulkan.get_buf_size(self.bufferSrc, ctypes.byref(self.device))
        print("Returned buffer size: ", m)

    # def sync(self):
    # res = _vkfft_vulkan.sync_app(self.app)

    def fft(self, cmd_buf, src):
        """
        Compute the forward FFT

        :return: the transformed array. For a R2C inplace transform, the complex view of the
            array is returned.
        """

        self.bufferSrc = _types.VkBuffer(getVulkanPtr(src))
        self.commandBufferFwd = _types.VkCommandBuffer(getVulkanPtr(cmd_buf))

        _vkfft_vulkan.fft(
            self.app,
            ctypes.byref(self.commandBufferFwd),
            ctypes.byref(self.bufferSrc),
            None,
        )

    def ifft(self, cmd_buf, src):
        """
        Compute the backward FFT

        :return: the transformed array. For a C2R inplace transform, the float view of the
            array is returned.
        """

        self.bufferSrc = _types.VkBuffer(getVulkanPtr(src))
        self.commandBufferRev = _types.VkCommandBuffer(getVulkanPtr(cmd_buf))

        _vkfft_vulkan.ifft(
            self.app,
            ctypes.byref(self.commandBufferRev),
            None,
            ctypes.byref(self.bufferSrc),
        )

    def setFFTWorkGroupSize(self, N):
        for wg in self.indirectHost:
            # print(wg.id, ':', wg.x, wg.y, wg.z, '->', end=' ')
            if wg.id > 0:
                try:
                    ax = "xyz"[wg.id]
                    setattr(wg, ax, N)
                    # print(wg.id)
                except (IndexError):
                    # the wg.id is set by VkFFT to either 0, 1, or 2,
                    # to specify which axis contains the batch number.
                    # this exception is to catch issues with the wg.id being any other number.
                    # However, this issue was likely solved by zeroing the values in
                    # vulkan_compute_lib.setIndirectBuffer().
                    print("Vulkan Workgroup ID error:", wg.id)
            # print(wg.x, wg.y, wg.z)


def vkfft_version():
    """
    Get VkFFT version
    :return: version as X.Y.Z
    """
    int_ver = _vkfft_vulkan.vkfft_version()
    return "%d.%d.%d" % (int_ver // 10000, (int_ver % 10000) // 100, int_ver % 100)


def vkfft_max_fft_dimensions():
    """
    Get the maximum number of dimensions VkFFT can handle. This is
    set at compile time. VkFFT default is 4, pyvkfft sets this to 8.
    Note that consecutive non-transformed are collapsed into a single
    axis, reducing the effective number of dimensions.

    :return: VKFFT_MAX_FFT_DIMENSIONS
    """
    return _vkfft_vulkan.vkfft_max_fft_dimensions()
