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

# ---------------- Global Parameters ----------------
NUM_SUPER = 32768         # Can push up to 32768+
NUM_INSTANCES = 32768     # Can push higher if VRAM allows
MAX_SLICE = 8
VIRT_WIDTH = 3840        # Virtual canvas oversampling (e.g., 4K)
VIRT_HEIGHT = 2160
FPS_TARGET = 60

fibTable = np.array([float(f) for f in range(1,129)], 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,131,137,139,149,151,157,163,167,173,179,181,191,
                       193,197,199,211,223,227,229,233,239,241,251,257,263,
                       269,271,277,281,283,293,307,311,313,317,331,337,347,349,
                       353,359,367,373,379,383,389,397,401,409,419,421,431,433,
                       439,443,449,457,461,463,467,479,487,491,499,503,509,521,
                       523,541,547,557,563,569,571,577,587,593,599,601,607,613,
                       617,619,631,641,643,647,653,659,661,673,677,683,691,701,
                       709,719,727,733,739,743,751,757,761,769,773,787,797,809], dtype=np.float32)

cycle = 0.0
omegaTime = 0.0
phi = 1.61803398875
phiInv = 0.61803398875

shader = None
window = None

# ---------------- Shader ----------------
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 int NUM_SUPER;
uniform int NUM_INSTANCES;
uniform int MAX_SLICE;
uniform float fibTable[128];
uniform float primeTable[128];

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;

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

    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 Initialization ----------------
def init_gl():
    global shader
    shader = compileProgram(
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )
    glUseProgram(shader)

    # Upload uniforms
    glUniform1i(glGetUniformLocation(shader, "NUM_SUPER"), NUM_SUPER)
    glUniform1i(glGetUniformLocation(shader, "NUM_INSTANCES"), NUM_INSTANCES)
    glUniform1i(glGetUniformLocation(shader, "MAX_SLICE"), MAX_SLICE)
    glUniform1fv(glGetUniformLocation(shader, "fibTable"), 128, fibTable)
    glUniform1fv(glGetUniformLocation(shader, "primeTable"), 128, primeTable)
    glUniform1f(glGetUniformLocation(shader, "phi"), phi)
    glUniform1f(glGetUniformLocation(shader, "phiInv"), phiInv)

# ---------------- Render Loop ----------------
def display():
    global cycle, omegaTime
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader)

    glUniform1f(glGetUniformLocation(shader, "cycle"), cycle)
    glUniform1f(glGetUniformLocation(shader, "omegaTime"), omegaTime)
    for instanceID in range(NUM_INSTANCES):
        glUniform1i(glGetUniformLocation(shader, "instanceID"), instanceID)
        glBegin(GL_POINTS)
        glVertex2f(0.0, 0.0)
        glEnd()

    glutSwapBuffers()
    cycle += 1.0
    omegaTime += 0.01
    glutPostRedisplay()

def main():
    global window
    glutInit()
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB)
    glutInitWindowSize(1080, 1080)
    glutInitWindowPosition(50, 50)
    window = glutCreateWindow(b"HDGL Prismatic BaseI Fabric RX480 TFLOPs")
    init_gl()
    glutDisplayFunc(display)
    glutMainLoop()

if __name__ == "__main__":
    main()
