<reponame>ISISNeutronMuon/SScanSS-2
import ctypes
from contextlib import suppress
from OpenGL import GL, error
from OpenGL.GL import shaders


DEFAULT_VERTEX_SHADER = """
#version 120
attribute vec3 position;
varying vec4 colour;

void main(void)
{	
  colour = gl_Color;
  gl_Position = gl_ModelViewProjectionMatrix * vec4(position, 1);
}
"""


DEFAULT_FRAGMENT_SHADER = """
#version 120

varying vec4 colour;
void main (void)
{
  gl_FragColor = colour;			
}
"""


GOURAUD_VERTEX_SHADER = """
#version 120

#define NUM_LIGHTS {0}
attribute vec3 position;
attribute vec3 vnormal;

/*******************************************************
*  Fixed.vert Fixed Function Equivalent Vertex Shader  *
*        Automatically Generated by GLSL ShaderGen     *
*          https://github.com/mojocorp/ShaderGen       *
*******************************************************/

vec4 Ambient;
vec4 Diffuse;
vec4 Specular;

void directionalLight(in int i, in vec3 normal)
{{
   float nDotVP;         // normal . light direction
   float nDotHV;         // normal . light half vector
   float pf;             // power factor

   nDotVP = max(0.0, dot(normal, normalize(vec3 (gl_LightSource[i].position))));
   nDotHV = max(0.0, dot(normal, vec3 (gl_LightSource[i].halfVector)));

   if (nDotVP == 0.0)
   {{
       pf = 0.0;
   }}
   else
   {{
       pf = pow(nDotHV, gl_FrontMaterial.shininess);
   }}

   Ambient  += gl_LightSource[i].ambient;
   Diffuse  += gl_LightSource[i].diffuse * nDotVP;
   Specular += gl_LightSource[i].specular * pf;
}}


void flight(in vec3 normal, in vec4 ecPosition, float alphaFade)
{{
    vec4 color;
    vec3 ecPosition3;
    vec3 eye;

    ecPosition3 = vec3(ecPosition) / ecPosition.w;
    eye = vec3 (0.0, 0.0, 1.0);

    // Clear the light intensity accumulators
    Ambient  = vec4 (0.0);
    Diffuse  = vec4 (0.0);
    Specular = vec4 (0.0);

    for (int i=0; i < NUM_LIGHTS; i++)
        directionalLight(i, normal);

    color = gl_FrontLightModelProduct.sceneColor +
            Ambient * gl_Color +
            Diffuse * gl_Color;

    // color += Specular * gl_FrontMaterial.specular;
    color = clamp( color, 0.0, 1.0 );

    gl_FrontColor = color;
    gl_FrontColor.a *= alphaFade;
}}


void main (void)
{{
    float alphaFade = 1.0;

    // Do fixed functionality vertex transform
    gl_Position = gl_ModelViewProjectionMatrix * vec4(position, 1); 

    // Eye-coordinate position of vertex, needed in various calculations
    vec4 ecPosition = gl_ModelViewMatrix * vec4(position, 1);

    vec3  transformedNormal = normalize(gl_NormalMatrix * vnormal);
    flight(transformedNormal, ecPosition, alphaFade);
}}
"""


GOURAUD_FRAGMENT_SHADER = """
#version 120
/*******************************************************
* Fixed.frag Fixed Function Equivalent Fragment Shader *
*        Automatically Generated by GLSL ShaderGen     *
*          https://github.com/mojocorp/ShaderGen       *
*******************************************************/

void main (void) 
{
    vec4 color;
    color = gl_Color;
    gl_FragColor = color;		
}
"""


class Shader:
    """Base class for a GLSL program

    :param vertex_shader: source code for vertex shaders
    :type vertex_shader: str
    :param fragment_shader: source code for fragment shaders
    :type fragment_shader: str
    """
    def __init__(self, vertex_shader, fragment_shader):
        self.id = shaders.compileProgram(shaders.compileShader(vertex_shader, GL.GL_VERTEX_SHADER),
                                         shaders.compileShader(fragment_shader, GL.GL_FRAGMENT_SHADER))

    def destroy(self):
        """Deletes the shader program"""
        GL.glDeleteProgram(self.id)

    def bind(self):
        """Sets program associated with this object as active program in the
        current OpenGL context"""
        GL.glUseProgram(self.id)

    def release(self):
        """Releases the active shader program in the current OpenGL context"""
        GL.glUseProgram(0)


class DefaultShader(Shader):
    """Creates a GLSL program the renders primitives with colour"""
    def __init__(self):
        super().__init__(DEFAULT_VERTEX_SHADER, DEFAULT_FRAGMENT_SHADER)


class GouraudShader(Shader):
    """Creates a GLSL program the renders primitive with Gouraud shading

    :param number_of_lights: number of lights in the scene
    :type number_of_lights: int
    """
    def __init__(self, number_of_lights):
        vertex_shader = GOURAUD_VERTEX_SHADER.format(number_of_lights)

        super().__init__(vertex_shader, GOURAUD_FRAGMENT_SHADER)


class VertexArray:
    def __init__(self, vertices, indices, normals):
        """Creates buffers for vertex, normal, and element attribute data

        :param vertices: N x 3 array of vertices
        :type vertices: numpy.ndarray
        :param indices: M x 1 array of vertices
        :type indices: numpy.ndarray
        :param normals: N x 3 array of normal
        :type normals: numpy.ndarray
        """
        self.count = len(indices)
        self.buffers = []

        self.vertex_buffer = GL.glGenBuffers(1)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vertex_buffer)
        GL.glBufferData(GL.GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL.GL_STATIC_DRAW)
        self.buffers.append(self.vertex_buffer)

        if len(normals) > 0:
            self.normal_buffer = GL.glGenBuffers(1)
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.normal_buffer)
            GL.glBufferData(GL.GL_ARRAY_BUFFER, normals.nbytes, normals.flatten(), GL.GL_STATIC_DRAW)
            self.buffers.append(self.normal_buffer)

        self.element_buffer = GL.glGenBuffers(1)
        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, self.element_buffer)
        GL.glBufferData(GL.GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices, GL.GL_STATIC_DRAW)
        self.buffers.append(self.element_buffer)

        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, 0)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, 0)

    def __del__(self):
        with suppress(error.Error, ctypes.ArgumentError):
            GL.glDeleteBuffers(len(self.buffers), self.buffers)

    def bind(self):
        """Binds the buffers associated with this object to the current OpenGL context"""
        GL.glEnableVertexAttribArray(0)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vertex_buffer)
        GL.glVertexAttribPointer(0, 3, GL.GL_FLOAT, GL.GL_FALSE, 12, ctypes.c_void_p(0))

        if len(self.buffers) == 3:
            GL.glEnableVertexAttribArray(1)
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.normal_buffer)
            GL.glVertexAttribPointer(1, 3, GL.GL_FLOAT, GL.GL_FALSE, 12, ctypes.c_void_p(0))

        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, self.element_buffer)

    def release(self):
        """Releases the buffers associated with this object from the current OpenGL context"""
        GL.glDisableVertexAttribArray(0)
        GL.glDisableVertexAttribArray(1)
        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, 0)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, 0)
