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

# ---------- Display & Virtual Canvas ----------
DISPLAY_W  = 2048  # 2K display
DISPLAY_H  = 1080
VIRT_W     = 16384 # 16K virtual canvas
VIRT_H     = 16384

NUM_SUPER     = 256
NUM_INSTANCES = 256

# ---------- Vertex Shader ----------
VERTEX_SRC = """
#version 330
layout(location = 0) in vec2 pos;
out vec2 texCoord;
void main(){
    texCoord = (pos + 1.0) * 0.5;
    gl_Position = vec4(pos,0,1);
}
"""

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

uniform float cycle;
uniform float omegaTime;
uniform float phi;
uniform int NUM_SUPER;
uniform int NUM_INSTANCES;
uniform float fibTable[128];
uniform float primeTable[128];
uniform vec2 virtOffset;

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

void main(){
    // Stage 5: rolling virtual canvas
    vec2 uv = texCoord + virtOffset / vec2(%d.0, %d.0);
    uv = clamp(uv,0.0,1.0);
    float r = length(uv - 0.5) * 2.0;

    float val = 0.0;
    for(int s=0; s<NUM_SUPER; s++){
        int idx = (s + int(cycle)) % NUM_INSTANCES;
        val += compute_harmonics(idx, r);
    }
    val /= float(NUM_SUPER);

    float phase = sin(cycle*0.01 + val);
    fragColor = vec4(val, phase, r, 1.0);
}
""" % (VIRT_W, VIRT_H)

# ---------- Globals ----------
window = None
shader = None
vao = None
cycle = 0.0
omega_time = 0.0
virt_offset = np.array([0.0,0.0],dtype=np.float32)

fib_table = np.array([((1.6180339887**n - (1-1.6180339887)**n)/np.sqrt(5)) for n in range(128)],dtype=np.float32)
prime_table = 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
],dtype=np.float32)

# ---------- OpenGL Initialization ----------
def init_gl():
    global shader, vao
    shader = compileProgram(
        compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )
    verts = np.array([-1,-1,1,-1,-1,1,1,-1,1,1,-1,1],dtype=np.float32)
    vao = glGenVertexArrays(1)
    glBindVertexArray(vao)
    vbo = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER,vbo)
    glBufferData(GL_ARRAY_BUFFER,verts.nbytes,verts,GL_STATIC_DRAW)
    glVertexAttribPointer(0,2,GL_FLOAT,GL_FALSE,0,None)
    glEnableVertexAttribArray(0)

    glUseProgram(shader)
    glUniform1f(glGetUniformLocation(shader,"phi"),1.6180339887)
    glUniform1fv(glGetUniformLocation(shader,"fibTable"),128,fib_table)
    glUniform1fv(glGetUniformLocation(shader,"primeTable"),128,prime_table)
    glUniform1i(glGetUniformLocation(shader,"NUM_SUPER"),NUM_SUPER)
    glUniform1i(glGetUniformLocation(shader,"NUM_INSTANCES"),NUM_INSTANCES)

# ---------- Display ----------
def display():
    global cycle, omega_time, virt_offset
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader)

    virt_offset[0] = (cycle*8.0) % VIRT_W
    virt_offset[1] = (cycle*4.0) % VIRT_H
    glUniform2f(glGetUniformLocation(shader,"virtOffset"),*virt_offset)
    glUniform1f(glGetUniformLocation(shader,"cycle"),cycle)
    glUniform1f(glGetUniformLocation(shader,"omegaTime"),omega_time)

    glBindVertexArray(vao)
    glDrawArrays(GL_TRIANGLES,0,6)
    glutSwapBuffers()

    cycle += 1.0
    omega_time += 0.05

# ---------- Idle ----------
def idle():
    glutPostRedisplay()

# ---------- Main ----------
def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE)
    glutInitWindowSize(DISPLAY_W,DISPLAY_H)
    glutCreateWindow(b"HDGL Ultimate Prismatic Waterfall RX480")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__=="__main__":
    main()
