/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
   Class to represent an array in GPU and CPU memory
*/

#include <stdio.h>
#include <stdlib.h>
#define HELPERGL_EXTERN_GL_FUNC_IMPLEMENTATION
#include <cuda_gl_interop.h>
#include <cuda_runtime.h>
#include <helper_cuda.h>
#include <helper_gl.h>

template <class T> class GpuArray
{
public:
    GpuArray();
    ~GpuArray();

    enum Direction {
        HOST_TO_DEVICE,
        DEVICE_TO_HOST,
    };

    // allocate and free
    void alloc(size_t size, bool vbo = false, bool doubleBuffer = false, bool useElementArray = false);
    void free();

    // swap buffers for double buffering
    void swap();

    // when using vbo, must map before getting device ptr
    void map();
    void unmap();

    void copy(Direction dir, uint start = 0, uint count = 0);
    void memset(T value, uint start = 0, uint count = 0);

    T     *getDevicePtr() { return m_dptr[m_currentRead]; }
    GLuint getVbo() { return m_vbo[m_currentRead]; }

    T     *getDeviceWritePtr() { return m_dptr[m_currentWrite]; }
    GLuint getWriteVbo() { return m_vbo[m_currentWrite]; }

    T *getHostPtr() { return m_hptr; }

    size_t getSize() const { return m_size; }

private:
    GLuint createVbo(size_t size, bool useElementArray);

    void allocDevice();
    void allocVbo(bool useElementArray);
    void allocHost();

    void freeDevice();
    void freeVbo();
    void freeHost();

    size_t                       m_size;
    T                           *m_dptr[2];
    GLuint                       m_vbo[2];
    struct cudaGraphicsResource *m_cuda_vbo_resource[2]; // handles OpenGL-CUDA exchange

    T *m_hptr;

    bool m_useVBO;
    bool m_doubleBuffer;
    uint m_currentRead, m_currentWrite;
};

template <class T>
GpuArray<T>::GpuArray()
    : m_size(0)
    , m_hptr(0)
    , m_currentRead(0)
    , m_currentWrite(0)
{
    m_dptr[0] = 0;
    m_dptr[1] = 0;

    m_vbo[0] = 0;
    m_vbo[1] = 0;

    m_cuda_vbo_resource[0] = NULL;
    m_cuda_vbo_resource[1] = NULL;
}

template <class T> GpuArray<T>::~GpuArray() { free(); }

template <class T> void GpuArray<T>::alloc(size_t size, bool vbo, bool doubleBuffer, bool useElementArray)
{
    m_size = size;

    m_useVBO       = vbo;
    m_doubleBuffer = doubleBuffer;

    if (m_doubleBuffer) {
        m_currentWrite = 1;
    }

    allocHost();

    if (vbo) {
        allocVbo(useElementArray);
    }
    else {
        allocDevice();
    }
}

template <class T> void GpuArray<T>::free()
{
    freeHost();

    if (m_vbo) {
        freeVbo();
    }

    if (m_dptr) {
        freeDevice();
    }
}

template <class T> void GpuArray<T>::allocHost() { m_hptr = (T *)new T[m_size]; }

template <class T> void GpuArray<T>::freeHost()
{
    if (m_hptr) {
        delete[] m_hptr;
        m_hptr = 0;
    }
}

template <class T> void GpuArray<T>::allocDevice()
{
    checkCudaErrors(cudaMalloc((void **)&m_dptr[0], m_size * sizeof(T)));

    if (m_doubleBuffer) {
        checkCudaErrors(cudaMalloc((void **)&m_dptr[1], m_size * sizeof(T)));
    }
}

template <class T> void GpuArray<T>::freeDevice()
{
    if (m_dptr[0]) {
        checkCudaErrors(cudaFree(m_dptr[0]));
        m_dptr[0] = 0;
    }

    if (m_dptr[1]) {
        checkCudaErrors(cudaFree(m_dptr[1]));
        m_dptr[1] = 0;
    }
}

template <class T> GLuint GpuArray<T>::createVbo(size_t size, bool useElementArray)
{
    GLuint vbo;
    glGenBuffers(1, &vbo);

    if (useElementArray) {
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, vbo);
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, size, 0, GL_DYNAMIC_DRAW);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0);
    }
    else {
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        glBufferData(GL_ARRAY_BUFFER, size, 0, GL_DYNAMIC_DRAW);
        glBindBuffer(GL_ARRAY_BUFFER, 0);
    }

    return vbo;
}

template <class T> void GpuArray<T>::allocVbo(bool useElementArray)
{
    m_vbo[0] = createVbo(m_size * sizeof(T), useElementArray);
    checkCudaErrors(cudaGraphicsGLRegisterBuffer(&m_cuda_vbo_resource[0], m_vbo[0], cudaGraphicsMapFlagsWriteDiscard));

    if (m_doubleBuffer) {
        m_vbo[1] = createVbo(m_size * sizeof(T), useElementArray);
        checkCudaErrors(
            cudaGraphicsGLRegisterBuffer(&m_cuda_vbo_resource[1], m_vbo[1], cudaGraphicsMapFlagsWriteDiscard));
    }
}

template <class T> void GpuArray<T>::freeVbo()
{
    if (m_vbo[0]) {
        checkCudaErrors(cudaGraphicsUnregisterResource(m_cuda_vbo_resource[0]));
        glDeleteBuffers(1, &m_vbo[0]);
        m_vbo[0] = 0;
    }

    if (m_vbo[1]) {
        checkCudaErrors(cudaGraphicsUnregisterResource(m_cuda_vbo_resource[1]));
        glDeleteBuffers(1, &m_vbo[1]);
        m_vbo[1] = 0;
    }
}

template <class T> void GpuArray<T>::swap() { std::swap(m_currentRead, m_currentWrite); }

template <class T> void GpuArray<T>::map()
{
    if (m_vbo[0]) {
        checkCudaErrors(cudaGraphicsMapResources(1, &m_cuda_vbo_resource[0], 0));
        size_t num_bytes;
        checkCudaErrors(cudaGraphicsResourceGetMappedPointer((void **)&m_dptr[0], &num_bytes, m_cuda_vbo_resource[0]));
    }

    if (m_doubleBuffer && m_vbo[1]) {
        checkCudaErrors(cudaGraphicsMapResources(1, &m_cuda_vbo_resource[1], 0));
        size_t num_bytes;
        checkCudaErrors(cudaGraphicsResourceGetMappedPointer((void **)&m_dptr[1], &num_bytes, m_cuda_vbo_resource[1]));
    }
}

template <class T> void GpuArray<T>::unmap()
{
    if (m_vbo[0]) {
        checkCudaErrors(cudaGraphicsUnmapResources(1, &m_cuda_vbo_resource[0], 0));
        m_dptr[0] = 0;
    }

    if (m_doubleBuffer && m_vbo[1]) {
        checkCudaErrors(cudaGraphicsUnmapResources(1, &m_cuda_vbo_resource[1], 0));
        m_dptr[1] = 0;
    }
}

template <class T> void GpuArray<T>::copy(Direction dir, uint start, uint count)
{
    if (count == 0) {
        count = (uint)m_size;
    }

    map();

    switch (dir) {
    case HOST_TO_DEVICE:
        checkCudaErrors(cudaMemcpy((void *)(m_dptr[m_currentRead] + start),
                                   (void *)(m_hptr + start),
                                   count * sizeof(T),
                                   cudaMemcpyHostToDevice));
        break;

    case DEVICE_TO_HOST:
        checkCudaErrors(cudaMemcpy((void *)(m_hptr + start),
                                   (void *)(m_dptr[m_currentRead] + start),
                                   count * sizeof(T),
                                   cudaMemcpyDeviceToHost));
        break;
    }

    unmap();
}

template <class T> void GpuArray<T>::memset(T value, uint start, uint count) {}