from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GL.shaders import compileProgram, compileShader
import numpy as np, time, math, sys

WIDTH, HEIGHT = 8192, 8192
window = None
shader_main, shader_overlay = None, None
fbo, tex = [], []
NUM_STRANDS = 8

def make_shader(vsrc, fsrc):
    return compileProgram(
        compileShader(vsrc, GL_VERTEX_SHADER),
        compileShader(fsrc, GL_FRAGMENT_SHADER)
    )

VERTEX = """
#version 330
layout(location=0) in vec2 pos;
out vec2 uv;
void main(){
    uv = pos*0.5+0.5;
    gl_Position = vec4(pos,0,1);
}
"""

def init_shaders():
    global shader_main, shader_overlay
    with open("prismatic_frag_main.glsl") as f: frag_main = f.read()
    with open("prismatic_frag_overlay.glsl") as f: frag_overlay = f.read()
    shader_main = make_shader(VERTEX, frag_main)
    shader_overlay = make_shader(VERTEX, frag_overlay)

def init_fbos():
    global fbo, tex
    for i in range(NUM_STRANDS):
        f = glGenFramebuffers(1)
        t = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, t)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, WIDTH, HEIGHT, 0,
                     GL_RGBA, GL_FLOAT, None)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glBindFramebuffer(GL_FRAMEBUFFER, f)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
                               GL_TEXTURE_2D, t, 0)
        fbo.append(f); tex.append(t)
    glBindFramebuffer(GL_FRAMEBUFFER, 0)

def upload_tables(shader):
    phi = (1+math.sqrt(5))/2
    fib = [int(round(((phi**i - (-phi)**(-i))/math.sqrt(5)))) % 997 for i in range(128)]
    primes = [p for p in range(2,800) if all(p%d for d in range(2,int(math.sqrt(p))+1))][:128]
    glUseProgram(shader)
    glUniform1f(glGetUniformLocation(shader,"phi"), phi)
    glUniform1f(glGetUniformLocation(shader,"phiInv"), 1.0/phi)
    glUniform1fv(glGetUniformLocation(shader,"fibTable"), 128, np.array(fib,dtype=np.float32))
    glUniform1fv(glGetUniformLocation(shader,"primeTable"), 128, np.array(primes,dtype=np.float32))

def render_pass(shader, target=None, tex_inputs=None):
    if target:
        glBindFramebuffer(GL_FRAMEBUFFER, target)
    else:
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
    glUseProgram(shader)
    if tex_inputs:
        for i,t in enumerate(tex_inputs):
            glActiveTexture(GL_TEXTURE0+i)
            glBindTexture(GL_TEXTURE_2D, t)
            glUniform1i(glGetUniformLocation(shader,f"tex{i}"), i)
    glDrawArrays(GL_TRIANGLES, 0, 6)
    glBindFramebuffer(GL_FRAMEBUFFER, 0)

def display():
    t = time.time()
    # Phase update
    glUseProgram(shader_main)
    glUniform1f(glGetUniformLocation(shader_main,"omegaTime"), t)
    # Compute all strands
    for i in range(NUM_STRANDS):
        render_pass(shader_main, target=fbo[i])
    # Overlay stage
    glUseProgram(shader_overlay)
    glUniform1f(glGetUniformLocation(shader_overlay,"omegaTime"), t)
    render_pass(shader_overlay, target=None, tex_inputs=tex)
    glutSwapBuffers()

def idle():
    glutPostRedisplay()

def main():
    global window
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE)
    glutInitWindowSize(WIDTH, HEIGHT)
    window = glutCreateWindow(b"8 strands of Superglyphs")
    init_shaders()
    init_fbos()
    upload_tables(shader_main)
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__ == "__main__":
    main()
