/* 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.
 */

/*
    This class renders particles using OpenGL and GLSL shaders
*/

#include <math.h>
#include <stdlib.h>
#define HELPERGL_EXTERN_GL_FUNC_IMPLEMENTATION
#include <helper_gl.h>

#include "SmokeRenderer.h"
#include "SmokeShaders.h"

#if defined(__APPLE__) || defined(MACOSX)
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
#include <GLUT/glut.h>
#else
#include <GL/freeglut.h>
#endif

#define USE_MBLUR         1
#define COLOR_ATTENUATION 1

SmokeRenderer::SmokeRenderer(int maxParticles)
    : mMaxParticles(maxParticles)
    , mNumParticles(0)
    , mPosVbo(0)
    , mVelVbo(0)
    , mColorVbo(0)
    , mIndexBuffer(0)
    , mParticleRadius(0.005f)
    , mDisplayMode(VOLUMETRIC)
    , mWindowW(800)
    , mWindowH(600)
    , mFov(40.0f)
    , m_downSample(2)
    , m_numSlices(32)
    , m_numDisplayedSlices(32)
    , m_sliceNo(0)
    , m_shadowAlpha(0.005f)
    , m_spriteAlpha(0.1f)
    , m_doBlur(false)
    , m_blurRadius(1.0f)
    , m_displayLightBuffer(false)
    , m_lightPos(5.0f, 5.0f, -5.0f)
    , m_lightTarget(0.0f, 0.0f, 0.0f)
    , m_lightColor(1.0f, 1.0f, 0.5f)
    , m_colorAttenuation(0.1f, 0.2f, 0.3f)
    , m_lightBufferSize(256)
    , m_srcLightTexture(0)
    , m_lightDepthTexture(0)
    , m_lightFbo(0)
    , m_imageTex(0)
    , m_depthTex(0)
    , m_imageFbo(0)
{
    // load shader programs
    m_simpleProg         = new GLSLProgram(particleVS, simplePS);
    m_particleProg       = new GLSLProgram(mblurVS, mblurGS, particlePS);
    m_particleShadowProg = new GLSLProgram(mblurVS, mblurGS, particleShadowPS);

    m_blurProg       = new GLSLProgram(passThruVS, blurPS);
    m_displayTexProg = new GLSLProgram(passThruVS, texture2DPS);

    // create buffer for light shadows
    createLightBuffer();

    glutReportErrors();
}

SmokeRenderer::~SmokeRenderer()
{
    delete m_particleProg;
    delete m_particleShadowProg;
    delete m_blurProg;
    delete m_displayTexProg;
    delete m_simpleProg;

    delete m_lightFbo;
    glDeleteTextures(2, m_lightTexture);
    glDeleteTextures(1, &m_lightDepthTexture);

    delete m_imageFbo;
    glDeleteTextures(1, &m_imageTex);
    glDeleteTextures(1, &m_depthTex);
}

// draw points from vertex buffer objects
void SmokeRenderer::drawPoints(int start, int count, bool sort)
{
    glBindBuffer(GL_ARRAY_BUFFER, mPosVbo);
    glVertexPointer(4, GL_FLOAT, 0, 0);
    glEnableClientState(GL_VERTEX_ARRAY);

    if (mColorVbo) {
        glBindBuffer(GL_ARRAY_BUFFER, mColorVbo);
        glColorPointer(4, GL_FLOAT, 0, 0);
        glEnableClientState(GL_COLOR_ARRAY);
    }

    if (mVelVbo) {
        glBindBuffer(GL_ARRAY_BUFFER, mVelVbo);
        glClientActiveTexture(GL_TEXTURE0);
        glTexCoordPointer(4, GL_FLOAT, 0, 0);
        glEnableClientState(GL_TEXTURE_COORD_ARRAY);
    }

    if (sort) {
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER_ARB, mIndexBuffer);
        glDrawElements(GL_POINTS, count, GL_UNSIGNED_INT, (void *)(start * sizeof(unsigned int)));
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER_ARB, 0);
    }
    else {
        glDrawArrays(GL_POINTS, start, count);
    }

    glDisableClientState(GL_VERTEX_ARRAY);
    glDisableClientState(GL_COLOR_ARRAY);

    glClientActiveTexture(GL_TEXTURE0);
    glDisableClientState(GL_TEXTURE_COORD_ARRAY);
}

// draw points using given shader program
void SmokeRenderer::drawPointSprites(GLSLProgram *prog, int start, int count, bool shadowed)
{
    glEnable(GL_DEPTH_TEST);
    glDepthMask(GL_FALSE); // don't write depth
    glEnable(GL_BLEND);

    prog->enable();
    prog->setUniform1f("pointRadius", mParticleRadius);

    if (shadowed) {
        prog->bindTexture("shadowTex", m_lightTexture[m_srcLightTexture], GL_TEXTURE_2D, 0);
    }

    // draw points
    drawPoints(start, count, true);

    prog->disable();

    glDepthMask(GL_TRUE);
    glDisable(GL_BLEND);
}

// calculate vectors for half-angle slice rendering
void SmokeRenderer::calcVectors()
{
    // get model view matrix
    glGetFloatv(GL_MODELVIEW_MATRIX, (float *)m_modelView.get_value());

    // calculate eye space light vector
    m_lightVector = normalize(m_lightPos);
    m_lightPosEye = m_modelView * vec4f(m_lightPos, 1.0);

    // calculate half-angle vector between view and light
    m_viewVector = -vec3f(m_modelView.get_row(2));

    if (dot(m_viewVector, m_lightVector) > 0) {
        m_halfVector   = normalize(m_viewVector + m_lightVector);
        m_invertedView = false;
    }
    else {
        m_halfVector   = normalize(-m_viewVector + m_lightVector);
        m_invertedView = true;
    }

    // calculate light view matrix
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    glLoadIdentity();
    gluLookAt(m_lightPos[0],
              m_lightPos[1],
              m_lightPos[2],
              m_lightTarget[0],
              m_lightTarget[1],
              m_lightTarget[2],
              0.0,
              1.0,
              0.0);

    // calculate light projection matrix
    glMatrixMode(GL_PROJECTION);
    glPushMatrix();
    glLoadIdentity();
    gluPerspective(45.0, 1.0, 1.0, 200.0);

    glGetFloatv(GL_MODELVIEW_MATRIX, (float *)m_lightView.get_value());
    glGetFloatv(GL_PROJECTION_MATRIX, (float *)m_lightProj.get_value());

    glMatrixMode(GL_PROJECTION);
    glPopMatrix();
    glMatrixMode(GL_MODELVIEW);
    glPopMatrix();

    // construct shadow matrix
    matrix4f scale;
    scale.set_scale(vec3f(0.5, 0.5, 0.5));
    matrix4f translate;
    translate.set_translate(vec3f(0.5, 0.5, 0.5));

    m_shadowMatrix = translate * scale * m_lightProj * m_lightView * inverse(m_modelView);

    // calc object space eye position
    m_eyePos = inverse(m_modelView) * vec4f(0.0, 0.0, 0.0, 1.0);

    // calc half vector in eye space
    m_halfVectorEye = m_modelView * vec4f(m_halfVector, 0.0);
}

// draw slice of particles from camera view
void SmokeRenderer::drawSlice(int i)
{
    m_imageFbo->Bind();
    glViewport(0, 0, m_imageW, m_imageH);

    glColor4f(1.0, 1.0, 1.0, m_spriteAlpha);

    if (m_invertedView) {
        // front-to-back
        glBlendFunc(GL_ONE_MINUS_DST_ALPHA, GL_ONE);
    }
    else {
        // back-to-front
        glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_ALPHA);
    }

    drawPointSprites(m_particleShadowProg, i * m_batchSize, m_batchSize, true);

    m_imageFbo->Disable();
}

// draw slice of particles from light's point of view
void SmokeRenderer::drawSliceLightView(int i)
{
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    glLoadMatrixf((GLfloat *)m_lightView.get_value());

    glMatrixMode(GL_PROJECTION);
    glPushMatrix();
    glLoadMatrixf((GLfloat *)m_lightProj.get_value());

    m_lightFbo->Bind();
    glViewport(0, 0, m_lightBufferSize, m_lightBufferSize);

    glColor4f(m_colorAttenuation[0] * m_shadowAlpha,
              m_colorAttenuation[1] * m_shadowAlpha,
              m_colorAttenuation[2] * m_shadowAlpha,
              1.0);
    glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_COLOR);

    drawPointSprites(m_particleProg, i * m_batchSize, m_batchSize, false);

    m_lightFbo->Disable();

    glMatrixMode(GL_PROJECTION);
    glPopMatrix();
    glMatrixMode(GL_MODELVIEW);
    glPopMatrix();
}

// draw particles as slices with shadowing
void SmokeRenderer::drawSlices()
{
    m_batchSize = mNumParticles / m_numSlices;

    // clear light buffer
    m_srcLightTexture = 0;
    m_lightFbo->Bind();
    m_lightFbo->AttachTexture(GL_TEXTURE_2D, m_lightTexture[m_srcLightTexture], GL_COLOR_ATTACHMENT0_EXT);
    glClearColor(1.0f - m_lightColor[0], 1.0f - m_lightColor[1], 1.0f - m_lightColor[2], 0.0f);
    glClear(GL_COLOR_BUFFER_BIT);
    m_lightFbo->Disable();

    // clear volume image
    m_imageFbo->Bind();
    glClearColor(0.0, 0.0, 0.0, 0.0);
    glClear(GL_COLOR_BUFFER_BIT);
    m_imageFbo->Disable();

    glActiveTexture(GL_TEXTURE0);
    glMatrixMode(GL_TEXTURE);
    glLoadMatrixf((GLfloat *)m_shadowMatrix.get_value());

    // render slices
    if (m_numDisplayedSlices > m_numSlices)
        m_numDisplayedSlices = m_numSlices;

    for (int i = 0; i < m_numDisplayedSlices; i++) {
        // draw slice from camera view, sampling light buffer
        drawSlice(i);
        // draw slice from light view to light buffer, accumulating shadows
        drawSliceLightView(i);

        if (m_doBlur) {
            blurLightBuffer();
        }
    }

    glActiveTexture(GL_TEXTURE0);
    glMatrixMode(GL_TEXTURE);
    glLoadIdentity();
}

// blur light buffer to simulate scattering effects
void SmokeRenderer::blurLightBuffer()
{
    m_lightFbo->Bind();
    m_lightFbo->AttachTexture(GL_TEXTURE_2D, m_lightTexture[1 - m_srcLightTexture], GL_COLOR_ATTACHMENT0_EXT);
    glViewport(0, 0, m_lightBufferSize, m_lightBufferSize);

    m_blurProg->enable();
    m_blurProg->bindTexture("tex", m_lightTexture[m_srcLightTexture], GL_TEXTURE_2D, 0);
    m_blurProg->setUniform2f("texelSize", 1.0f / (float)m_lightBufferSize, 1.0f / (float)m_lightBufferSize);
    m_blurProg->setUniform1f("blurRadius", m_blurRadius);
    glDisable(GL_DEPTH_TEST);
    drawQuad();
    m_blurProg->disable();

    m_srcLightTexture = 1 - m_srcLightTexture;

    m_lightFbo->Disable();
}

// display texture to screen
void SmokeRenderer::displayTexture(GLuint tex)
{
    m_displayTexProg->enable();
    m_displayTexProg->bindTexture("tex", tex, GL_TEXTURE_2D, 0);
    drawQuad();
    m_displayTexProg->disable();
}

// composite final volume image on top of scene
void SmokeRenderer::compositeResult()
{
    glViewport(0, 0, mWindowW, mWindowH);
    glDisable(GL_DEPTH_TEST);
    glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_ALPHA);
    glEnable(GL_BLEND);
    displayTexture(m_imageTex);
    glDisable(GL_BLEND);
}

void SmokeRenderer::render()
{
    switch (mDisplayMode) {
    case POINTS:
        glColor3f(1.0, 1.0, 1.0);
        m_simpleProg->enable();
        drawPoints(0, mNumParticles, false);
        m_simpleProg->disable();
        break;

    case SPRITES:
        glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_COLOR);
        glColor4f(1.0, 1.0, 1.0, m_spriteAlpha);
        drawPointSprites(m_particleProg, 0, mNumParticles, false);
        break;

    case VOLUMETRIC:
        drawSlices();
        compositeResult();
        break;

    case NUM_MODES:
        break;
    }

    if (m_displayLightBuffer) {
        // display light buffer to screen
        glViewport(0, 0, m_lightBufferSize, m_lightBufferSize);
        glDisable(GL_DEPTH_TEST);
        displayTexture(m_lightTexture[m_srcLightTexture]);
        glViewport(0, 0, mWindowW, mWindowH);
    }

    glutReportErrors();
}

// render scene depth to texture
// (this is to ensure that particle are correctly occluded in the low-resolution
// render buffer)
void SmokeRenderer::beginSceneRender(Target target)
{
    if (target == LIGHT_BUFFER) {
        m_lightFbo->Bind();
        glViewport(0, 0, m_lightBufferSize, m_lightBufferSize);

        glMatrixMode(GL_MODELVIEW);
        glPushMatrix();
        glLoadMatrixf((GLfloat *)m_lightView.get_value());

        glMatrixMode(GL_PROJECTION);
        glPushMatrix();
        glLoadMatrixf((GLfloat *)m_lightProj.get_value());
    }
    else {
        m_imageFbo->Bind();
        glViewport(0, 0, m_imageW, m_imageH);
    }

    glColorMask(GL_FALSE, GL_FALSE, GL_FALSE, GL_FALSE);
    glDepthMask(GL_TRUE);
    glClear(GL_DEPTH_BUFFER_BIT);
}

void SmokeRenderer::endSceneRender(Target target)
{
    if (target == LIGHT_BUFFER) {
        m_lightFbo->Disable();
        glMatrixMode(GL_PROJECTION);
        glPopMatrix();
        glMatrixMode(GL_MODELVIEW);
        glPopMatrix();
    }
    else {
        m_imageFbo->Disable();
    }

    glViewport(0, 0, mWindowW, mWindowH);
    glColorMask(GL_TRUE, GL_TRUE, GL_TRUE, GL_TRUE);
}

// create an OpenGL texture
GLuint SmokeRenderer::createTexture(GLenum target, int w, int h, GLint internalformat, GLenum format)
{
    GLuint texid;
    glGenTextures(1, &texid);
    glBindTexture(target, texid);

    glTexParameteri(target, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(target, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(target, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(target, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    glTexImage2D(target, 0, internalformat, w, h, 0, format, GL_FLOAT, 0);
    return texid;
}

// create buffers for off-screen rendering
void SmokeRenderer::createBuffers(int w, int h)
{
    if (m_imageFbo) {
        glDeleteTextures(1, &m_imageTex);
        glDeleteTextures(1, &m_depthTex);
        delete m_imageFbo;
    }

    mWindowW = w;
    mWindowH = h;

    m_imageW = w / m_downSample;
    m_imageH = h / m_downSample;

    // create fbo for image buffer
    GLint format = GL_RGBA16F_ARB;
    // GLint format = GL_LUMINANCE16F_ARB;
    // GLint format = GL_RGBA8;
    m_imageTex = createTexture(GL_TEXTURE_2D, m_imageW, m_imageH, format, GL_RGBA);
    m_depthTex = createTexture(GL_TEXTURE_2D, m_imageW, m_imageH, GL_DEPTH_COMPONENT24_ARB, GL_DEPTH_COMPONENT);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);

    m_imageFbo = new FramebufferObject();
    m_imageFbo->AttachTexture(GL_TEXTURE_2D, m_imageTex, GL_COLOR_ATTACHMENT0_EXT);
    m_imageFbo->AttachTexture(GL_TEXTURE_2D, m_depthTex, GL_DEPTH_ATTACHMENT_EXT);
    m_imageFbo->IsValid();
}

void SmokeRenderer::setLightColor(vec3f c)
{
    m_lightColor = c;

    // set light texture border color
    GLfloat borderColor[4] = {1.0f - m_lightColor[0], 1.0f - m_lightColor[1], 1.0f - m_lightColor[2], 0.0f};

    glBindTexture(GL_TEXTURE_2D, m_lightTexture[0]);
    glTexParameterfv(GL_TEXTURE_2D, GL_TEXTURE_BORDER_COLOR, borderColor);

    glBindTexture(GL_TEXTURE_2D, m_lightTexture[1]);
    glTexParameterfv(GL_TEXTURE_2D, GL_TEXTURE_BORDER_COLOR, borderColor);

    glBindTexture(GL_TEXTURE_2D, 0);
}

// create FBOs for light buffer
void SmokeRenderer::createLightBuffer()
{
    GLint format = GL_RGBA16F_ARB;
    // GLint format = GL_RGBA8;
    // GLint format = GL_LUMINANCE16F_ARB;

    m_lightTexture[0] = createTexture(GL_TEXTURE_2D, m_lightBufferSize, m_lightBufferSize, format, GL_RGBA);
    // make shadows clamp to light color at edges
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);

    m_lightTexture[1] = createTexture(GL_TEXTURE_2D, m_lightBufferSize, m_lightBufferSize, format, GL_RGBA);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);

    m_lightDepthTexture = createTexture(
        GL_TEXTURE_2D, m_lightBufferSize, m_lightBufferSize, GL_DEPTH_COMPONENT24_ARB, GL_DEPTH_COMPONENT);

    m_lightFbo = new FramebufferObject();
    m_lightFbo->AttachTexture(GL_TEXTURE_2D, m_lightTexture[m_srcLightTexture], GL_COLOR_ATTACHMENT0_EXT);
    m_lightFbo->AttachTexture(GL_TEXTURE_2D, m_lightDepthTexture, GL_DEPTH_ATTACHMENT_EXT);
    m_lightFbo->IsValid();
}

void SmokeRenderer::setWindowSize(int w, int h)
{
    mAspect      = (float)mWindowW / (float)mWindowH;
    mInvFocalLen = tan(mFov * 0.5f * NV_PI / 180.0f);

    createBuffers(w, h);
}

void SmokeRenderer::drawQuad()
{
    glBegin(GL_QUADS);
    glTexCoord2f(0.0f, 0.0f);
    glVertex2f(-1.0f, -1.0f);
    glTexCoord2f(1.0f, 0.0f);
    glVertex2f(1.0f, -1.0f);
    glTexCoord2f(1.0f, 1.0f);
    glVertex2f(1.0f, 1.0f);
    glTexCoord2f(0.0f, 1.0f);
    glVertex2f(-1.0f, 1.0f);
    glEnd();
}

void SmokeRenderer::drawVector(vec3f v)
{
    glBegin(GL_LINES);
    glVertex3f(0.0f, 0.0f, 0.0f);
    glVertex3fv((float *)&v[0]);
    glEnd();
}

// render vectors to screen for debugging
void SmokeRenderer::debugVectors()
{
    glColor3f(1.0f, 1.0f, 0.0f);
    drawVector(m_lightVector);

    glColor3f(0.0f, 1.0f, 0.0f);
    drawVector(m_viewVector);

    glColor3f(0.0f, 0.0f, 1.0f);
    drawVector(-m_viewVector);

    glColor3f(1.0f, 0.0f, 0.0f);
    drawVector(m_halfVector);
}