import os.path
from ctypes import sizeof

import numpy as np
from scipy.constants import N_A, c, k
from scipy.fft import next_fast_len

from radis.gpu.params import (
    init_G_params,
    init_L_params,
    init_Q,
    set_G_params,
    set_L_params,
    set_pTQ,
)
from radis.gpu.structs import initData_t, iterData_t, workGroupSizeArray_t
from radis.misc.utils import getProjectRoot

shader_path = os.path.join(getProjectRoot(), "gpu", "vulkan", "shaders")


def next_fast_len_even(n):
    # Find an FFT length that can be processed quickly.
    # The length needs to be divisible by 2 because of the
    # length doubling for prevention of circular convolution.
    n = next_fast_len(n)
    while n & 1:
        n = next_fast_len(n + 1)
    return n


from radis.gpu.vulkan.vulkan_compute_lib import GPUApplication, GPUBuffer


class gpuApp(GPUApplication):
    def __init__(
        self,
        vmin,
        Nv,
        dv,
        dxG,
        dxL,
        v0,
        da,
        na,
        S0,
        El,
        gamma_arr,
        iso,
        Mm_arr,
        Q_intp_list,
        verbose=0,
        device_id=0,
    ):

        """
        Initialize gpuApp object required for GPU calculations

        Parameters
        ----------
        vmin : float
            Minimum value frequency/wavenumber axis
        Nv : int
            Total number of frequency/wavenumber points.
        dv : float
            Stepsize of frequency/wavenumber axis (called wstep elsewhere in RADIS).
        dxG : float
            Relative grid spacing for Gaussian lineshapes.
        dxL : float
            Relative grid spacing for Lorentzian lineshapes.
        v0 : numpy.ndarray[np.float32]
            Array of line center frequencies (in cm-1).
        da : numpy.ndarray[np.float32]
            Pressure shift  (in cm-1.atm-1).
        na : numpy.ndarray[np.float32]
            Temperature dependency of Lorentzian widths
        S0 : numpy.ndarray[np.float32]
            Line intensity scaling factors.
        El : numpy.ndarray[np.float32]
            Lower level energy levels.
        gamma_arr : numpy.ndarray[np.float32]
            (m,n) shaped array with Lorentzian width parameters, with n the number of lines in the database
            and m the number of collision partners included. This is usually at least two,
            with the first (m=0) always self broadening and the last (m=-1) always air broadening.
        iso : numpy.ndarray[np.uint32]
            Index of isotopologue.
        Mm_arr : numpy.ndarray
            Molecular masses for all isotopologues in the database (Mm_arr[0] is always 0).
        Q_intp_list : list
            List of Q branch interpolators.
        verbose : bool, optional
            Print verbosity level. Default is 0.
        device_id: int, str
            Select the GPU device. If ``int``, specifies the device index, which is printed for convenience during GPU initialization with backend='vulkan' (default).
            If ``str``, return the first device that includes the specified string (case in-sesitive). If not found, return the device at index 0.
            default = 0

        """

        self.backend = "vulkan"

        # Initialize the vulkan GPUApplication:
        super().__init__(deviceID=device_id, path=shader_path, verbose=verbose)

        ## Next, the GPU is made aware of a number of parameters.
        ## Parameters that don't change during iteration are stored
        ## in init_h.

        ## When using the GPU, we distinguish objects for host (=CPU side) use (_h)
        ## and device (=GPU side) use (_d).

        ## The init_d/init_h pair is created by first making a GPUBuffer object
        ## init_d. Then that object makes the init_h object, which guarantees
        ## that data written to init_h can be transferred to init_d.

        ## init_d is a Uniform buffer, which means that the contents don't change
        ## during execution of the shader (=GPU code). This means that iter_d
        ## is also a Uniform buffer, because it only changes before/after execution
        ## of the shader.

        ## Radis configures Uniform buffers as using memory that is visible to both
        ## host (CPU) and device (GPU). For this reason, no explicit transfers between
        ## host and device memory are required.

        if verbose >= 2:
            print("Copying initialization parameters to device memory...")

        self.init_d = GPUBuffer(
            sizeof(initData_t), usage="uniform", binding=0
        )  # host_visible, device_local, (host_cached)

        self.init_h = self.init_d.getHostStructPtr(initData_t)

        self.init_h.v_min = vmin
        self.init_h.dv = dv
        self.init_h.N_v = Nv
        self.init_h.N_v_FT = next_fast_len_even(2 * self.init_h.N_v)
        self.init_h.N_x_FT = self.init_h.N_v_FT // 2 + 1
        self.init_h.dxG = dxG
        self.init_h.dxL = dxL
        self.init_h.N_lines = int(len(v0))
        self.init_h.N_coll = gamma_arr.shape[0]

        log_c2Mm_arr = np.array(
            [0]
            + [
                0.5 * np.log(8 * k * np.log(2) / (c**2 * Mm * 1e-3 / N_A))
                for Mm in Mm_arr[1:]
            ]
        )
        for i in range(len(log_c2Mm_arr)):
            self.init_h.log_c2Mm[i] = log_c2Mm_arr[i]

        init_Q(Q_intp_list)
        log_2vMm = np.log(v0) + log_c2Mm_arr.take(iso)

        init_G_params(log_2vMm.astype(np.float32), verbose)
        init_L_params(na, gamma_arr, verbose)  # TODO: do this on GPU?

        # Parameters that *do* change are stored in iter_h.
        # We only assign values to them in the iterate() method,
        # but already allocate the buffers here.

        self.iter_d = GPUBuffer(
            sizeof(iterData_t), usage="uniform", binding=1
        )  # host_visible, device_local, (host_cached)
        self.iter_h = self.iter_d.getHostStructPtr(iterData_t)

        if verbose >= 2:
            print("done!")

        if verbose >= 2:
            print("Allocating device memory and copying data...")

        ## Next the database memory is initialized. The database is potentially the largest buffer,
        ## Which could take up gigabytes of device memory. Radis tells the Vulkan API to allocate
        ## this buffer as 'storage' buffer. In dedicated GPU's, this is the largest memory heap,
        ## but usually not directly visible to the host. In order to copy memory from device to the host,
        ## internally a staging buffer is created. the staging buffer is small host memory where data is temporarily copied to,
        ## which can be used to transfer data from host to device. This transfer thus consists of two separate transfers:
        ## 1. transfer data from RAM to staging buffer, 2. transfer data from staging buffer to device memory.
        ## The staging buffer is typically smaller than the data to-be-transferred, so this process repeats until all data is copied.
        ## The staging buffer size is given by the chunksize keyword, which is only used the first time when the buffer is initialized, and ignored
        ## during consecutive calls.
        ## In the code below we copy multiple arrays to the device, so for each array the staging buffer transfers are repeated internally until each array is fully copied.

        ## By using database_d.copyToBuffer(), the transfer from staging buffer to device buffer happens immediately.
        ## Alternatively, one could split up the transfer by first using .fromArray() to copy data to the staging buffer, and then
        ## call .transferStagingBuffer() separately. We will use this later for buffers that are copied every iteration.

        database_arrays = [iso, v0, da, S0, El, na, gamma_arr]
        N_db = np.sum(
            [
                np.sum(arr.shape[:-1]) if len(arr.shape) > 1 else 1
                for arr in database_arrays
            ]
        )
        self.database_d = GPUBuffer(N_db * v0.nbytes, usage="storage", binding=2)

        byte_offset = 0
        for arr in database_arrays:
            byte_offset += self.database_d.copyToBuffer(
                arr, device_offset=byte_offset, chunksize=32 * 1024 * 1024
            )

        ## The S_klm_d buffer is purely a device buffer, which stores the result of the lineshape-distribution algorithm
        ## before the FT is applied. Due to changing N_G and N_L per iteration, this buffer may also need to change size.
        ## Because allocating new buffers is costly and may result in fragmentation of device memory, we only change the buffer size
        ## if it needs to be increased. When we do, we increase it by some margin (default is 50% extra), so as to not have to re-allocate everytime
        ## the N_G*N_L size increases. Increasing the buffer size will happen during the iteration step,
        ## but to allow for dynamic resizing we specify it as a FFT buffer by using fftSize. fftSize specifies the size of the FT, which may be
        ## smaller than the actual buffer size. In fact, in order to use in-place transforms with VkFFT, the length of the buffer is that of the reverse FT,
        ## i.e. N_v_FT//2+1 of type complex64 (8 byte), instead of N_v_FT of type float32 (4 byte). This means that the buffer size is 2 elements larger per
        ## nu-axis.

        ## The S_klm_d buffer is (implicitly) created with batchSize = 1. During iteration the batchSize will be set to N_G * N_L.

        self.S_klm_d = GPUBuffer(
            fftSize=self.init_h.N_v_FT, usage="storage", binding=3
        )  # req: large, device_local

        ## The spectrum_d buffer receives the resulting spectrum. This is a storage buffer like S_klm_d, but does not have to change size dynamically.
        ## It *does* have to transfer data from device buffers to host buffers, and since it's contents change during shader execution, we cannot use Uniform for this.
        ## To make this transfer as efficient as possible, the device-to-host (staging buffe) transfer is recorded in the command buffer (see ahead),
        ## while the staging buffer to RAM is done in python code after the comman buffer is done executing.

        self.spectrum_d = GPUBuffer(
            fftSize=self.init_h.N_v_FT, usage="storage", binding=4
        )  # req: device_local, host_visible, host_cached, (large)

        ## The indirect_d buffer is a special buffer that is used to specify the number of workgroups (i.e number of instances of a shader thread)
        ## dynamically. This is needed because the forward FFT will change size as N_G or N_L change. To address this, the FFT plan is created with
        ## the batchSize number of the highest N_G*N_L*margin(=1.5), but less threads are launched depending on the current value of N_G*N_L.
        ## We tell the gpuApp specifically that this is our indirect buffer, so that VkFFT can fill it with the workgroup sizes it would need for the batchSize
        ## it was created for. We then every iteration set the workgroup sizes to the right number, so that we don't perform superfluous FFT's.

        self.indirect_d = GPUBuffer(
            sizeof(workGroupSizeArray_t), usage="indirect", binding=10
        )  # host_visible, device_local, (host_cached)
        self.setIndirectBuffer(self.indirect_d, workGroupSizeArray_t)

        # Write command buffer:
        N_tpb = 128  # threads per block
        threads = (N_tpb, 1, 1)

        # Finally it is time to write the command buffer. The command buffer is a list of commands that the GPU needs to execute.
        # Timestamp labels can be added as desired, to help with benchmarking GPU performance. The command buffer is written at this step,
        # but will not be ran until the iterate() stage.

        # The command buffer starts with copying the new iter_d and indirect_d data from the staging buffer to the device.
        # This is however not needed explicitly, because we took care to use memory that is already device_local.
        # Then the S_klm_d and spectrum_d buffers are zeroed. The "cmdFillLDM.spv" shader is executed, which is a separate file in the shader folder.
        # This shader populates the S_klm_d array. Next the forward FFT is performed, followed by cmdApplyLineshapes.spv (another user defined shader),
        # and finally the contents of spectrum_d are transferred to the staging buffer on the host side again.

        self.appendCommands(
            [
                # self.cmdAddTimestamp("Transfer staging buffers"),
                # self.indirect_d.cmdTransferStagingBuffer("H2D"),
                # self.iter_d.cmdTransferStagingBuffer("H2D"),
                self.cmdAddTimestamp("Zeroing buffers"),
                self.S_klm_d.cmdClearBuffer(),
                self.spectrum_d.cmdClearBuffer(),
                self.cmdAddTimestamp("Line addition"),
                self.cmdScheduleShader(
                    "cmdFillLDM.spv", (self.init_h.N_lines // N_tpb + 1, 1, 1), threads
                ),
                self.cmdAddTimestamp("FFT fwd"),
                self.cmdFFT(self.S_klm_d, name="FFTa"),
                self.cmdAddTimestamp("selfly conv."),
                self.cmdScheduleShader(
                    "cmdApplyLineshapes.spv",
                    (self.init_h.N_x_FT // N_tpb + 1, 1, 1),
                    threads,
                ),
                # self.cmdScheduleShader('cmdTestApplyLineshapesP.spv', (self.init_h.N_x_FT // N_tpb + 1, N_G*N_L, 1), threads),
                self.cmdAddTimestamp("FFT inv"),
                self.cmdIFFT(self.spectrum_d, name="FFTb"),
                self.cmdAddTimestamp("Tranfer staging buffer (result)"),
                self.spectrum_d.cmdTransferStagingBuffer("D2H"),
                self.cmdAddTimestamp("End"),
            ]
        )

        self.writeCommandBuffer()

        ## During iteration, we only run .setBatchSize() to update the Vulkan buffers.
        ## Here we specify what needs to happen when .setBatchSize() is called.
        ## that is two things: the size of S_klm_d needs to be updated, and the workgroup size (indirect_d buffer contents)
        ## must be updated.

        ## If S_klm_d.setBatchSize() needs to increase the buffer size, the internal forward FFT plan (stored in ._fftAppFwd)
        ## will be removed such that it will be rewritten with the new maximum sizes. Note that this happens much less often
        ## than a change in N_G*N_L, because of the 50% extra margin. When it does happen, the performance of that particular iteration
        ## will drop, but this is only once every so often.

        self.registerBatchSizeUpdateFunction(self.S_klm_d.setBatchSize)
        self.registerBatchSizeUpdateFunction(self.setFwdFFTWorkGroupSize)

        if verbose >= 2:
            print("done!")

    def iterate(
        self,
        p,
        T,
        mole_fraction,
        verbose=0,
        # TODO: for GPU instrument functions (not currently supported):
        l=1.0,
        slit_FWHM=0.0,
    ):
        """
        Parameters
        ----------
        p : float
            pressure [bar]
        T : float
            temperature [K]
        mole_fraction : float

        Other Parameters
        ----------------
        verbose : int, optional
            The default is 0.


        Returns
        -------
        abscoeff_h : numpy.ndarray[np.float32]
            array with absorbtion coefficients in (cm.-1)
        iter_h : radis.gpu.structs.iterData_t
            structue with parameters used for computation of abscoeff_h.
        times : dict
            dictionary with computation cumulative computation times for
            different stages of the GPU computation. The ``'total'`` key
            gives the total time.
        """

        if verbose >= 2:
            print("Copying iteration parameters to device...")

        ## Update the contents of iter_h, which as a device_local buffer also
        ## immediately updates the GPU side buffer

        set_pTQ(p, T, mole_fraction, self.iter_h, l=l, slit_FWHM=slit_FWHM)
        set_G_params(self.init_h, self.iter_h)
        set_L_params(self.init_h, self.iter_h)

        if verbose >= 2:
            print("Running compute pipelines...")

        ## Update the batch size, which internally re-allocates buffers and re-plan the FFT if required,
        ## as well as update the workgroup size in the indirect_d buffer
        self.setBatchSize(self.iter_h.N_G * self.iter_h.N_L)

        ## RUN THE COMMAND BUFFER
        self.run()

        ## get the timestamp values
        gpu_times = self.get_timestamps()

        ## Copy the results into a numpy array. Remember that the device-to-host transfer is recorded in the command buffer,
        ## here we only have to copy the contents of the staging buffer to a numpy array using spectrum_d.toArray().
        abscoeff_h = np.zeros(self.init_h.N_v, dtype=np.float32)
        self.spectrum_d.toArray(abscoeff_h)

        if verbose == 1:
            print("Finished calculating spectrum!")

        return abscoeff_h, gpu_times

    def get_griddims(self):
        # use separate getter so as not to increase the app.iter_h reference count
        return self.iter_h.N_L, self.iter_h.N_G

    def __del__(self):
        # print('>>> Deleting gpuApp...')
        self.init_h = None
        self.iter_h = None
        self.free()
