#ifdef ENABLE_OPTIX
#include <Denoiser.h>
#include <optix_function_table_definition.h>
#include <optix_stubs.h>
#include <cuda_helper.hpp>


__global__ void drawTexture(uint2 screen_size, cudaSurfaceObject_t surface, float4* colorArray) {
    CUDA::ThreadInfo ti;
    if (ti.threadId >= screen_size.x * screen_size.y)
        return;

    unsigned int i = ti.threadId % screen_size.x;
    unsigned int j = ti.threadId / screen_size.x;

    float4 color = colorArray[ti.threadId];

    surf2Dwrite(color, surface, i * sizeof(float4), j);
}

Denoiser::Denoiser() {
    cudaStreamCreate(&stream);
    optixInit();

    OptixDenoiserOptions options = {};
    options.guideAlbedo = 0;
    options.guideNormal = 0;
    options.denoiseAlpha = OPTIX_DENOISER_ALPHA_MODE_COPY;

    checkOptixError(optixDeviceContextCreate(nullptr, nullptr, &optixContext));
    checkOptixError(optixDeviceContextSetLogCallback(optixContext, context_log_cb, nullptr, 3));
    checkOptixError(optixDenoiserCreate(optixContext, OPTIX_DENOISER_MODEL_KIND_LDR, &options, &denoiser));
    OptixDenoiserSizes denoiser_sizes;
    checkOptixError(optixDenoiserComputeMemoryResources(denoiser, 1500, 1500, &denoiser_sizes));

    m_scratch_size = static_cast<uint32_t>(denoiser_sizes.withoutOverlapScratchSizeInBytes);
    m_state_size = static_cast<uint32_t>(denoiser_sizes.stateSizeInBytes);
    CUDA::checkCudaError(cudaMalloc(reinterpret_cast<void**>(&m_scratch), m_scratch_size));
    CUDA::checkCudaError(cudaMalloc(reinterpret_cast<void**>(&m_state), m_state_size));
}

Denoiser::~Denoiser() {
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_scratch)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_state)));

    CUDA::checkCudaError(cudaStreamDestroy(stream));
    checkOptixError(optixDenoiserDestroy(denoiser));
    checkOptixError(optixDeviceContextDestroy(optixContext));
}

void Denoiser::denoiseTexture(cudaSurfaceObject_t surface, float4* data, unsigned int width, unsigned int height) {
    OptixImage2D input = {
            (CUdeviceptr) data,       width, height, (unsigned int) (width * sizeof(float4)), sizeof(float4),
            OPTIX_PIXEL_FORMAT_FLOAT4};

    OptixImage2D output = {
            0, width, height, (unsigned int) (width * sizeof(float4)), sizeof(float4), OPTIX_PIXEL_FORMAT_FLOAT4};
    const uint64_t frame_byte_size = width * height * sizeof(float4);
    CUDA::checkCudaError(cudaMalloc(reinterpret_cast<void**>(&output.data), frame_byte_size));

    OptixDenoiserLayer layer = {input, {}, output, OPTIX_DENOISER_AOV_TYPE_NONE};
    OptixDenoiserGuideLayer m_guideLayer = {};

    checkOptixError(optixDenoiserSetup(denoiser, stream, input.width, input.height, m_state, m_state_size, m_scratch,
                                       m_scratch_size));

    checkOptixError(optixDenoiserInvoke(denoiser, stream, &m_params, m_state, m_state_size, &m_guideLayer, &layer, 1, 0,
                                        0, m_scratch, m_scratch_size));

    drawTexture<<<CUDA::getBlockCount(width * height), CUDA::THREADS_PER_BLOCK>>>(
            make_uint2(input.width, input.height), surface, reinterpret_cast<float4*>(layer.output.data));

    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(layer.output.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(layer.previousOutput.data)));

    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.albedo.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.normal.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.flow.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.flowTrustworthiness.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.previousOutputInternalGuideLayer.data)));
    CUDA::checkCudaError(cudaFree(reinterpret_cast<void*>(m_guideLayer.outputInternalGuideLayer.data)));
}
#endif
