-- comp_scene.lua
comp_scene = {}

local openGL = require("opengl")
local ffi = require("ffi")
local sf = require("util.shaderfunctions")
local mm = require("util.matrixmath")

-- Types from:
-- https://github.com/nanoant/glua/blob/master/init.lua
local glIntv     = ffi.typeof('GLint[?]')
local glUintv    = ffi.typeof('GLuint[?]')
local glCharv    = ffi.typeof('GLchar[?]')
local glSizeiv   = ffi.typeof('GLsizei[?]')
local glFloatv   = ffi.typeof('GLfloat[?]')
local glConstCharpp = ffi.typeof('const GLchar *[1]')

local g = 0.0
local vao = 0
local prog_display = 0
local prog_accel = 0
local prog_acceltiled = 0
local prog_integrate = 0

local pt_vert = [[
#version 430
layout(location = 0) in vec4 vposition;
void main() {
   gl_Position = vposition;
}
]]

local bb_geom = [[
#version 430
layout(location = 0) uniform mat4 View;
layout(location = 1) uniform mat4 Projection;
layout (points) in;
layout (triangle_strip, max_vertices = 4) out;
out vec2 txcoord;
void main() {
   vec4 pos = View*gl_in[0].gl_Position;
   txcoord = vec2(-1,-1);
   gl_Position = Projection*(pos+0.2*vec4(txcoord,0,0));
   EmitVertex();
   txcoord = vec2( 1,-1);
   gl_Position = Projection*(pos+0.2*vec4(txcoord,0,0));
   EmitVertex();
   txcoord = vec2(-1, 1);
   gl_Position = Projection*(pos+0.2*vec4(txcoord,0,0));
   EmitVertex();
   txcoord = vec2( 1, 1);
   gl_Position = Projection*(pos+0.2*vec4(txcoord,0,0));
   EmitVertex();
}
]]

local rad_frag = [[
#version 330
in vec2 txcoord;
layout(location = 0) out vec4 FragColor;
void main() {
   float s = (1/(1+15.*dot(txcoord, txcoord))-1/16.);
   FragColor = s*vec4(0.3,0.3,1.0,1);
}
]]

local accel_comp = [[
#version 430
layout(local_size_x=256) in;

layout(location = 0) uniform float dt;
layout(std430, binding=0) buffer pblock { vec4 positions[]; };
layout(std430, binding=1) buffer vblock { vec4 velocities[]; };

void main() {
   int N = int(gl_NumWorkGroups.x*gl_WorkGroupSize.x);
   int index = int(gl_GlobalInvocationID);

   vec3 position = positions[index].xyz;
   vec3 velocity = velocities[index].xyz;
   vec3 acceleration = vec3(0,0,0);
   for(int i = 0;i<N;++i) {
       vec3 other = positions[i].xyz;
       vec3 diff = position - other;
       float invdist = 1.0/(length(diff)+0.001);
       acceleration -= diff*0.1*invdist*invdist*invdist;
   }
   velocities[index] = vec4(velocity+dt*acceleration,0);
}
]]

local accel_tiled_comp = [[
#version 430
layout(local_size_x=256) in;

layout(location = 0) uniform float dt;
layout(std430, binding=0) buffer pblock { vec4 positions[]; };
layout(std430, binding=1) buffer vblock { vec4 velocities[]; };

shared vec4 tmp[gl_WorkGroupSize.x];
void main() {
   int N = int(gl_NumWorkGroups.x*gl_WorkGroupSize.x);
   int index = int(gl_GlobalInvocationID);
   vec3 position = positions[index].xyz;
   vec3 velocity = velocities[index].xyz;
   vec3 acceleration = vec3(0,0,0);
   for(int tile = 0;tile<N;tile+=int(gl_WorkGroupSize.x)) {
       tmp[gl_LocalInvocationIndex] = positions[tile + int(gl_LocalInvocationIndex)];
       groupMemoryBarrier();
       barrier();
       for(int i = 0;i<gl_WorkGroupSize.x;++i) {
           vec3 other = tmp[i].xyz;
           vec3 diff = position - other;
           float invdist = 1.0/(length(diff)+0.001);
           acceleration -= diff*0.1*invdist*invdist*invdist;
       }
       groupMemoryBarrier();
       barrier();
   }
   velocities[index] = vec4(velocity+dt*acceleration,0);
}
]]

local integ_comp = [[
#version 430
layout(local_size_x=256) in;

layout(location = 0) uniform float dt;
layout(std430, binding=0) buffer pblock { vec4 positions[]; };
layout(std430, binding=1) buffer vblock { vec4 velocities[]; };

void main() {
   int index = int(gl_GlobalInvocationID);
   vec4 position = positions[index];
   vec4 velocity = velocities[index];
   position.xyz += dt*velocity.xyz;
   positions[index] = position;
}
]]

local particles = 8*1024
function init_point_attributes()
    pos_array = ffi.new("float[?]", particles*4)
    vel_array = ffi.new("float[?]", particles*4)
    for i=0,particles-1 do
        pos_array[4*i+0] = 2*math.random()-1
        pos_array[4*i+1] = 0.4*math.random()-0.2
        pos_array[4*i+2] = 2*math.random()-1
        pos_array[4*i+3] = 1
    end
    for i=0,particles*4-1 do
        vel_array[i] = 0
    end
    
    local vboIds = ffi.new("int[2]")
    gl.GenBuffers(2, vboIds)
    
    local vboP = vboIds[0]
    local vboV = vboIds[1]
    gl.BindBuffer(GL.SHADER_STORAGE_BUFFER, vboV)
    gl.BufferData(GL.SHADER_STORAGE_BUFFER, ffi.sizeof(vel_array), vel_array, GL.STATIC_DRAW)
    
    gl.BindBuffer(GL.ARRAY_BUFFER, vboP)
    gl.BufferData(GL.ARRAY_BUFFER, ffi.sizeof(pos_array), pos_array, GL.STATIC_DRAW)
    
    gl.EnableVertexAttribArray(0)
    gl.VertexAttribPointer(0, 4, GL.FLOAT, GL.FALSE, 0, nil);
    
    -- causes crash
    --gl.BindBuffersBase(GL.SHADER_STORAGE_BUFFER, 0, 2, vboIds)
    gl.BindBufferBase(GL.SHADER_STORAGE_BUFFER, 0, vboP)
    gl.BindBufferBase(GL.SHADER_STORAGE_BUFFER, 1, vboV)
    
    local dt = 1/60
    
    gl.UseProgram(prog_accel)
    gl.Uniform1f(0, dt)
    gl.UseProgram(prog_acceltiled)
    gl.Uniform1f(0, dt)
    gl.UseProgram(prog_integrate)
    gl.Uniform1f(0, dt)
end

function comp_scene.initGL()
    local vaoId = ffi.new("int[1]")
    gl.GenVertexArrays(1, vaoId)
    vao = vaoId[0]
    gl.BindVertexArray(vao)

    prog_display = sf.make_shader_from_source({
        vsrc = pt_vert,
        gsrc = bb_geom,
        fsrc = rad_frag,
        })

    prog_accel = sf.make_shader_from_source({
        compsrc = accel_comp,
        })
    prog_acceltiled = sf.make_shader_from_source({
        compsrc = accel_tiled_comp,
        })
    prog_integrate = sf.make_shader_from_source({
        compsrc = integ_comp,
        })

    init_point_attributes()
    gl.BindVertexArray(0)

    gl.Disable(GL.DEPTH_TEST)
    gl.Enable(GL.BLEND)
    gl.BlendFunc(GL.ONE, GL.ONE)
end

function draw_points()
    gl.BindVertexArray(vao)
    gl.DrawArrays(GL.POINTS, 0, particles)
    gl.BindVertexArray(0)
end

function comp_scene.render_for_one_eye(mview, proj)
    -- Cast the cdata ptr to a table for further manipulation here in Lua
    local m0 = ffi.cast("float*", mview)
    -- Once we have done so, the table is 0-indexed. here we clumsily jam it into
    -- a Lua-style, 1-indexed table.
    local m = {}
    for i=0,15 do m[i+1] = m0[i] end

    --glh_rotate(m, g*0.1, 0,1,0)
    mm.glh_translate(m, -0.5, 0.1*math.sin(g), -0.5)

    gl.UseProgram(prog_display)
    
    -- Override matrices to example
    local e_p = {}
    mm.glh_perspective(e_p, 90.0, 4.0/3.0, 0.1, 100.0)
    local e_m = {}
    for i=1,16 do e_m[i]=1 end
    mm.make_translation_matrix(e_m, 0, 0, -30)
    gl.UniformMatrix4fv(0, 1, GL.FALSE, glFloatv(16, e_m))
    gl.UniformMatrix4fv(1, 1, GL.FALSE, glFloatv(16, e_p))

    draw_points()
    sf.CHECK_GL_ERROR()

    gl.UseProgram(0)
end

function comp_scene.timestep(dt)
    g = g + dt
    
    --gl.UseProgram(prog_acceltiled)
    --gl.DispatchCompute(particles/256, 1, 1)
    gl.UseProgram(prog_accel)
    gl.DispatchCompute(particles/256, 1, 1)

    gl.UseProgram(prog_integrate)
    gl.DispatchCompute(particles/256, 1, 1)
end

return comp_scene
