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         # slots per strand
NUM_STRANDS = 512
PARTICLES = 50000

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 = """
#version 330
out vec4 fragColor;

uniform sampler1D texStrand0; // RGBA = strands 0-3
uniform sampler1D texStrand1; // RGBA = strands 4-7
uniform sampler1D texStrand2; // RGBA = strands 8-11
uniform sampler1D texStrand3; // RGBA = strands 12-15
uniform sampler1D texStrand4; // RGBA = strands 16-19
uniform sampler1D texStrand5; // RGBA = strands 20-23
uniform sampler1D texStrand6; // RGBA = strands 24-27
uniform sampler1D texStrand7; // RGBA = strands 28-31

uniform float omegaTime;
uniform float phi;
uniform int NUM_SLOTS;

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

vec4 fetch_strand_block(sampler1D tex, int id, float r){
    vec4 data = texture(tex, float(id)/float(NUM_SLOTS));
    return vec4(
        prismatic_recursion(data.r,r),
        prismatic_recursion(data.g,r),
        prismatic_recursion(data.b,r),
        prismatic_recursion(data.a,r)
    );
}

void main(){
    float r = length(gl_FragCoord.xy - vec2(512.0,384.0))/512.0;
    vec4 val0=vec4(0.0), val1=vec4(0.0), val2=vec4(0.0), val3=vec4(0.0);
    vec4 val4=vec4(0.0), val5=vec4(0.0), val6=vec4(0.0), val7=vec4(0.0);
    
    for(int s=0; s<NUM_SLOTS; s++){
        val0 += fetch_strand_block(texStrand0, s, r);
        val1 += fetch_strand_block(texStrand1, s, r);
        val2 += fetch_strand_block(texStrand2, s, r);
        val3 += fetch_strand_block(texStrand3, s, r);
        val4 += fetch_strand_block(texStrand4, s, r);
        val5 += fetch_strand_block(texStrand5, s, r);
        val6 += fetch_strand_block(texStrand6, s, r);
        val7 += fetch_strand_block(texStrand7, s, r);
    }
    
    vec4 combined = (val0+val1+val2+val3+val4+val5+val6+val7)/float(NUM_SLOTS*8);
    fragColor = combined; // final RGBA encodes all strands in superposition
}
"""

# ------------------ 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):
    # strand_values: NUM_STRANDS x NUM_SLOTS
    textures=[]
    for i in range(0, NUM_STRANDS,4):
        block = np.zeros((NUM_SLOTS,4),dtype=np.float32)
        for j in range(4):
            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_RGBA32F,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

# ------------------ INIT ------------------
def init_gl():
    global shader, tex_handles
    shader=compileProgram(
        compileShader(VERTEX_SHADER,GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SHADER,GL_FRAGMENT_SHADER)
    )
    # generate NUM_STRANDS arrays
    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))  # simple superposition

    packed=pack_strands_to_rgba(strand_values)
    tex_handles=[create_1d_texture(p) for p in packed]

    glUseProgram(shader)
    for i, tex_name in enumerate([
        "texStrand0","texStrand1","texStrand2","texStrand3",
        "texStrand4","texStrand5","texStrand6","texStrand7"
    ]):
        glUniform1i(glGetUniformLocation(shader, tex_name), i)
    glUniform1f(glGetUniformLocation(shader,"phi"),1.6180339887)
    glUniform1i(glGetUniformLocation(shader,"NUM_SLOTS"),NUM_SLOTS)

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

    def display():
        glClear(GL_COLOR_BUFFER_BIT)
        for i, tex in enumerate(tex_handles):
            glActiveTexture(GL_TEXTURE0+i)
            glBindTexture(GL_TEXTURE_1D, tex)
        glUseProgram(shader)
        glDrawArrays(GL_TRIANGLES,0,3)
        glutSwapBuffers()
        glutPostRedisplay()

    glutDisplayFunc(display)
    glutMainLoop()

if __name__=="__main__":
    main()
