-- vs.lua
-- A first-order particle motion system, takes a time as input
-- and simply calculates point positions using a vertex shader.
-- Draws GL_POINTS. TODO: Use instanced rendering

local ffi = require("ffi")
local sf = require("util.shaderfunctions2")

require("util.objfile")

-- Naming this table 'vs' causes errors. Not sure why.
verts = {}
verts.__index = verts

function verts.new(...)
    local self = setmetatable({}, verts)
    if self.init ~= nil and type(self.init) == "function" then
        self:init(...)
    end 
    return self
end

local vert_header = [[
#version 310 es

in vec4 vPosition;
in vec4 vColor;
in vec4 vTexCoord;

out vec3 vfColor;
out vec3 vfTexCoord;
out vec3 vfWorldPos;
#line 1
]]

local vert_body = [[
#ifdef GL_ES
precision mediump float;
precision mediump int;
#endif
// spiral

uniform mat4 modelMtx;
uniform mat4 viewMtx;
uniform mat4 projMtx;

uniform mediump float time;
uniform int numParticles;
uniform float colorPhase;
uniform float colorSaturation;

vec3 pal( in float t, in vec3 a, in vec3 b, in vec3 c, in vec3 d )
{
    return a + b*cos( 6.28318*(c*t+d) );
}

vec3 getColor(float i)
{
    float colorPhase = -.4;
    return pal( i+colorPhase, vec3(0.5,0.5,0.5),vec3(0.5,0.5,0.5),vec3(1.0,1.0,0.5),vec3(0.8,0.90,0.30) );
}

// Quaternion rotation of a vector
vec3 qtransform(vec4 q, vec3 v)
{
    return v + 2.0*cross(cross(v, q.xyz) + q.w*v, q.xyz);
}

// http://www.geeks3d.com/20141201/how-to-rotate-a-vertex-by-a-quaternion-in-glsl/
vec4 quatFromAxisAngle(vec3 axis, float angle)
{
    float half_angle = angle / 2.;
    vec4 q;
    q.x = axis.x * sin(half_angle);
    q.y = axis.y * sin(half_angle);
    q.z = axis.z * sin(half_angle);
    q.w = cos(half_angle);
    return q;
}

const float PI = acos(-1.);

void main()
{
    int index = gl_InstanceID;
    float t = 1. / float(index);
    float ix = float(index) / float(numParticles);
    t += .1*float(index);
    t *= .1;

    vec4 position = vec4(vec3(0.),1.);
    float r = .2*pow(1.2,1.+t)-.2;
    //r += .1*pow(abs(sin(6.*2.*PI*t)));

    position.xyz = vec3(r*cos(2.*PI*t), r*sin(2.*PI*t),0.);
    position.z += .01*sin(100.*t+time);

    mat4 mvpMatrix = projMtx * viewMtx * modelMtx;
    mat3 normMtx = mat3(modelMtx);
    float th = 2. * acos(-1.) * t;
    vec4 q = normalize(quatFromAxisAngle(vec3(0.,0.,1.), -th));


    vec4 mPos = vPosition;
    vec3 mPos3 = qtransform(q,mPos.xyz);
    mPos = vec4(mPos3, 1.);

    vfColor = abs(normMtx * qtransform(q,vColor.rgb));
    vfTexCoord = mix(vec3(.8), getColor(fract(ix+colorPhase)), colorSaturation);
    vec3 wpos = (1.-ix)*.13*mPos.xyz + 4.*position.xyz;
    vfWorldPos = (modelMtx * vec4(wpos,1.)).xyz;
    gl_Position = mvpMatrix * vec4(wpos,1.);
}

]]

local basic_frag = [[
#version 310 es

#ifdef GL_ES
precision mediump float;
precision mediump int;
#endif

in vec3 vfWorldPos;
in vec3 vfColor;
out vec4 fragColor;

void main()
{
    fragColor = vec4(vfColor, 1.0);
}
]]

local lighting_frag = [[
#version 310 es

#ifdef GL_ES
precision mediump float;
precision mediump int;
#endif

uniform float time;
uniform float fogDensity;
uniform vec4 fogColor;

in vec3 vfWorldPos;
in vec3 vfColor; // normal
in vec3 vfTexCoord;

out vec4 fragColor;

float shininess = 125.;

vec3 directionalLight(vec3 normal)
{
    //vec3 lightDir = normalize(vec3(cos(time),1.,sin(time)));
    vec3 lightDir = normalize(vec3(0.,0.,1.));
    float intensity = max(dot(normal,lightDir), 0.);

#if 1
    float spec = 0.;
    if (intensity > 0.)
    {
        vec3 h = normalize(lightDir);
        float intSpec = max(dot(h,normal), 0.);
        spec = pow(intSpec, shininess);
    }
    return (intensity + spec)*vfTexCoord;
#endif    
}

vec3 pointLight(vec3 pos, vec3 normal)
{
    float r = 12.5;
    float speed = 2.;
    //vec3 lightPos = vec3(r*sin(speed*time), 5.*abs(sin(7.*time)), r*cos(speed*time));
    //vec3 lightPos = vec3(0., 4.*abs(sin(time)), 0.);
    vec3 lightPos = vec3(0., 0., 10.);

    vec3 lightDir = normalize(vec3(lightPos - pos));
    vec3 eye = vec3(0.);

    float spec = 0.;

    float intensity = max(dot(normal,lightDir), 0.);
    if (intensity > 0.)
    {
        vec3 h = normalize(lightDir + eye);
        float intSpec = max(dot(h,normal), 0.);
        spec = pow(intSpec, shininess);
    }
 
    return (intensity + spec)*vfTexCoord;
}

vec3 spotLight(vec3 pos, vec3 normal)
{
    float r = 12.5;
    float speed = 2.;
    //vec3 lightPos = vec3(r*sin(speed*time), 15., r*cos(speed*time));
    //vec3 lightPos = vec3(0., 4.*abs(sin(time)), 0.);
    vec3 lightPos = vec3(10., 15., 10.);

    vec3 spotDir = normalize(vec3(0.,-.0,1.));
    vec3 lightDir = normalize(vec3(lightPos - pos));
    vec3 eye = vec3(0.);

    float spec = 0.;

    float intensity = 0.;
    float spotCutoff = .75;
    if (dot(-spotDir,lightDir) > spotCutoff)
    {
        intensity = max(dot(normal,lightDir), 0.);
        if (intensity > 0.)
        {
            vec3 h = normalize(lightDir + eye);
            float intSpec = max(dot(h,normal), 0.);
            spec = pow(intSpec, shininess);
        }
    }
    return (intensity + spec)*vfTexCoord;
}

void main()
{
    vec3 vWorldPos = vfWorldPos;
    vec3 vVertCol = vfColor; // normal
    vVertCol += .001*vfTexCoord;

    vec3 normal = normalize(vVertCol);

    vec4 sceneCol = vec4(directionalLight(normal), 1.);
    //vec4 sceneCol = vec4(pointLight(vWorldPos, normal), 1.);
    //vec4 sceneCol = vec4(spotLight(vWorldPos, normal), 1.);

    // Fog
    // http://www.ozone3d.net/tutorials/glsl_fog/p04.php
    float density = fogDensity;
    const float LOG2 = 1.442695;
    float z = gl_FragCoord.z / gl_FragCoord.w;
    float fogFactor = exp2( -density * density * z * z * LOG2 );
    fogFactor = clamp(fogFactor, 0.0, 1.0);
    fragColor = mix(fogColor, sceneCol, fogFactor );
}
]]

function verts:init()
    self.vao = 0
    self.vbos = {}
    self.prog = 0
    self.npts = 1024*2
    self.time = 0

    self.src = vert_body
    self.descr = "Immediate particle system"
    self.shadertype = "vertex"
    self.dynamic_scale = true
    self.num_tris = 0
    self.chosen_frag = lighting_frag
    self.obj = nil
    self.fogDensity = .25
    self.colorPhase = 0
    self.colorSaturation = 0
end

function verts:setDataDirectory(dir)
    self.data_dir = dir
end

function verts:stateString()
    return self.npts..' points'
end

function verts:buildShader(src)
    -- Do we need to bind vao to compile a program?
    --gl.glBindVertexArray(self.vao)
    gl.glDeleteProgram(self.prog)

    self.prog,err = sf.make_shader_from_source({
        vsrc = vert_header..src,
        fsrc = self.chosen_frag,
        })
    if err then print(err) end
    --gl.glBindVertexArray(0)
    return err
end

local vertex_lists = {
    tetra = {
        verts = {
            0,2.598,0,
            0,0,1.732,
            -1.5,0,-.866,
            1.5,0,-.866,
        },
        norms = {
            1,1,1,
            1,0,0,
            0,1,0,
            0,0,1,
        },
        tris = {
            0,1,2,
            0,2,3,
            0,1,3,
            1,3,2,
        },
    },
    square = {
        verts = {
            0,0,0,
            1,0,0,
            1,1,0,
            0,1,0,
        },
        norms = {
        --[[0,0,1,
            1,0,1,
            1,1,1,
            0,1,0,]]
            0,0,1,
            0,0,1,
            0,0,1,
            0,0,1,
        },
        tris = {
            0,1,2,
            0,2,3,
        },
    },
    halfcube = {
        verts = {
            0,0,0,
            1,0,0,
            1,1,0,
            0,1,0,

            0,0,0,
            1,0,0,
            1,0,1,
            0,0,1,

            0,0,0,
            0,1,0,
            0,1,1,
            0,0,1,
        },
        norms = {
            0,0,1,
            0,0,1,
            0,0,1,
            0,0,1,

            0,1,0,
            0,1,0,
            0,1,0,
            0,1,0,

            1,0,0,
            1,0,0,
            1,0,0,
            1,0,0,
        },
        tris = {
            0,1,2,
            0,2,3,
            4,5,6,
            4,6,7,
            8,9,10,
            8,10,11,
        },
    },
    cube = {
        verts = {
            0,0,0,
            1,0,0,
            1,1,0,
            0,1,0,

            0,0,0,
            1,0,0,
            1,0,1,
            0,0,1,

            0,0,0,
            0,1,0,
            0,1,1,
            0,0,1,

            0,0,1,
            1,0,1,
            1,1,1,
            0,1,1,

            0,1,0,
            1,1,0,
            1,1,1,
            0,1,1,

            1,0,0,
            1,1,0,
            1,1,1,
            1,0,1,
        },
        norms = {
            0,0,1,
            0,0,1,
            0,0,1,
            0,0,1,

            0,1,0,
            0,1,0,
            0,1,0,
            0,1,0,

            1,0,0,
            1,0,0,
            1,0,0,
            1,0,0,
            -- negative, whatever
            0,0,1,
            0,0,1,
            0,0,1,
            0,0,1,

            0,1,0,
            0,1,0,
            0,1,0,
            0,1,0,

            1,0,0,
            1,0,0,
            1,0,0,
            1,0,0,
        },
        tris = {
            0,1,2,
            0,2,3,
            4,5,6,
            4,6,7,
            8,9,10,
            8,10,11,
            12,13,14,
            12,14,15,
            16,17,18,
            16,18,19,
            20,21,22,
            20,22,23,
        },
    }
}


-- TODO: Load obj for this
function verts:initSingleInstanceAttributes(dat)
    local glIntv = ffi.typeof('GLint[?]')
    local glUintv = ffi.typeof('GLuint[?]')
    local glFloatv = ffi.typeof('GLfloat[?]')

    local verts = glFloatv(#dat.verts, dat.verts)
    local cols = glFloatv(#dat.norms, dat.norms)

    local vpos_loc = gl.glGetAttribLocation(self.prog, "vPosition")
    local vcol_loc = gl.glGetAttribLocation(self.prog, "vColor")

    local vvbo = glIntv(0)
    gl.glGenBuffers(1, vvbo)
    gl.glBindBuffer(GL.GL_ARRAY_BUFFER, vvbo[0])
    gl.glBufferData(GL.GL_ARRAY_BUFFER, ffi.sizeof(verts), verts, GL.GL_STATIC_DRAW)
    gl.glVertexAttribPointer(vpos_loc, 3, GL.GL_FLOAT, GL.GL_FALSE, 0, nil)
    table.insert(self.vbos, vvbo)

    local cvbo = glIntv(0)
    gl.glGenBuffers(1, cvbo)
    gl.glBindBuffer(GL.GL_ARRAY_BUFFER, cvbo[0])
    gl.glBufferData(GL.GL_ARRAY_BUFFER, ffi.sizeof(cols), cols, GL.GL_STATIC_DRAW)
    gl.glVertexAttribPointer(vcol_loc, 3, GL.GL_FLOAT, GL.GL_FALSE, 0, nil)
    table.insert(self.vbos, cvbo)

    gl.glEnableVertexAttribArray(vpos_loc)
    gl.glEnableVertexAttribArray(vcol_loc)
    print(vpos_loc)
    print(vcol_loc)

    local tris = glIntv(#dat.tris, dat.tris)
    self.num_tris = #dat.tris
    local qvbo = glIntv(0)
    gl.glGenBuffers(1, qvbo)
    gl.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, qvbo[0])
    gl.glBufferData(GL.GL_ELEMENT_ARRAY_BUFFER, ffi.sizeof(tris), tris, GL.GL_STATIC_DRAW)
    table.insert(self.vbos, qvbo)
end

function verts:initSingleInstanceAttributesfromObj(filename)
    local glIntv = ffi.typeof('GLint[?]')
    local glUintv = ffi.typeof('GLuint[?]')
    local glFloatv = ffi.typeof('GLfloat[?]')

    -- Load obj file
    filename = 'models/'..filename..'.obj'
    if self.data_dir then filename = self.data_dir .. "/" .. filename end
    print('initSingleInstanceAttributesfromObj', filename)

    self.obj = objfile.new()
    print("Loading obj "..filename)
    self.obj:loadmodel(filename, true)
    if #self.obj.olist == 0 then return end
    
    local str = (#self.obj.vertlist/3).." vertices*3, "
    str = str..(#self.obj.normlist/3).." normals*3, "
    str = str..(#self.obj.texlist/3).." texcoords*3, "
    str = str..(#self.obj.idxlist/3).." ints*3 for triangle indices."
    print(str)

    -- Delete the old
    gl.glBindVertexArray(self.vao)
    for _,v in pairs(self.vbos) do
        gl.glDeleteBuffers(1,v)
    end
    self.vbos = {}

    --print("Number of objects: "..#self.obj.olist)
    local v = self.obj.vertlist
    local n = self.obj.normlist
    local i = self.obj.idxlist

    local verts = glFloatv(#v,v)
    local cols = glFloatv(#n,n)
    --local verts = glFloatv(#dat.verts, dat.verts)
    --local cols = glFloatv(#dat.norms, dat.norms)

    local vpos_loc = gl.glGetAttribLocation(self.prog, "vPosition")
    local vcol_loc = gl.glGetAttribLocation(self.prog, "vColor")

    local vvbo = glIntv(0)
    gl.glGenBuffers(1, vvbo)
    gl.glBindBuffer(GL.GL_ARRAY_BUFFER, vvbo[0])
    gl.glBufferData(GL.GL_ARRAY_BUFFER, ffi.sizeof(verts), verts, GL.GL_STATIC_DRAW)
    gl.glVertexAttribPointer(vpos_loc, 3, GL.GL_FLOAT, GL.GL_FALSE, 0, nil)
    table.insert(self.vbos, vvbo)

    local cvbo = glIntv(0)
    gl.glGenBuffers(1, cvbo)
    gl.glBindBuffer(GL.GL_ARRAY_BUFFER, cvbo[0])
    gl.glBufferData(GL.GL_ARRAY_BUFFER, ffi.sizeof(cols), cols, GL.GL_STATIC_DRAW)
    gl.glVertexAttribPointer(vcol_loc, 3, GL.GL_FLOAT, GL.GL_FALSE, 0, nil)
    table.insert(self.vbos, cvbo)

    gl.glEnableVertexAttribArray(vpos_loc)
    gl.glEnableVertexAttribArray(vcol_loc)

    local tris = glIntv(#i, i)
    self.num_tris = #i
    local qvbo = glIntv(0)
    gl.glGenBuffers(1, qvbo)
    gl.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, qvbo[0])
    gl.glBufferData(GL.GL_ELEMENT_ARRAY_BUFFER, ffi.sizeof(tris), tris, GL.GL_STATIC_DRAW)
    table.insert(self.vbos, qvbo)

    gl.glBindVertexArray(0)
end

function verts:initGL()
    local vaoId = ffi.new("int[1]")
    gl.glGenVertexArrays(1, vaoId)
    self.vao = vaoId[0]
    gl.glBindVertexArray(self.vao)
    local err = self:buildShader(vert_body)
    self:initSingleInstanceAttributes(vertex_lists.cube)
    gl.glBindVertexArray(0)
    if err then print(err) end
end

function verts:exitGL()
    gl.glBindVertexArray(self.vao)
    for _,v in pairs(self.vbos) do
        gl.glDeleteBuffers(1,v)
    end
    self.vbos = {}
    gl.glDeleteProgram(self.prog)
    local vaoId = ffi.new("GLuint[1]", self.vao)
    gl.glDeleteVertexArrays(1, vaoId)
    gl.glBindVertexArray(0)
end

function verts:renderEye(model, view, proj)
    if self.prog == 0 then return end
    local um_loc = gl.glGetUniformLocation(self.prog, "modelMtx")
    local uv_loc = gl.glGetUniformLocation(self.prog, "viewMtx")
    local up_loc = gl.glGetUniformLocation(self.prog, "projMtx")
    local un_loc = gl.glGetUniformLocation(self.prog, "numParticles")
    local ut_loc = gl.glGetUniformLocation(self.prog, "time")
    local uf_loc = gl.glGetUniformLocation(self.prog, "fogDensity")
    local ucp_loc = gl.glGetUniformLocation(self.prog, "colorPhase")
    local ucs_loc = gl.glGetUniformLocation(self.prog, "colorSaturation")

    gl.glUseProgram(self.prog)

    local glFloatv = ffi.typeof('GLfloat[?]')
    gl.glUniformMatrix4fv(um_loc, 1, GL.GL_FALSE, glFloatv(16, model))
    gl.glUniformMatrix4fv(uv_loc, 1, GL.GL_FALSE, glFloatv(16, view))
    gl.glUniformMatrix4fv(up_loc, 1, GL.GL_FALSE, glFloatv(16, proj))
    gl.glUniform1i(un_loc, self.npts)
    gl.glUniform1f(ut_loc, self.time)
    gl.glUniform1f(uf_loc, self.fogDensity)
    gl.glUniform1f(ucp_loc, self.colorPhase)
    gl.glUniform1f(ucs_loc, self.colorSaturation)

    --gl.glEnable(GL.GL_BLEND)
    --gl.glBlendFunc(GL.GL_SRC_ALPHA, GL.GL_ONE_MINUS_SRC_ALPHA)
    gl.glBindVertexArray(self.vao)
    if self.obj then
        for k,v in pairs(self.obj.olist) do
            if v[1] > 0 then
                gl.glDrawElementsInstanced(GL.GL_TRIANGLES, #self.obj.idxlist, GL.GL_UNSIGNED_INT, nil, self.npts)
            end
        end
    else
        gl.glDrawElementsInstanced(GL.GL_TRIANGLES, self.num_tris, GL.GL_UNSIGNED_INT, nil, self.npts)
    end
    gl.glBindVertexArray(0)
    --gl.glDisable(GL.GL_BLEND)

    gl.glUseProgram(0)
end

function verts:render_for_one_eye(view, proj)
    if self.prog ~= 0 then
        gl.glUseProgram(self.prog)

        local ut_loc = gl.glGetUniformLocation(self.prog, "time")
        gl.glUniform1f(ut_loc, self.time)
        local un_loc = gl.glGetUniformLocation(self.prog, "numParticles")
        gl.glUniform1i(un_loc, self.npts)
        local uf_loc = gl.glGetUniformLocation(self.prog, "fogDensity")
        gl.glUniform1f(uf_loc, self.fogDensity)

        local umv_loc = gl.glGetUniformLocation(self.prog, "mvmtx")
        local upr_loc = gl.glGetUniformLocation(self.prog, "prmtx")

        local glFloatv = ffi.typeof('GLfloat[?]')
        gl.glUniformMatrix4fv(umv_loc, 1, GL.GL_FALSE, glFloatv(16, view))
        gl.glUniformMatrix4fv(upr_loc, 1, GL.GL_FALSE, glFloatv(16, proj))
        gl.glBindVertexArray(self.vao)
        gl.glDrawElementsInstanced(GL.GL_TRIANGLES, 4*3, GL.GL_UNSIGNED_INT, nil, self.npts)

        gl.glBindVertexArray(0)
        gl.glUseProgram(0)
    end
end

function verts:timestep(absTime, dt)
end

function verts:onSingleTouch(pointerid, action, x, y)
    --print("points.onSingleTouch",pointerid, action, x, y)
end

function verts:scaleBufferSize(k)
    self.dynamic_scale =  false
    self.npts = self.npts * k
end

function verts:charkeypressed(key)
    if key == '-' then
        self:scaleBufferSize(.5)
    elseif key == '=' then
        self:scaleBufferSize(2)
    end
end

function verts:keypressed(key, scancode, action, mods)
    local func_table = {
        [298] = function (x) self:scaleBufferSize(.5*.5) end, -- F9
        [299] = function (x) self:scaleBufferSize(.5) end, -- F10
        [300] = function (x) self:scaleBufferSize(2) end, -- F11
        [301] = function (x) self:scaleBufferSize(2*2) end, -- F12
    }
    local f = func_table[key]
    if f then f() handled = true end
end

function verts:handleCommand(args)
    if not args then return {'nil args'} end
    if #args == 0 then return {'empty args'} end
    if #args == 1 then return {args[1]..'no args'} end

    local key = args[1]
    local val = args[2]
    if key and val then
        if key == 'dynscale' then
            self.dynamic_scale = true -- TODO size
            return {'set dynamic scale'}
        elseif key == 'frag' then
            if val == 'basic' then
                self.chosen_frag = basic_frag
                return {'basic frag shader'}
            elseif val == 'lighting' then
                self.chosen_frag = lighting_frag
                return {'lighting frag shader'}
            end
        elseif key == 'shape' then
            local vlist = vertex_lists[val]
            if vlist then
                gl.glBindVertexArray(self.vao)
                for _,v in pairs(self.vbos) do
                    gl.glDeleteBuffers(1,v)
                end
                self.vbos = {}
                self:initSingleInstanceAttributes(vlist)
                gl.glBindVertexArray(0)
                return {'set shape to '..val}
            end
        elseif key == 'obj' then
            self:initSingleInstanceAttributesfromObj(val)
            return {'set obj to '..val}
        elseif key == 'instances' then
            local n = tonumber(val)
            if n then
                self.dynamic_scale = false
                self.npts = n
                return {'set n points'}
            end
        elseif key == 'fog' then
            local n = tonumber(val)
            if n then
                self.fogDensity = n
                return {'set fog density'}
            end
        end
    end
end

-- Return a list of strings(lines) describing commands
function verts:commandHelp()
    return {
        'vs - vertex shader help',
        'Commands:',
        '   dynscale <sz> ',
        '   frag <basic|lighting>',
        '   shape <tetra|cube|halfcube|square>',
        '   fog <density> '..self.fogDensity,
        'Instances: '..tostring(self.npts),
    }
end

return verts
