import numpy as np
from OpenGL.GL import *
from OpenGL.GL.shaders import compileProgram, compileShader
import glfw

# ---------------- Globals ----------------
NUM_SLOTS = 4096  # full prismatic lattice
NUM_INSTANCES = 4096
MAX_SLICE = 8

# ---------------- Fibonacci Table ----------------
def generate_fib_table_gpu(n):
    # We will store normalized Fibonacci numbers in float32
    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]
    # final normalization
    fibs /= fibs.max() if fibs.max() != 0 else 1.0
    return fibs

# ---------------- Shader Source (inline) ----------------
VERTEX_SRC = """
#version 430
layout(location=0) in vec3 position;
void main() {
    gl_Position = vec4(position, 1.0);
    gl_PointSize = 1.0;
}
"""

FRAGMENT_SRC = f"""
#version 430
out vec4 fragColor;

uniform float cycle;
uniform float omegaTime;
uniform float phi;
uniform int instanceID;
uniform int MAX_SLICE;

layout(std430, binding=0) buffer FibBuffer {{
    float fibTable[];
}};

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

void main() {{
    float r = length(gl_FragCoord.xy / vec2(800.0,600.0) - 0.5) * 2.0;
    float val = 0.0;

    for(int s=0; s<{NUM_SLOTS}; s++) {{
        int idx = (instanceID * {NUM_SLOTS} + s) % {NUM_INSTANCES};
        val += prismatic_recursion(idx, r);
    }}
    val /= float({NUM_SLOTS});

    float phase = sin(cycle*0.01 + val);
    float slice = mod(float(instanceID), float(MAX_SLICE));
    fragColor = vec4(val, phase, r, slice/float(MAX_SLICE));
}}
"""

# ---------------- OpenGL / GLFW Setup ----------------
def init_glfw():
    if not glfw.init():
        raise Exception("GLFW init failed")
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 4)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    window = glfw.create_window(800, 600, "HDGL Superglyphs RX480", None, None)
    glfw.make_context_current(window)
    return window

def init_shaders():
    shader = compileProgram(
        compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )
    return shader

def create_fib_ssbo(fib_table):
    ssbo = glGenBuffers(1)
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo)
    glBufferData(GL_SHADER_STORAGE_BUFFER, fib_table.nbytes, fib_table, GL_STATIC_DRAW)
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssbo)
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0)
    return ssbo

def main():
    window = init_glfw()
    shader = init_shaders()

    # Generate GPU-side Fibonacci table
    fib_table = generate_fib_table_gpu(NUM_SLOTS)
    ssbo = create_fib_ssbo(fib_table)

    # Dummy full-screen quad (for fragment shader execution)
    quad = np.array([
        -1,-1,0,
         1,-1,0,
        -1, 1,0,
        -1, 1,0,
         1,-1,0,
         1, 1,0
    ], dtype=np.float32)

    vao = glGenVertexArrays(1)
    glBindVertexArray(vao)
    vbo = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, vbo)
    glBufferData(GL_ARRAY_BUFFER, quad.nbytes, quad, GL_STATIC_DRAW)
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,0,None)

    glUseProgram(shader)
    phi_loc = glGetUniformLocation(shader,"phi")
    cycle_loc = glGetUniformLocation(shader,"cycle")
    omega_loc = glGetUniformLocation(shader,"omegaTime")
    instance_loc = glGetUniformLocation(shader,"instanceID")
    max_slice_loc = glGetUniformLocation(shader,"MAX_SLICE")

    t = 0.0
    while not glfw.window_should_close(window):
        glfw.poll_events()
        glClearColor(0.0,0.0,0.0,1.0)
        glClear(GL_COLOR_BUFFER_BIT)

        glUniform1f(phi_loc, 1.6180339887)
        glUniform1f(cycle_loc, t)
        glUniform1f(omega_loc, t*0.3)
        glUniform1i(instance_loc, 0)
        glUniform1i(max_slice_loc, MAX_SLICE)

        glDrawArrays(GL_TRIANGLES, 0, 6)

        glfw.swap_buffers(window)
        t += 0.02

    glfw.terminate()

if __name__=="__main__":
    main()
