// A simple fluid solver.. perhaps extended with some other stuff (hopefully).
// Martin Michelsen
// Caltech CS 101c, spring 2010

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include <GL/glew.h>

#if defined (__APPLE__) || defined(MACOSX)
#include <GLUT/glut.h>
#else
#include <GL/glut.h>
#endif

#include <cuda_gl_interop.h>

#include <lab8_uninitkernel.cu>

// first, some error handling functions
inline void _safecall(cudaError err, const char *file, const unsigned int line)
{
    if (err != cudaSuccess)
    {
        fprintf(stderr, "at %s/%u: CUDA runtime error: %d/%s\n",
                file, line, err, cudaGetErrorString(err));
        exit(-1);
    }
}
inline void _checkmsg(const char *msg, const char *file, const unsigned int line)
{
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
    {
        fprintf(stderr, "at %s/%u: %s: %d/%s\n",
                file, line, msg, err, cudaGetErrorString(err));
        exit(-1);
    }
}
#define safecall(a) _safecall(a, __FILE__, __LINE__)
#define checkmsg(a) _checkmsg(a, __FILE__, __LINE__)

#define iDivUp(a, b) ((a % b != 0) ? (a / b + 1) : (a / b))

// simulation setup parameters
// these are default parameters. they can be changed (TODO) via the command line.
long fieldDimension = 256; // dimension of simulation grid (it's a cube)
long windowWidth = 512; // dimensions of viewing window
long windowHeight = 512;

// view settings (from lab 6)
float3 viewRotation = make_float3(0.5, 0.5, 0.0);
float3 viewTranslation = make_float3(0.0, 0.0, -8.0f);
float3 lightPos = make_float3(2.0f, 2.0f, 2.0f);
int renderFlags = 7;
float iso = 0.5f;

cudaArray *dPressure;

// OpenGL pixel buffer
GLuint pbo = 0;

////////////////////////////////////////////////////////////////////////////////
// update & display functions

// display results using OpenGL (called by GLUT)
void display()
{
    // use OpenGL to build view matrix
    GLfloat modelView[16];
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    glLoadIdentity();
    glRotatef(-viewRotation.x, 1.0f, 0.0f, 0.0f);
    glRotatef(-viewRotation.y, 0.0f, 1.0f, 0.0f);
    glTranslatef(-viewTranslation.x, -viewTranslation.y, -viewTranslation.z);
    glGetFloatv(GL_MODELVIEW_MATRIX, modelView);
    glPopMatrix();

    // transpose matrix to conform with OpenGL's notation
    GLfloat invViewMatrix[12];
    invViewMatrix[0] = modelView[0];
    invViewMatrix[1] = modelView[4];
    invViewMatrix[2] = modelView[8];
    invViewMatrix[3] = modelView[12];
    invViewMatrix[4] = modelView[1];
    invViewMatrix[5] = modelView[5];
    invViewMatrix[6] = modelView[9];
    invViewMatrix[7] = modelView[13];
    invViewMatrix[8] = modelView[2];
    invViewMatrix[9] = modelView[6];
    invViewMatrix[10] = modelView[10];
    invViewMatrix[11] = modelView[14];

    // set necessary constants in hardware
    safecall(cudaMemcpyToSymbol(dInvViewMatrix, invViewMatrix, sizeof(float4) * 3));

    // map PBO to get CUDA device pointer
    unsigned long* renderOutput;
    safecall(cudaGLMapBufferObject((void**)&renderOutput, pbo));

    // call CUDA kernel, writing results to PBO
    dim3 blockSize(32, 8, 1);
    dim3 gridSize = dim3(iDivUp(windowWidth, blockSize.x), iDivUp(windowHeight, blockSize.y), 1);
    RenderPressureIsosurfaceK<<< gridSize, blockSize >>>(renderOutput, windowWidth, windowHeight);
    checkmsg("RenderPressureIsosurfaceK failed");
    safecall(cudaGLUnmapBufferObject(pbo));

    // display results
    glClear(GL_COLOR_BUFFER_BIT);

    // draw image from PBO
    glDisable(GL_DEPTH_TEST);
    glRasterPos2i(0, 0);
    glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
    glDrawPixels(windowWidth, windowHeight, GL_RGBA, GL_UNSIGNED_BYTE, 0);
    glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, 0);

    glutSwapBuffers();
    glutReportErrors();

    glutPostRedisplay();
}

////////////////////////////////////////////////////////////////////////////////
// initialization functions

// "initialize" our data structures
void initCuda()
{
    // create 3D cudaArrays for the velocity and pressure textures
    cudaExtent volumeSize = make_cudaExtent(fieldDimension, fieldDimension, fieldDimension);
    cudaChannelFormatDesc channelDescPressure = cudaCreateChannelDesc<float>();
    safecall(cudaMalloc3DArray(&dPressure, &channelDescPressure, volumeSize));

    // set texture parameters
    pressureTex.normalized = true;
    pressureTex.filterMode = cudaFilterModeLinear;
    pressureTex.addressMode[0] = cudaAddressModeClamp;
    pressureTex.addressMode[1] = cudaAddressModeClamp;
    pressureTex.addressMode[2] = cudaAddressModeClamp;

    // bind arrays to 3D textures
    safecall(cudaBindTextureToArray(pressureTex, dPressure, channelDescPressure));
}

// deallocate our memory
void cleanup()
{
    safecall(cudaFreeArray(dPressure));
    safecall(cudaGLUnregisterBufferObject(pbo));    
    glDeleteBuffersARB(1, &pbo);
}

// initialize some parameters we use for rendering
void initGLParameters()
{
    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glOrtho(0.0, 1.0, 0.0, 1.0, 0.0, 1.0);
}

// initialize (or re-initialize) the pixel buffer
void initPixelBuffer()
{
    if (pbo) {
        // delete old buffer
        safecall(cudaGLUnregisterBufferObject(pbo));
        glDeleteBuffersARB(1, &pbo);
    }

    // create pixel buffer object for display
    glGenBuffersARB(1, &pbo);
    glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
    glBufferDataARB(GL_PIXEL_UNPACK_BUFFER_ARB, windowWidth * windowHeight * sizeof(GLubyte) * 4, 0, GL_STREAM_DRAW_ARB);
    glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, 0);
    safecall(cudaGLRegisterBufferObject(pbo));
}

// set constants on device
void setConsts()
{
    safecall(cudaMemcpyToSymbol(dLightPos, &lightPos, sizeof(float3)));
    safecall(cudaMemcpyToSymbol(dRenderFlags, &renderFlags, sizeof(int)));
    safecall(cudaMemcpyToSymbol(dFieldDimension, &fieldDimension, sizeof(int)));
    safecall(cudaMemcpyToSymbol(dIso, &iso, sizeof(float)));
}

////////////////////////////////////////////////////////////////////////////////
// user interaction functions

void idle()
{
    // do nothing
}

void keyboard(unsigned char key, int x, int y)
{
    //float ds = 0.005f;
    float dl = 0.1f;

    switch (key)
    {
        // esc = exit
      case 27:
        exit(0);
        break;

        // [] = move light along z
      case ']':
        lightPos.z -= dl;
        break;
      case '[':
        lightPos.z += dl;
        break;

        // r, t, y = enable/disable various types of lights
      case 'r':
        if (renderFlags & RENDER_AMBIENT) renderFlags &= ~RENDER_AMBIENT;
        else renderFlags |= RENDER_AMBIENT;
        printf("ambient lighting: %s\n", (renderFlags & RENDER_AMBIENT) ? "enabled" : "disabled");
        break;
      case 't':
        if (renderFlags & RENDER_DIFFUSE) renderFlags &= ~RENDER_DIFFUSE;
        else renderFlags |= RENDER_DIFFUSE;
        printf("diffuse lighting: %s\n", (renderFlags & RENDER_DIFFUSE) ? "enabled" : "disabled");
        break;
      case 'y':
        if (renderFlags & RENDER_SPECULAR) renderFlags &= ~RENDER_SPECULAR;
        else renderFlags |= RENDER_SPECULAR;
        printf("specular lighting: %s\n", (renderFlags & RENDER_SPECULAR) ? "enabled" : "disabled");
        break;

        // +, - = 
      case '+':
      case '=':
        iso += 0.01f;
        printf("isosurface: f = %f\n", iso);
        break;
      case '-':
        iso -= 0.01f;
        printf("isosurface: f = %f\n", iso);
        break;

        // change color function
      case 'c':
        renderFlags = (renderFlags & ~0x18) | ((renderFlags + 0x08) & 0x18);
        if ((renderFlags & 0x18) == 0x00) printf("color function: position spectrum\n");
        if ((renderFlags & 0x18) == 0x08) printf("color function: normals\n");
        if ((renderFlags & 0x18) == 0x10) printf("color function: red\n");
        if ((renderFlags & 0x18) == 0x18) printf("color function: radial spectrum\n");
        break;
    }

    setConsts();
    glutPostRedisplay();
}

void special(int key, int x, int y)
{
    float dl = 0.1f;

    switch (key)
    {
        // arrow keys: move light along x/y
      case GLUT_KEY_UP:
        lightPos.y += dl;
        break;
      case GLUT_KEY_DOWN:
        lightPos.y -= dl;
        break;
      case GLUT_KEY_LEFT:
        lightPos.x -= dl;
        break;
      case GLUT_KEY_RIGHT:
        lightPos.x += dl;
        break;
    }

    setConsts();
    glutPostRedisplay();
}

// click locations and button state (for tracking mouse motion)
int ox, oy;
int buttonState = 0;

void mouse(int button, int state, int x, int y)
{
    if (state == GLUT_DOWN) buttonState |= 1 << button;
    else if (state == GLUT_UP) buttonState = 0;

    ox = x;
    oy = y;
    glutPostRedisplay();
}

void motion(int x, int y)
{
    float dx, dy;
    dx = x - ox;
    dy = y - oy;

    // TODO: change this stuff (use move modes)
    if (buttonState == 3) {
        // left+middle = zoom
        viewTranslation.z += dy / 100.0;
    }
    else if (buttonState & 2) {
        // middle = translate
        viewTranslation.x += dx / 100.0;
        viewTranslation.y -= dy / 100.0;
    }
    else if (buttonState & 1) {
        // left = rotate
        viewRotation.x += dy / 5.0;
        viewRotation.y += dx / 5.0;
    }

    ox = x;
    oy = y;
    glutPostRedisplay();
}

void reshape(int x, int y)
{
    // update windowWidth, windowHeight
    windowWidth = x;
    windowHeight = y;

    // reinitialize display state with new size
    initPixelBuffer();

    glViewport(0, 0, windowWidth, windowHeight);
    glMatrixMode(GL_MODELVIEW);
    glLoadIdentity();
}

////////////////////////////////////////////////////////////////////////////////
// Program main

// This function returns the best GPU (with maximum GFLOPS)
// (taken from cutil_inline_runtime.h)
int cutGetMaxGflopsDeviceId()
{
    int current_device   = 0, sm_per_multiproc = 0;
    int max_compute_perf = 0, max_perf_device  = 0;
    int device_count     = 0, best_SM_arch     = 0;
    int arch_cores_sm[3] = {1, 8, 32};
    cudaDeviceProp deviceProp;

    cudaGetDeviceCount(&device_count);

    // Find the best major SM Architecture GPU device
    while (current_device < device_count)
    {
        cudaGetDeviceProperties(&deviceProp, current_device);
        if (deviceProp.major > 0 && deviceProp.major < 9999)
            best_SM_arch = max(best_SM_arch, deviceProp.major);
        current_device++;
    }

    // Find the best CUDA capable GPU device
    for (current_device = 0; current_device < device_count; current_device++)
    {
        cudaGetDeviceProperties(&deviceProp, current_device);
        if (deviceProp.major == 9999 && deviceProp.minor == 9999)
            sm_per_multiproc = 1;
        else if (deviceProp.major <= 2)
            sm_per_multiproc = arch_cores_sm[deviceProp.major];
        else
            sm_per_multiproc = arch_cores_sm[2];

        int compute_perf = deviceProp.multiProcessorCount * sm_per_multiproc * deviceProp.clockRate;
        if (compute_perf > max_compute_perf)
        {
            // If we find GPU with SM major > 2, search only these
            if (best_SM_arch > 2)
            {
                // If our device==dest_SM_arch, choose this, or else pass
                if (deviceProp.major == best_SM_arch)
                {
                    max_compute_perf = compute_perf;
                    max_perf_device = current_device;
                }
            } else {
                max_compute_perf = compute_perf;
                max_perf_device = current_device;
            }
        }
    }
    return max_perf_device;
}

// (taken from cutil_inline_runtime.h)
void cutilDeviceInit(int deviceID)
{
    int deviceCount;
    safecall(cudaGetDeviceCount(&deviceCount));
    if (deviceCount == 0)
    {
        fprintf(stderr, "error: no devices supporting CUDA\n");
        exit(-1);
    }

    if (deviceID < 0) deviceID = 0;
    if (deviceID > deviceID - 1) deviceID = deviceCount - 1;
    cudaDeviceProp deviceProp;
    safecall(cudaGetDeviceProperties(&deviceProp, deviceID));
    if (deviceProp.major < 1)
    {
        fprintf(stderr, "error: GPU device does not support CUDA\n");
        exit(-1);
    }
    printf("using CUDA device %d: %s\n", deviceID, deviceProp.name);
    safecall(cudaSetDevice(deviceID));
}

int main(int argc, char* argv[])
{
    int deviceID = -1;

    // parse arguments
    // TODO: parse more arguments
    for (int x = 1; x < argc; x++)
    {
        if (argv[x][0] != '-') continue;
        switch (argv[x][1])
        {
          case 's':
          case 'S':
            fieldDimension = atoi(&argv[x][2]);
            break;
          case 'd':
          case 'D':
            deviceID = atoi(&argv[x][2]);
            break;
        }
    }

    // use command-line specified CUDA device, otherwise use device with highest Gflops/s
    if (deviceID != -1)
        cutilDeviceInit(deviceID);
    else {
        deviceID = cutGetMaxGflopsDeviceId();
        cudaSetDevice(deviceID);
    }

    struct cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, deviceID);
    if (!strncmp("Tesla", prop.name, 5))
    {
        printf("This sample needs a card capable of OpenGL and display.\n");
        printf("Please choose a different device with the -device=x argument.\n");
        exit(0);
    }

    srand(time(NULL));

    // initialize GLUT callback functions
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_RGB | GLUT_DOUBLE);
    glutInitWindowSize(windowWidth, windowHeight);
    glutCreateWindow("Uninitialized buffer rendering");
    glutDisplayFunc(display);
    glutKeyboardFunc(keyboard);
    glutSpecialFunc(special);
    glutMouseFunc(mouse);
    glutMotionFunc(motion);
    glutReshapeFunc(reshape);
    glutIdleFunc(idle);

    glewInit();
    if (!glewIsSupported("GL_VERSION_2_0 GL_ARB_pixel_buffer_object"))
    {
        fprintf(stderr, "Required OpenGL extensions missing.");
        exit(-1);
    }

    initCuda();
    initGLParameters();
    initPixelBuffer();

    /* Give c, plane initial values */
    setConsts();

    atexit(cleanup);

    glutMainLoop();

    cudaThreadExit();
    return 0;
}
