#include "SynthesisPipeline.h"

SynthesisPipeline::SynthesisPipeline(sptr<Msl> net, sptr<Camera> cam, uint nSamples,
                                     glm::vec2 depthRange, int encodeDim, int coordChns,
                                     float enhanceSigma, float enhanceFe)
    : _nRays(cam->res().x * cam->res().y),
      _nSamples(nSamples),
      _enhanceSigma(enhanceSigma),
      _enhanceFe(enhanceFe),
      _cam(cam),
      _inferPipeline(new InferPipeline(net, _nRays, nSamples, depthRange, encodeDim, coordChns)),
//      _enhancement(new Enhancement(cam->res())),
      _rays(new CudaArray<glm::vec3>(_nRays)),
      _colors(new CudaArray<glm::vec4>(_nRays)) {
    _glResultBuffer = _createGlResultBuffer(_nRays);
    _glResultTextures.push_back(_createGlResultTexture(_cam->res()));
}

void SynthesisPipeline::getRays(sptr<CudaArray<glm::vec3>> o_rays, View& view) {

}

void SynthesisPipeline::run(View &view) {
    CudaEvent eStart, eGenRays, eInferred, eEnhanced;

    cudaEventRecord(eStart);

    _genRays(view);

    cudaEventRecord(eGenRays);

    _inferPipeline->run(_colors, _rays, view.t(), true);

    cudaEventRecord(eInferred);

    _enhance();

    cudaEventRecord(eEnhanced);

    CHECK_EX(cudaDeviceSynchronize());

    float timeTotal, timeGenRays, timeInfer, timeEnhance;
    cudaEventElapsedTime(&timeTotal, eStart, eEnhanced);
    cudaEventElapsedTime(&timeGenRays, eStart, eGenRays);
    cudaEventElapsedTime(&timeInfer, eGenRays, eInferred);
    cudaEventElapsedTime(&timeEnhance, eInferred, eEnhanced);
    {
        std::ostringstream sout;
        sout << typeid(*this).name() << " => Total: " << timeTotal
             << "ms (Gen rays: " << timeGenRays << "ms, Infer: " << timeInfer
             << "ms, Enhance: " << timeEnhance << "ms)";
        Logger::instance.info(sout.str());
    }

    // Copy result from Cuda array to OpenGL buffer
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, _glResultBuffer);
    void *bufferData = glMapBuffer(GL_PIXEL_UNPACK_BUFFER, GL_WRITE_ONLY);
    cudaMemcpy(bufferData, _colors->getBuffer(), _colors->size(), cudaMemcpyDeviceToHost);
    glUnmapBuffer(GL_PIXEL_UNPACK_BUFFER);
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);

    _uploadResultToTextures();
}

GLuint SynthesisPipeline::getGlResultTexture(int index) { return _glResultTextures[index]; }

void SynthesisPipeline::_genRays(View &view) { view.transVectors(_rays, _cam->localRays()); }

void SynthesisPipeline::_enhance() { _enhancement->run(_colors); }

void SynthesisPipeline::_uploadResultToTextures() {
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, _glResultBuffer);
    glBindTexture(GL_TEXTURE_2D, _glResultTextures[0]);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, _cam->res().x, _cam->res().y, GL_RGBA, GL_FLOAT, 0);
    glBindTexture(GL_TEXTURE_2D, 0);
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
}

GLuint SynthesisPipeline::_createGlResultTexture(glm::uvec2 res) {
    GLuint textureID;
    glEnable(GL_TEXTURE_2D);
    glGenTextures(1, &textureID);
    glBindTexture(GL_TEXTURE_2D, textureID);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, res.x, res.y, 0, GL_RGBA, GL_FLOAT, nullptr);
    glBindTexture(GL_TEXTURE_2D, 0);
    glDisable(GL_TEXTURE_2D);
    return textureID;
}

GLuint SynthesisPipeline::_createGlResultBuffer(uint elements) {
    GLuint glBuffer;
    glGenBuffers(1, &glBuffer);
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, glBuffer);
    glBufferData(GL_PIXEL_UNPACK_BUFFER, elements * sizeof(glm::vec4), nullptr, GL_STREAM_DRAW);
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
    return glBuffer;
}