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

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

// CUDA Library Headers
#include <cuda_gl_interop.h>
#include <curand.h>

// CUDA utilities and system includes
#include <helper_cuda.h>
#include <rendercheck_gl.h>

// System includes
#include <iomanip>
#include <math.h>
#include <sstream>
#include <stdexcept>

// Includes
#include "rng.h"

// standard utility and system includes
#include <helper_timer.h>

// SDK information
static const char *printfFile = "randomFog.txt";

// RNG instance
RNG *g_pRng = NULL;

// CheckRender instance (for QA)
CheckRender *g_pCheckRender = NULL;

// Simple struct which contains the position and color of a vertex
struct SVertex
{
    GLfloat x, y, z;
    GLfloat r, g, b;
};

// Data for the vertices
SVertex *g_pVertices = NULL;
int      g_nVertices;          // Size of the vertex array
int      g_nVerticesPopulated; // Number currently populated

// Control the randomness
int nSkip1 = 0; // Number of samples to discard between x,y
int nSkip2 = 0; // Number of samples to discard between y,z
int nSkip3 = 0; // Number of samples to discard between z,x

// Control the display
enum Shape_t { Sphere, SphericalShell, Cube, Plane };
Shape_t g_currentShape = Sphere;
bool    g_bShowAxes    = true;
bool    g_bTenXZoom    = false;
bool    g_bAutoRotate  = true;
int     g_lastShapeX   = 1024;
int     g_lastShapeY   = 1024;
float   g_xRotated     = 0.0f;
float   g_yRotated     = 0.0f;

const float PI = 3.14159265359f;

void createCube(void)
{
    int startVertex = 0;

    for (int i = startVertex; i < g_nVerticesPopulated; i++) {
        g_pVertices[i].x = (g_pRng->getNextU01() - .5f) * 2;

        for (int j = 0; j < nSkip1; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].y = (g_pRng->getNextU01() - .5f) * 2;

        for (int j = 0; j < nSkip2; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].z = (g_pRng->getNextU01() - .5f) * 2;

        for (int j = 0; j < nSkip3; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].r = 1.0f;
        g_pVertices[i].g = 1.0f;
        g_pVertices[i].b = 1.0f;
    }
}

void createPlane(void)
{
    int startVertex = 0;

    for (int i = startVertex; i < g_nVerticesPopulated; i++) {
        g_pVertices[i].x = (g_pRng->getNextU01() - .5f) * 2;

        for (int j = 0; j < nSkip1; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].y = (g_pRng->getNextU01() - .5f) * 2;

        for (int j = 0; j < nSkip2; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].z = 0.0f;

        g_pVertices[i].r = 1.0f;
        g_pVertices[i].g = 1.0f;
        g_pVertices[i].b = 1.0f;
    }
}

void createSphere(void)
{
    int startVertex = 0;

    for (int i = startVertex; i < g_nVerticesPopulated; i++) {
        float r;
        float rho;
        float theta;

        if (g_currentShape == Sphere) {
            r = g_pRng->getNextU01();
            r = powf(r, 1.f / 3.f);

            for (int j = 0; j < nSkip3; j++) {
                g_pRng->getNextU01();
            }
        }
        else {
            r = 1.0f;
        }

        rho = g_pRng->getNextU01() * PI * 2.0f;

        for (int j = 0; j < nSkip1; j++) {
            g_pRng->getNextU01();
        }

        theta = (g_pRng->getNextU01() * 2.0f) - 1.0f;
        theta = asin(theta);

        for (int j = 0; j < nSkip2; j++) {
            g_pRng->getNextU01();
        }

        g_pVertices[i].x = r * fabs(cos(theta)) * cos(rho);
        g_pVertices[i].y = r * fabs(cos(theta)) * sin(rho);
        g_pVertices[i].z = r * sin(theta);

        g_pVertices[i].r = 1.0f;
        g_pVertices[i].g = 1.0f;
        g_pVertices[i].b = 1.0f;
    }
}

void createAxes(void)
{
    // z axis:
    g_pVertices[200000].x = 0.0f;
    g_pVertices[200000].y = 0.0f;
    g_pVertices[200000].z = -1.5f;
    g_pVertices[200001].x = 0.0f;
    g_pVertices[200001].y = 0.0f;
    g_pVertices[200001].z = 1.5f;
    g_pVertices[200000].r = 1.0f;
    g_pVertices[200000].g = 0.0f;
    g_pVertices[200000].b = 0.0f;
    g_pVertices[200001].r = 0.0f;
    g_pVertices[200001].g = 1.0f;
    g_pVertices[200001].b = 1.0f;
    // y axis:
    g_pVertices[200002].x = 0.0f;
    g_pVertices[200002].y = -1.5f;
    g_pVertices[200002].z = 0.0f;
    g_pVertices[200003].x = 0.0f;
    g_pVertices[200003].y = 1.5f;
    g_pVertices[200003].z = 0.0f;
    g_pVertices[200002].r = 0.0f;
    g_pVertices[200002].g = 1.0f;
    g_pVertices[200002].b = 0.0f;
    g_pVertices[200003].r = 1.0f;
    g_pVertices[200003].g = 0.0f;
    g_pVertices[200003].b = 1.0f;
    // x axis:
    g_pVertices[200004].x = -1.5f;
    g_pVertices[200004].y = 0.0f;
    g_pVertices[200004].z = 0.0f;
    g_pVertices[200005].x = 1.5f;
    g_pVertices[200005].y = 0.0f;
    g_pVertices[200005].z = 0.0f;
    g_pVertices[200004].r = 0.0f;
    g_pVertices[200004].g = 0.0f;
    g_pVertices[200004].b = 1.0f;
    g_pVertices[200005].r = 1.0f;
    g_pVertices[200005].g = 1.0f;
    g_pVertices[200005].b = 0.0f;
}

void drawPoints(void)
{
    if (g_bShowAxes) {
        glDrawArrays(GL_LINE_STRIP, 200000, 2);
        glDrawArrays(GL_LINE_STRIP, 200002, 2);
        glDrawArrays(GL_LINE_STRIP, 200004, 2);
    }

    glDrawArrays(GL_POINTS, 0, g_nVerticesPopulated);
}

void drawText(void)
{
    using std::string;
    using std::stringstream;

    glPushMatrix();
    glLoadIdentity();
    glRasterPos2f(-1.2f, 1.2f);

    string       infoString;
    stringstream ss;
    g_pRng->getInfoString(infoString);
    ss << " skip1=" << nSkip1;
    ss << " skip2=" << nSkip2;
    ss << " skip3=" << nSkip3;
    ss << " points=" << g_nVerticesPopulated;
    infoString.append(ss.str());

    for (unsigned int i = 0; i < infoString.size(); i++) {
        glutBitmapCharacter(GLUT_BITMAP_HELVETICA_12, infoString[i]);
    }

    glPopMatrix();
}

void reshape(int x, int y)
{
    float xScale;
    float yScale;

    g_lastShapeX = x;
    g_lastShapeY = y;

    // Check if shape is visible
    if (x == 0 || y == 0) {
        return;
    }

    // Set a new projection matrix
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();

    // Adjust fit
    if (y > x) {
        xScale = 1.0f;
        yScale = (float)y / x;
    }
    else {
        xScale = (float)x / y;
        yScale = 1.0f;
    }

    // Angle of view:40 degrees
    // Near clipping plane distance: 10.0 (default)
    // Far clipping plane distance: 10.0 (default)
    if (g_bTenXZoom) {
        glOrtho(-.15f * xScale, .15f * xScale, -.15f * yScale, .15f * yScale, -5.0f, 5.0f);
    }
    else {
        glOrtho(-1.5f * xScale, 1.5f * xScale, -1.5f * yScale, 1.5f * yScale, -10.0f, 10.0f);
    }

    // Use the whole window for rendering
    glViewport(0, 0, x, y);
    glMatrixMode(GL_MODELVIEW);
}

void display(void)
{
    glClear(GL_COLOR_BUFFER_BIT);
    glLoadIdentity();
    glTranslatef(0.0f, 0.0f, -4.0f);
    glRotatef(g_yRotated, 0.0f, 1.0f, 0.0f);
    glRotatef(g_xRotated, 1.0f, 0.0f, 0.0f);
    drawPoints();
    drawText();
    glFlush();
    glutSwapBuffers();
}

void idle(void)
{
    if (g_bAutoRotate) {
        g_yRotated += 0.1f;

        if (g_yRotated >= 360.0f) {
            g_yRotated -= 360.0f;
        }

        g_xRotated += 0.05f;

        if (g_xRotated >= 360.0f) {
            g_xRotated -= 360.0f;
        }

        display();
    }
}

void reCreate(void)
{
    switch (g_currentShape) {
    case Sphere:
    case SphericalShell:
        createSphere();
        break;

    case Cube:
        createCube();
        break;

    default:
        createPlane();
    }

    display();
}

void cleanup(int code)
{
    if (g_pRng) {
        delete g_pRng;
        g_pRng = NULL;
    }

    if (g_pVertices) {
        delete[] g_pVertices;
        g_pVertices = NULL;
    }

    if (g_pCheckRender) {
        delete g_pCheckRender;
        g_pCheckRender = NULL;
    }

    exit(code);
}

void glutClose() { cleanup(EXIT_SUCCESS); }

void keyboard(unsigned char key, int x, int y)
{
    switch (key) {
    // Select shape
    case 's':
    case 'S':
        g_currentShape = Sphere;
        createSphere();
        display();
        break;

    case 'e':
    case 'E':
        g_currentShape = SphericalShell;
        createSphere();
        display();
        break;

    case 'b':
    case 'B':
        g_currentShape = Cube;
        createCube();
        display();
        break;

    case 'p':
    case 'P':
        g_currentShape = Plane;
        createPlane();
        display();
        break;

    // Rotation
    case 'a':
    case 'A':
        g_bAutoRotate = !g_bAutoRotate;
        break;

    case 'i':
    case 'I':
        g_xRotated -= 1.0f;

        if (g_xRotated <= 0.0f) {
            g_xRotated += 360.0f;
        }

        display();
        break;

    case ',':
        g_xRotated += 1.0f;

        if (g_xRotated >= 360.0f) {
            g_xRotated -= 360.0f;
        }

        display();
        break;

    case 'j':
    case 'J':
        g_yRotated -= 1.0f;

        if (g_yRotated <= 0.0f) {
            g_yRotated += 360.0f;
        }

        display();
        break;

    case 'l':
    case 'L':
        g_yRotated += 1.0f;

        if (g_yRotated >= 360.0f) {
            g_yRotated -= 360.0f;
        }

        display();
        break;

    // Zoom
    case 't':
    case 'T':
        g_bTenXZoom = !g_bTenXZoom;
        reshape(g_lastShapeX, g_lastShapeY);
        reCreate();
        break;

    // Axes
    case 'z':
    case 'Z':
        g_bShowAxes = !g_bShowAxes;
        reCreate();
        break;

    // RNG
    case 'x':
    case 'X':
        g_pRng->selectRng(RNG::Pseudo);
        reCreate();
        break;

    case 'c':
    case 'C':
        g_pRng->selectRng(RNG::Quasi);
        reCreate();
        break;

    case 'v':
    case 'V':
        g_pRng->selectRng(RNG::ScrambledQuasi);
        reCreate();
        break;

    case 'r':
    case 'R':
        g_pRng->resetSeed();
        reCreate();
        break;

    case ']':
        g_pRng->incrementDimensions();
        reCreate();
        break;

    case '[':
        g_pRng->resetDimensions();
        reCreate();
        break;

    case '1':
        nSkip1++;
        reCreate();
        break;

    case '2':
        nSkip2++;
        reCreate();
        break;

    case '3':
        nSkip3++;
        reCreate();
        break;

    case '!':
        nSkip1 = 0;
        nSkip2 = 0;
        nSkip3 = 0;
        reCreate();
        break;

    // Number of vertices
    case '+':
        g_nVerticesPopulated += 8000;

        if (g_nVerticesPopulated > g_nVertices) {
            g_nVerticesPopulated = g_nVertices;
        }

        reCreate();
        break;

    case '-':
        g_nVerticesPopulated -= 8000;

        if (g_nVerticesPopulated < 8000) {
            g_nVerticesPopulated = 8000;
        }

        reCreate();
        break;

    // Quit
    case 27:
    case 'q':
    case 'Q':
#if defined(__APPLE__) || defined(MACOSX)
        exit(EXIT_SUCCESS);
#else
        glutDestroyWindow(glutGetWindow());
        return;
#endif
    }
}

void showHelp(void)
{
    using std::left;
    using std::setw;
    using std::stringstream;

    stringstream ss;

    ss << "\nRandom number visualization\n\n";
    ss << "On creation, randomFog generates 200,000 random coordinates in "
          "spherical coordinate space (radius, angle rho, angle theta) with "
          "curand's XORWOW algorithm. The coordinates are normalized for a "
          "uniform distribution through the sphere.\n\n";
    ss << "The X axis is drawn with blue in the negative direction and yellow "
          "positive.\n"
       << "The Y axis is drawn with green in the negative direction and magenta "
          "positive.\n"
       << "The Z axis is drawn with red in the negative direction and cyan "
          "positive.\n\n";
    ss << "The following keys can be used to control the output:\n\n";
    ss << left;
    ss << "\t" << setw(10) << "s"
       << "Generate a new set of random numbers and display as spherical "
          "coordinates (Sphere)\n";
    ss << "\t" << setw(10) << "e"
       << "Generate a new set of random numbers and display on a spherical "
          "surface (shEll)\n";
    ss << "\t" << setw(10) << "b"
       << "Generate a new set of random numbers and display as cartesian "
          "coordinates (cuBe/Box)\n";
    ss << "\t" << setw(10) << "p"
       << "Generate a new set of random numbers and display on a cartesian plane "
          "(Plane)\n\n";
    ss << "\t" << setw(10) << "i,l,j"
       << "Rotate the negative Z-axis up, right, down and left respectively\n";
    ss << "\t" << setw(10) << "a"
       << "Toggle auto-rotation\n";
    ss << "\t" << setw(10) << "t"
       << "Toggle 10x zoom\n";
    ss << "\t" << setw(10) << "z"
       << "Toggle axes display\n\n";
    ss << "\t" << setw(10) << "x"
       << "Select XORWOW generator (default)\n";
    ss << "\t" << setw(10) << "c"
       << "Select Sobol' generator\n";
    ss << "\t" << setw(10) << "v"
       << "Select scrambled Sobol' generator\n";
    ss << "\t" << setw(10) << "r"
       << "Reset XORWOW (i.e. reset to initial seed) and regenerate\n";
    ss << "\t" << setw(10) << "]"
       << "Increment the number of Sobol' dimensions and regenerate\n";
    ss << "\t" << setw(10) << "["
       << "Reset the number of Sobol' dimensions to 1 and regenerate\n\n";
    ss << "\t" << setw(10) << "+"
       << "Increment the number of displayed points by 8,000 (up to maximum "
          "200,000)\n";
    ss << "\t" << setw(10) << "-"
       << "Decrement the number of displayed points by 8,000 (down to minimum "
          "8,000)\n\n";
    ss << "\t" << setw(10) << "q/[ESC]"
       << "Quit the application.\n\n";
    puts(ss.str().c_str());
}

int main(int argc, char **argv)
{
    using std::runtime_error;

    try {
        bool bQA = false;

        // Open the log file
        printf("Random Fog\n");
        printf("==========\n\n");

        // Check QA mode
        if (checkCmdLineFlag(argc, (const char **)argv, "qatest")) {
            bQA = true;

            findCudaDevice(argc, (const char **)argv);

            g_pCheckRender = new CheckBackBuffer(g_lastShapeX, g_lastShapeY, 4, false);
        }
        else {
#if defined(__linux__)
            setenv("DISPLAY", ":0", 0);
#endif
            // Initialize GL
            glutInit(&argc, argv);
            glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB);
            // TODO use width/height?
            glutInitWindowSize(1000, 1000);
            // Create a window with rendering context and everything else we need
            glutCreateWindow("Random Fog");

            if (!isGLVersionSupported(2, 0)) {
                fprintf(stderr, "This sample requires at least OpenGL 2.0\n");
                exit(EXIT_WAIVED);
            }

            // Select CUDA device with OpenGL interoperability
            findCudaDevice(argc, (const char **)argv);
        }

        // Create vertices
        g_nVertices          = 200000;
        g_nVerticesPopulated = 200000;
        g_pVertices          = new SVertex[g_nVertices + 6];

        // Setup the random number generators
        g_pRng = new RNG(12345, 1, 100000);
        printf("CURAND initialized\n");

        // Compute the initial vertices and indices, starting in spherical mode
        createSphere();
        createAxes();

        showHelp();

        if (bQA) {
            g_pCheckRender->setExecPath(argv[0]);
            g_pCheckRender->dumpBin(g_pVertices, g_nVerticesPopulated * sizeof(SVertex), "randomFog.bin");

            if (g_pCheckRender->compareBin2BinFloat("randomFog.bin",
                                                    "ref_randomFog.bin",
                                                    g_nVerticesPopulated * sizeof(SVertex) / sizeof(float),
                                                    0.25f,
                                                    0.35f)) {
                cleanup(EXIT_SUCCESS);
            }
            else {
                cleanup(EXIT_FAILURE);
            }
        }
        else {
            // As we do not yet use a depth buffer, we cannot fill our sphere
            glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);
            // Enable the vertex array functionality:
            glEnableClientState(GL_VERTEX_ARRAY);
            // Enable the color array functionality (so we can specify a color for
            // each vertex)
            glEnableClientState(GL_COLOR_ARRAY);
            // Pass the vertex pointer:
            glVertexPointer(3, // 3 components per vertex (x,y,z)
                            GL_FLOAT,
                            sizeof(SVertex),
                            g_pVertices);
            // Pass the color pointer
            glColorPointer(3, // 3 components per vertex (r,g,b)
                           GL_FLOAT,
                           sizeof(SVertex),
                           &g_pVertices[0].r); // Pointer to the first color
            // Point size for point mode
            glPointSize(1.0f);
            glLineWidth(2.0f);
            glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
            // Notify glut which messages we require:
            glutDisplayFunc(display);
            glutReshapeFunc(reshape);
            glutKeyboardFunc(keyboard);
            glutIdleFunc(idle);

#if defined(__APPLE__) || defined(MACOSX)
            atexit(glutClose);
#else
            glutCloseFunc(glutClose);
#endif

            // Let's get started!
            glutMainLoop();
        }
    }
    catch (runtime_error &e) {
        printf("runtime error (%s)\n", e.what());
    }

    exit(EXIT_SUCCESS);
}
