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

# ---------------- CONFIG ----------------
NUM_SUPER = 256          # adjustable superpositions
NUM_INSTANCES = 512      # adjustable instance count
MAX_SLICE = 8            # RGBA virtual channels
VIRT_WIDTH = 3840        # oversampled virtual canvas width
VIRT_HEIGHT = 2160       # oversampled virtual canvas height
FPS = 60.0

phi = 1.61803398875
phiInv = 1.0 / phi

fibTable = np.array([1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144] * 12, dtype=np.float32)
primeTable = np.array([
    2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,
    59,61,67,71,73,79,83,89,97,101,103,107,109,113,127
]*4, dtype=np.float32)

# ---------- Shader source ----------
FRAGMENT_SRC = """
#version 330
in vec2 texCoord;
out vec4 fragColor;

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

uniform float fibTable[128];
uniform float primeTable[128];
uniform int MAX_SLICE;

// Superposition helper
float prismatic_recursion(int id, float r){
    float phi_harm = pow(phi, float(mod(id,16)));
    float fib_harm = fibTable[id % 128];
    float dyadic = float(1 << int(mod(float(id),16.0)));
    float prime_harm = primeTable[id % 128];
    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 * prime_harm * Omega) * r_dim;
}

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

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

    float phase = sin(cycle*0.01 + val);

    // RGBA packing with MAX_SLICE channels
    float slice = mod(float(instanceID), float(MAX_SLICE));
    fragColor = vec4(val, phase, r, slice/float(MAX_SLICE));
}
"""

# ---------- Globals ----------
shader = None
window = None
cycle = 0.0
start_time = time.time()

# ---------- OpenGL Initialization ----------
def init_gl():
    global shader
    shader = compileProgram(
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER),
        compileShader("""
        #version 330
        layout(location=0) in vec2 position;
        out vec2 texCoord;
        void main() {
            texCoord = position * 0.5 + 0.5;
            gl_Position = vec4(position, 0.0, 1.0);
        }
        """, GL_VERTEX_SHADER)
    )

    # Setup VAO/VBO for fullscreen quad
    quad = np.array([
        [-1, -1],
        [ 1, -1],
        [ 1,  1],
        [-1,  1]
    ], dtype=np.float32)

    global vao
    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, 2, GL_FLOAT, GL_FALSE, 0, None)

# ---------- Display Loop ----------
def display():
    global cycle
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader)

    t = time.time() - start_time
    glUniform1f(glGetUniformLocation(shader, "cycle"), cycle)
    glUniform1f(glGetUniformLocation(shader, "omegaTime"), t)
    glUniform1f(glGetUniformLocation(shader, "phi"), phi)
    glUniform1f(glGetUniformLocation(shader, "phiInv"), phiInv)

    glUniform1fv(glGetUniformLocation(shader, "fibTable"), 128, fibTable)
    glUniform1fv(glGetUniformLocation(shader, "primeTable"), 128, primeTable)
    glUniform1i(glGetUniformLocation(shader, "MAX_SLICE"), MAX_SLICE)

    for i in range(NUM_INSTANCES):
        glUniform1i(glGetUniformLocation(shader, "instanceID"), i)
        glBindVertexArray(vao)
        glDrawArrays(GL_TRIANGLE_FAN, 0, 4)

    glutSwapBuffers()
    cycle += 1.0

# ---------- Main ----------
def main():
    global window
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGBA)
    glutInitWindowSize(1080, 1080)
    glutInitWindowPosition(100,100)
    glutCreateWindow("HDGL Prismatic 🌈 Fabric".encode("utf-8"))

    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(display)
    glutMainLoop()

if __name__ == "__main__":
    main()
