import glfw
from OpenGL.GL import *
import math
import numpy

# global variables for window info, camera info
resx = 1000   # resolution of window horizontal direction
resy = 600   # resoltuion vertical direction
prevx = 0.0
prevy = 0.0

VertexArrayID = -1;
vertexbuffer = -1;


computeProgramID = -1
computerRenderProgramID = -1;
texture = -1
texture3D = -1


PI = numpy.pi

g_vertex_buffer_data = numpy.array([
    -1.0, -1.0, 0.0,
     1.0, -1.0, 0.0,
    -1.0,  1.0, 0.0,
     1.0,  1.0, 0.0,
], numpy.float32)

g_uv_buffer_data = numpy.array([
    0.0, 0.0,
    1.0, 0.0,
    0.0, 1.0,
    1.0, 1.0,
], numpy.float32)

vertexbuffer_quad = -1
uvbuffer_quad = -1


pos = numpy.array([-2.0, -2.0, 0.0], numpy.float32);
f = numpy.array([0,1,0], numpy.float32)
u = numpy.array([0,0,1], numpy.float32)
r = numpy.array([1,0,0], numpy.float32)
phi = numpy.pi/4;
theta = numpy.pi/2;
fov = 60.0;

Nx = 512
Ny = 512
Nz = 512

voxels = numpy.zeros(Nx*Ny*Nz, dtype=numpy.uint8)

window = -1

def main():
    global VertexArrayID, texture, texture3D, window, PI, f,r, u, fov, theta, phi, pos
    global vertexbuffer_quad, uvbuffer_quad, computeProgramID, computerRenderProgramID

    # DOWNLOAD THIS FILE FIRST: https://www.digitalrocksportal.org/projects/16/images/65564/download/
    with(open("segmented_castle_512.ubc", "rb")) as f:
        voxels = numpy.frombuffer(f.read(), dtype=numpy.uint8)
    voxels = 255*(1 - voxels)

    print(len(voxels), type(voxels))
    print(voxels[:100])


    # Initialize the library
    if not glfw.init():
        return

    glfw.window_hint(glfw.SAMPLES, 4);
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 4);
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3);
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE);

    # Create a windowed mode window and its OpenGL context
    window = glfw.create_window(resx, resy, "Hello World", None, None)
    if not window:
        glfw.terminate()
        return

    # Make the window's context current
    glfw.make_context_current(window)

    # Tell gLFW which callback function to use on input
    glfw.set_key_callback(window, key_callback)
    glfw.set_mouse_button_callback(window, mousebutton_callback)
    glfw.set_cursor_pos_callback(window, mouseposition_callback)
    glfw.set_scroll_callback(window, mousewheel_callback)
    glfw.set_window_size_callback(window, windowsize_callback)

    glViewport(0, 0, resx, resy);
    glClearColor(0.2, 0.2, 0.2, 1.0)
    glfw.swap_interval(0);
    VertexArrayID = glGenVertexArrays(1);
    glBindVertexArray(VertexArrayID);

    computerRenderProgramID = loadShaders("vertex_shader_render.vs", "fragment_shader_render.fs")

    vertexbuffer_quad = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer_quad)
    glBufferData(GL_ARRAY_BUFFER, g_vertex_buffer_data, GL_STATIC_DRAW)

    uvbuffer_quad = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, uvbuffer_quad)
    glBufferData(GL_ARRAY_BUFFER, g_uv_buffer_data, GL_STATIC_DRAW)


    texture = glGenTextures(1)
    glActiveTexture(GL_TEXTURE0)
    glBindTexture(GL_TEXTURE_2D, texture)
    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_RGBA32F, resx, resy, 0, GL_RED, GL_FLOAT, None);
    glBindImageTexture(0, texture, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F);
    computeProgramID = loadComputeShader("compute_shader.cs", texture);


    texture3D = glGenTextures(1);
    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_3D, texture3D);
    glTexImage3D(GL_TEXTURE_3D, 0, GL_RED, Nx, Ny, Nz, 0, GL_RED, GL_UNSIGNED_BYTE, voxels);
    glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);



    # Loop until the user closes the window
    while not glfw.window_should_close(window):
        # Render here, e.g. using pyOpenGL

        cameraStuff()

        glClear(GL_COLOR_BUFFER_BIT)



        aspect = resx/resy;
        FOV = fov*PI/180.0;
        tan_psi = numpy.tan(FOV/2.0);
        tan_phi = numpy.tan(FOV/2.0)*aspect;



        glUseProgram(computeProgramID);
        loc = glGetUniformLocation(computeProgramID, "orig");
        if (loc != -1): glUniform3f(loc, pos[0], pos[1], pos[2]);
        loc = glGetUniformLocation(computeProgramID, "N");
        if (loc != -1): glUniform3i(loc, Nx, Ny, Nz);
        loc = glGetUniformLocation(computeProgramID, "chosenFace");
        if (loc != -1): glUniform1i(loc, -1);

        loc = glGetUniformLocation(computeProgramID, "chosenVoxel");
        if (loc != -1): glUniform3i(loc, 0, 0, 0);
        loc = glGetUniformLocation(computeProgramID, "f");
        if (loc != -1): glUniform3f(loc, f[0], f[1], f[2]);
        loc = glGetUniformLocation(computeProgramID, "r");
        if (loc != -1): glUniform3f(loc, r[0], r[1], r[2]);
        loc = glGetUniformLocation(computeProgramID, "u");
        if (loc != -1): glUniform3f(loc, u[0], u[1], u[2]);

        loc = glGetUniformLocation(computeProgramID, "tan_phi");
        if (loc != -1): glUniform1f(loc, tan_phi);

        loc = glGetUniformLocation(computeProgramID, "tan_psi");
        if (loc != -1): glUniform1f(loc, tan_psi);




        glDispatchCompute(int(resx)//8, int(resy)//8, 1);





        glUseProgram(computerRenderProgramID);

        glUseProgram(computerRenderProgramID);

        glEnableVertexAttribArray(0);
        glEnableVertexAttribArray(1);
        glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer_quad);
        glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,0,None);
        glBindBuffer(GL_ARRAY_BUFFER, uvbuffer_quad);
        glVertexAttribPointer(1,2,GL_FLOAT,GL_FALSE,0,None);
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

        glDisableVertexAttribArray(1);
        glDisableVertexAttribArray(0);



        # Swap front and back buffers
        glfw.swap_buffers(window)

        # Poll for and process events, call relevant callback function
        glfw.poll_events()

    glfw.terminate()

t1 = glfw.get_time()
dt = 1.0/60.0

def cameraStuff():
    global t1, pos, f, u, r, window, dt
    t2 = glfw.get_time()
    speed = (t2-t1)*2;
    alpha = 0.01
    dt = alpha*(t2-t1) + (1-alpha)*dt
    glfw.set_window_title(window, "fps = {:.3f}".format(1.0/dt))
    t1 = t2


    if (glfw.get_key(window, glfw.KEY_LEFT_CONTROL) == glfw.PRESS):
        speed *= 10;

    if (glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS):
        speed *= 10;


    deltaMoveForward = speed*(glfw.get_key(window, glfw.KEY_W) - glfw.get_key(window, glfw.KEY_S));
    deltaMoveRight   = speed*(glfw.get_key(window, glfw.KEY_D) - glfw.get_key(window, glfw.KEY_A));
    deltaMoveUp      = speed*(glfw.get_key(window, glfw.KEY_E) - glfw.get_key(window, glfw.KEY_Q));


    sinp = numpy.sin(phi)
    cosp = numpy.cos(phi)
    sint = numpy.sin(theta)
    cost = numpy.cos(theta);
    f = numpy.array((cosp*sint, sinp*sint, cost), numpy.float32)
    r = numpy.array((sinp, -cosp, 0.0), numpy.float32)
    u = numpy.array((-cosp*cost, -sinp*cost, sint), numpy.float32)

    pos = pos + deltaMoveForward*f + deltaMoveRight*r + deltaMoveUp*u;


def loadComputeShader(computeShaderPath, texture):
    ComputeShaderID   = glCreateShader(GL_COMPUTE_SHADER);

    with open(computeShaderPath, "r") as f:
        shaderCode = f.read()

    glShaderSource(ComputeShaderID, shaderCode)
    glCompileShader(ComputeShaderID)

    Result = glGetShaderiv(ComputeShaderID, GL_COMPILE_STATUS);

    if Result == GL_FALSE:
        print(glGetShaderInfoLog(ComputeShaderID))


    ProgramID = glCreateProgram()
    glAttachShader(ProgramID, ComputeShaderID);
    glLinkProgram(ProgramID);

    Result = glGetProgramiv(ProgramID, GL_LINK_STATUS);

    if Result == GL_FALSE:
        print(glGetProgramInfoLog(ProgramID))

    glDeleteShader(ComputeShaderID);


    glUseProgram(ProgramID);

    glUniform1i(glGetUniformLocation(ProgramID, "framebufferImage"), 0);
    glUniform1i(glGetUniformLocation(ProgramID, "voxelTextureSampler"), 1);
    glUniform1i(glGetUniformLocation(ProgramID, "voxelImage"), 3);
    glUniform1i(glGetUniformLocation(ProgramID, "cubemapSampler"), 2);


    return ProgramID


def loadShaders(vertexShaderPath, fragmentShaderPath):
    VertexShaderID   = glCreateShader(GL_VERTEX_SHADER);

    with open(vertexShaderPath, "r") as f:
        shaderCode = f.read()

    glShaderSource(VertexShaderID, shaderCode)
    glCompileShader(VertexShaderID)

    Result = glGetShaderiv(VertexShaderID, GL_COMPILE_STATUS);

    if Result == GL_FALSE:
        print(glGetShaderInfoLog(VertexShaderID))


    FragmentShaderID = glCreateShader(GL_FRAGMENT_SHADER);

    with open(fragmentShaderPath, "r") as f:
        shaderCode = f.read()

    glShaderSource(FragmentShaderID, shaderCode)
    glCompileShader(FragmentShaderID)

    Result = glGetShaderiv(FragmentShaderID, GL_COMPILE_STATUS);

    if Result == GL_FALSE:
        print(glGetShaderInfoLog(FragmentShaderID))

    ProgramID = glCreateProgram()
    glAttachShader(ProgramID, VertexShaderID);
    glAttachShader(ProgramID, FragmentShaderID);
    glLinkProgram(ProgramID);

    Result = glGetProgramiv(ProgramID, GL_LINK_STATUS);

    if Result == GL_FALSE:
        print(glGetProgramInfoLog(ProgramID))

    glDeleteShader(VertexShaderID);
    glDeleteShader(FragmentShaderID);

    return ProgramID



# Callback function called every time the window size change
# Adjusts the camera width and heigh so that the scale stays the same
# Resets projection matrix
def windowsize_callback(window, width, height):
    global resx, resy

    resx = width;
    resy = height;

    glViewport(0, 0, resx, resy);
    glBindTexture(GL_TEXTURE_2D, texture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, resx, resy, 0, GL_RED, GL_FLOAT, None);



# Callback function called every time a keyboard key is pressed, released or held down
def key_callback(window, key, scancode, action, mods):
    pass
    #print("Key: key = {}, scancode = {}, action = {}, mods = {}".format(key, scancode, action, mods))


# Callback function called every time a mouse button pressed or released
def mousebutton_callback(window, button, action, mods):
    global prevx, prevy

    # get current cursor position, used in mouseposition_callback
    # to know how far the mouse moved
    prevx, prevy = glfw.get_cursor_pos(window)



    #print("Mousebutton: button = {}, action = {}, mods = {}, prevx = {}, prevy = {}".format(button, action, mods, prevx, prevy))

# Callback function called every time a the mouse is moved
def mouseposition_callback(window, xpos, ypos):
    global prevx, prevy, phi, theta

    # move the camera if the first mouse button is held down
    # the cursor will stay on at the same location relative to world coordinates after movement

    if  glfw.get_mouse_button(window, 0):
        anglesPerPixel = 0.1 * PI / 180.0;
        phi = numpy.fmod(phi - (xpos-prevx) * anglesPerPixel + 2*PI, 2*PI);
        theta = max(1.0*PI/180.0, min(179.0*PI/180.0, theta + (ypos-prevy) * anglesPerPixel));

        prevx = xpos
        prevy = ypos

    #print("Mousepos: xpos = {}, ypos = {}".format(xpos, ypos))

# Callback function called every time a the mouse scroll wheel is moved
# yoffset = up-down
# xoffset = left-right
def mousewheel_callback(window, xoffset, yoffset):
    global prevx, prevy, fov
    zoomFactor = math.pow(0.95,yoffset);

    fov = max(15, min(120.0, fov*zoomFactor));

if __name__ == "__main__":
    main()