/*
 * Copyright 1993-2015 NVIDIA Corporation.  All rights reserved.
 *
 * Please refer to the NVIDIA end user license agreement (EULA) associated
 * with this source code for terms and conditions that govern your use of
 * this software. Any use, reproduction, disclosure, or distribution of
 * this software and related documentation outside the terms of the EULA
 * is strictly prohibited.
 *
 */

/*
    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);
}
