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

# ------------------ CONFIG ------------------
NUM_SLOTS = 4096
NUM_STRANDS = 32
BLOCK_SIZE = 4          # 4 strands per RGBA
WINDOW_WIDTH = 1024
WINDOW_HEIGHT = 768

# ------------------ GLSL ------------------
VERTEX_SHADER = """
#version 330
layout(location=0) in vec3 position;
void main() {
    gl_Position = vec4(position,1.0);
}
"""

FRAGMENT_SHADER_BLOCK = """
#version 330
out vec4 fragColor;
uniform sampler1D strandTex;
uniform float omegaTime;
uniform int NUM_SLOTS;
uniform int blockIndex; // 0..7

float prismatic_recursion(float val, float r){
    return sqrt(val) * (0.5 + 0.5*sin(omegaTime + r*10.0));
}

void main(){
    float r = length(gl_FragCoord.xy - vec2(512.0,384.0))/512.0;
    vec4 accum = vec4(0.0);
    for(int s=0; s<NUM_SLOTS; s++){
        vec4 data = texture(strandTex, float(s)/float(NUM_SLOTS));
        accum += vec4(
            prismatic_recursion(data.r,r),
            prismatic_recursion(data.g,r),
            prismatic_recursion(data.b,r),
            prismatic_recursion(data.a,r)
        );
    }
    fragColor = accum / float(NUM_SLOTS);
}
"""

FRAGMENT_SHADER_FINAL = """
#version 330
out vec4 fragColor;
uniform sampler2D blockTextures[8]; // 8 RGBA textures
void main(){
    vec2 uv = gl_FragCoord.xy / vec2(1024.0,768.0);
    vec4 combined = vec4(0.0);
    for(int i=0;i<8;i++){
        combined += texture(blockTextures[i], uv);
    }
    fragColor = combined / 8.0;
}
"""

# ------------------ TEXTURE HELPERS ------------------
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: fibs/=fibs[i]
    fibs /= fibs.max() if fibs.max()!=0 else 1.0
    return fibs

def generate_prime_table(n):
    primes = np.zeros(n, dtype=np.float32)
    primes[0] = 2.0
    count = 1
    candidate = 3
    while count<n:
        is_prime=True
        for p in primes[:count]:
            if candidate % p==0:
                is_prime=False
                break
        if is_prime:
            primes[count]=float(candidate)
            count+=1
        candidate+=2
    primes/=primes.max() if primes.max()!=0 else 1.0
    return primes

def pack_strands_to_rgba(strand_values):
    textures=[]
    for i in range(0, NUM_STRANDS, BLOCK_SIZE):
        block = np.zeros((NUM_SLOTS,4), dtype=np.float32)
        for j in range(BLOCK_SIZE):
            if i+j < NUM_STRANDS:
                block[:,j] = strand_values[i+j]
        textures.append(block)
    return textures

def create_1d_texture(data):
    tex = glGenTextures(1)
    glBindTexture(GL_TEXTURE_1D, tex)
    glTexImage1D(GL_TEXTURE_1D, 0, GL_RGBA16F, len(data), 0, GL_RGBA, GL_FLOAT, data)
    glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    return tex

def create_2d_framebuffer(tex):
    fbo = glGenFramebuffers(1)
    glBindFramebuffer(GL_FRAMEBUFFER, fbo)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, tex, 0)
    status = glCheckFramebufferStatus(GL_FRAMEBUFFER)
    if status != GL_FRAMEBUFFER_COMPLETE:
        print("Framebuffer incomplete:", status)
    glBindFramebuffer(GL_FRAMEBUFFER,0)
    return fbo

# ------------------ INIT ------------------
def init_gl():
    global shader_block, shader_final, tex_handles, fbo_handles, fb_tex
    shader_block = compileProgram(
        compileShader(VERTEX_SHADER, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SHADER_BLOCK, GL_FRAGMENT_SHADER)
    )
    shader_final = compileProgram(
        compileShader(VERTEX_SHADER, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SHADER_FINAL, GL_FRAGMENT_SHADER)
    )

    # Generate strands
    strand_values=[]
    fibs=generate_fib_table(NUM_SLOTS)
    primes=generate_prime_table(NUM_SLOTS)
    for s in range(NUM_STRANDS):
        strand_values.append((fibs+primes)*(s+1))

    # Pack strands into RGBA blocks
    packed = pack_strands_to_rgba(strand_values)
    tex_handles = [create_1d_texture(p) for p in packed]

    # Offscreen framebuffer for multi-pass
    fb_tex = glGenTextures(1)
    glBindTexture(GL_TEXTURE_2D, fb_tex)
    glTexImage2D(GL_TEXTURE_2D,0,GL_RGBA16F,WINDOW_WIDTH,WINDOW_HEIGHT,0,GL_RGBA,GL_FLOAT,None)
    glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MIN_FILTER,GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MAG_FILTER,GL_LINEAR)
    fbo_handles = create_2d_framebuffer(fb_tex)

# ------------------ MAIN ------------------
def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGBA)
    glutInitWindowSize(WINDOW_WIDTH,WINDOW_HEIGHT)
    glutCreateWindow(b"32-Strand Full Superglyphs Engine")
    init_gl()

    omegaTime = 0.0

    def display():
        nonlocal omegaTime
        omegaTime += 0.01

        # --- Render each block to FBO ---
        glBindFramebuffer(GL_FRAMEBUFFER, fbo_handles)
        glViewport(0,0,WINDOW_WIDTH,WINDOW_HEIGHT)
        glClear(GL_COLOR_BUFFER_BIT)
        for i, tex in enumerate(tex_handles):
            glUseProgram(shader_block)
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_1D, tex)
            glUniform1i(glGetUniformLocation(shader_block,"NUM_SLOTS"),NUM_SLOTS)
            glUniform1f(glGetUniformLocation(shader_block,"omegaTime"), omegaTime)
            glUniform1i(glGetUniformLocation(shader_block,"blockIndex"), i)
            glDrawArrays(GL_TRIANGLES,0,3)

        glBindFramebuffer(GL_FRAMEBUFFER,0)

        # --- Final composite pass ---
        glClear(GL_COLOR_BUFFER_BIT)
        glUseProgram(shader_final)
        for i in range(8):
            glActiveTexture(GL_TEXTURE0+i)
            glBindTexture(GL_TEXTURE_2D, fb_tex)
            glUniform1i(glGetUniformLocation(shader_final,f"blockTextures[{i}]"), i)
        glDrawArrays(GL_TRIANGLES,0,3)

        glutSwapBuffers()
        glutPostRedisplay()

    glutDisplayFunc(display)
    glutMainLoop()

if __name__=="__main__":
    main()
