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

# ----------------- Globals -----------------
window = None
shader = None
vao = None

width, height = 1280, 720
cycle = 0.0
start_time = time.time()

NUM_SLOTS = 2048  # Saturated but stable on RX 480
fib_table = np.zeros(NUM_SLOTS, dtype=np.float32)
prime_table = np.zeros(NUM_SLOTS, dtype=np.float32)

# ----------------- Math Tables -----------------
def generate_fib_table(n):
    fibs = np.zeros(n, dtype=np.float32)
    fibs[0], fibs[1] = 0.0, 1.0
    for i in range(2, n):
        fibs[i] = (fibs[i-1] + fibs[i-2])
        if fibs[i] > 1e6:  # rescale dynamically
            fibs /= fibs[i]
    fibs /= fibs.max() if fibs.max() != 0 else 1.0
    return fibs

def generate_prime_table(n):
    primes = []
    candidate = 2
    while len(primes) < n:
        if all(candidate % p for p in primes):
            primes.append(candidate)
        candidate += 1
    arr = np.array(primes[:n], dtype=np.float32)
    arr /= arr.max()
    return arr

# ----------------- GLSL Shaders -----------------
VERTEX_SRC = """
#version 330
layout(location = 0) in vec2 position;
out vec2 texCoord;
void main(){
    texCoord = (position + 1.0) * 0.5;
    gl_Position = vec4(position, 0.0, 1.0);
}
"""

FRAGMENT_SRC = f"""
#version 330

in vec2 texCoord;
out vec4 fragColor;

uniform float cycle;
uniform float omegaTime;
uniform float phi;
uniform float phiInv;

uniform float fibTable[{NUM_SLOTS}];
uniform float primeTable[{NUM_SLOTS}];

float prismatic_recursion(int id, float r){{
    float phi_harm   = pow(phi, float(id % 16));
    float fib_harm   = fibTable[id % {NUM_SLOTS}];
    float dyadic     = float(1 << (id % 16));
    float prime_harm = primeTable[id % {NUM_SLOTS}];
    float Omega      = 0.5 + 0.5*sin(omegaTime + float(id)*0.01);
    float r_dim      = pow(r, float((id % 7)+1));
    return sqrt(phi_harm * fib_harm * dyadic * prime_harm * Omega) * r_dim;
}}

void main(){{
    float r = length(texCoord - 0.5) * 2.0;
    float val = 0.0;

    for(int s = 0; s < {NUM_SLOTS}; s++) {{
        val += prismatic_recursion(s, r);
    }}
    val /= float({NUM_SLOTS});

    float phase = sin(cycle*0.01 + val);
    fragColor = vec4(val, phase, r, 1.0);
}}
"""

# ----------------- OpenGL Setup -----------------
def init_shaders():
    global shader
    shader = compileProgram(
        compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )

def init_buffers():
    global vao
    vertices = np.array([
        -1.0, -1.0,
         1.0, -1.0,
        -1.0,  1.0,
         1.0,  1.0,
    ], dtype=np.float32)

    vao = glGenVertexArrays(1)
    vbo = glGenBuffers(1)
    glBindVertexArray(vao)

    glBindBuffer(GL_ARRAY_BUFFER, vbo)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)

    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, None)

def display():
    global cycle
    glClear(GL_COLOR_BUFFER_BIT)

    glUseProgram(shader)
    cycle = time.time() - start_time

    glUniform1f(glGetUniformLocation(shader, "cycle"), cycle)
    glUniform1f(glGetUniformLocation(shader, "omegaTime"), cycle * 0.25)
    glUniform1f(glGetUniformLocation(shader, "phi"), (1.0 + 5**0.5) / 2.0)
    glUniform1f(glGetUniformLocation(shader, "phiInv"), 2.0 / (1.0 + 5**0.5))

    glUniform1fv(glGetUniformLocation(shader, "fibTable"), NUM_SLOTS, fib_table)
    glUniform1fv(glGetUniformLocation(shader, "primeTable"), NUM_SLOTS, prime_table)

    glBindVertexArray(vao)
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)

    glutSwapBuffers()
    glutPostRedisplay()

def main():
    global fib_table, prime_table

    fib_table = generate_fib_table(NUM_SLOTS)
    prime_table = generate_prime_table(NUM_SLOTS)

    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB)
    glutInitWindowSize(width, height)
    glutCreateWindow(b"HDGL SuperGlyphs 3 - RX480 Saturation")

    init_shaders()
    init_buffers()

    glutDisplayFunc(display)
    glutMainLoop()

if __name__ == "__main__":
    main()
