import math
import random
import sys
import time

import numpy as np
import OpenGL
from OpenGL.GL import *
import pygame
from pygame.locals import *

def createCircles(numCircles):
    circles = []
    minRadius = 5
    maxRadius = 20
    for i in range(numCircles):
        cx = random.random()*width
        cy = random.random()*height
        r = random.random()*(maxRadius - minRadius) + minRadius
        circles.append((cx, cy, r))
    return circles

def createShader(vsStr, fsStr):
    vs = glCreateShader(GL_VERTEX_SHADER)
    glShaderSource(vs, vsStr)
    glCompileShader(vs)
    check = glGetShaderiv(vs, GL_COMPILE_STATUS)
    if not(check):
        raise RuntimeError(glGetShaderInfoLog(vs))

    fs = glCreateShader(GL_FRAGMENT_SHADER)
    glShaderSource(fs, fsStr)
    glCompileShader(fs)
    check = glGetShaderiv(fs, GL_COMPILE_STATUS)
    if not(check):
        raise RuntimeError(glGetShaderInfoLog(fs))

    program = glCreateProgram()
    glAttachShader(program, vs)
    glAttachShader(program, fs)
    glLinkProgram(program)
    check = glGetProgramiv(program, GL_LINK_STATUS)
    if not(check):
        raise RuntimeError(glGetProgramInfoLog(program))

    return program

def updateBufferData(circles, vbo, tbo, ibo, xbo, vertices, texCoords, indices, radii, transforms, instances):
    for i, circle in enumerate(circles):
        cx, cy, r = circle
        vertices[i] = [-r, -r, r, -r, -r, r, r, r]
        texCoords[i] = [-1, -1, 1, -1, -1, 1, 1, 1]
        indices[i] = np.array([0, 1, 2, 2, 3, 1], dtype=np.float32)+(4*i)
        transforms[i] = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, cx, cy, 0, 1]
        instances[i] = [i]*4
        radii[i] = r
    glBindBuffer(GL_ARRAY_BUFFER, vbo)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    glBindBuffer(GL_ARRAY_BUFFER, 0)
    glBindBuffer(GL_ARRAY_BUFFER, tbo)
    glBufferData(GL_ARRAY_BUFFER, texCoords.nbytes, texCoords, GL_STATIC_DRAW)
    glBindBuffer(GL_ARRAY_BUFFER, 0)
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices, GL_STATIC_DRAW)
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0)
    glBindBuffer(GL_ARRAY_BUFFER, xbo)
    glBufferData(GL_ARRAY_BUFFER, instances.nbytes, instances, GL_STATIC_DRAW)
    glBindBuffer(GL_ARRAY_BUFFER, 0)

if __name__ == "__main__":
    width = 500
    height = 500
    title = "Pseudo-Instancing Demo"
    pygame.init()
    screen = pygame.display.set_mode((width, height), DOUBLEBUF|OPENGL|HWSURFACE)

    vsStr = """#version 120
    attribute vec2 a_position;
    attribute vec2 a_texCoord;
    attribute int a_instance;
    uniform mat4 u_transform[500];
    uniform vec4 u_color[500];
    uniform float u_radius[500];
    varying vec4 v_color;
    varying float v_radius;
    varying vec2 v_texCoord;
    void main()
    {
        gl_Position = gl_ProjectionMatrix * u_transform[a_instance] * vec4(a_position, 0, 1);
        v_texCoord = a_texCoord;
        v_color = u_color[a_instance];
        v_radius = u_radius[a_instance];
    }
    """

    fsStr = """#version 120
    varying vec4 v_color;
    varying float v_radius;
    varying vec2 v_texCoord;
    void main()
    {
        float d = v_radius - length(v_texCoord*vec2(v_radius, v_radius));
        float t = clamp(d, 0, 1);
        gl_FragColor = vec4(1, 1, 1, t) * v_color;
    }
    """

    numCircles = 500
    circles = createCircles(numCircles)
    shader = createShader(vsStr, fsStr)
    vbo, tbo, ibo, xbo = glGenBuffers(4)
    vertices = np.empty([numCircles, 8], dtype=np.float32)
    texCoords = np.empty([numCircles, 8], dtype=np.float32)
    indices = np.empty([numCircles, 6], dtype=np.int32)
    transforms = np.empty([numCircles, 16], dtype=np.float32)
    colors = np.random.rand(numCircles, 4)
    radii = np.empty(numCircles, dtype=np.float32)
    instances = np.empty([numCircles, 4], dtype=np.int32)
    glClearColor(1.0, 1.0, 1.0, 1.0)
    glMatrixMode(GL_PROJECTION)
    glOrtho(0, width, 0, height, -1, 1)
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

    while True:
        start = time.time()
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.QUIT:
                sys.exit()

        updateBufferData(circles, vbo, tbo, ibo, xbo, vertices, texCoords, indices, radii, transforms, instances)

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glUseProgram(shader)
        a_position = glGetAttribLocation(shader, "a_position")
        a_texCoord = glGetAttribLocation(shader, "a_texCoord")
        a_instance = glGetAttribLocation(shader, "a_instance")

        u_transform = glGetUniformLocation(shader, "u_transform")
        u_color = glGetUniformLocation(shader, "u_color")
        u_radius = glGetUniformLocation(shader, "u_radius")
        glUniformMatrix4fv(u_transform, numCircles, False, transforms)
        glUniform4fv(u_color, numCircles, colors)
        glUniform1fv(u_radius, numCircles, radii)

        glEnableVertexAttribArray(a_position)
        glEnableVertexAttribArray(a_texCoord)
        glEnableVertexAttribArray(a_instance)
        glBindBuffer(GL_ARRAY_BUFFER, vbo)
        glVertexAttribPointer(a_position, 2, GL_FLOAT, False, 0, None)
        glBindBuffer(GL_ARRAY_BUFFER, tbo)
        glVertexAttribPointer(a_texCoord, 2, GL_FLOAT, False, 0, None)
        glBindBuffer(GL_ARRAY_BUFFER, xbo)
        glVertexAttribPointer(a_instance, 1, GL_FLOAT, False, 0, None)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo)
        glDrawElements(GL_TRIANGLES, 6*numCircles, GL_UNSIGNED_INT, None)
        glDisableVertexAttribArray(a_position)
        glDisableVertexAttribArray(a_texCoord)
        glDisableVertexAttribArray(a_instance)
        glUseProgram(0)

        pygame.display.flip()
        end = time.time()
        fps = int(round(1/(end - start)))
        pygame.display.set_caption(title + ": " + str(fps))