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

# ---------- Shader Sources ----------
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_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];

float compute_harmonics(int id, float r){
    float phi_harm = pow(phi, float(id % 16));
    float fib_harm = fibTable[id % 128];
    float dyadic = float(1 << (id % 16));
    float prime_harm = primeTable[id % 128];
    float Omega = 0.5 + 0.5*sin(omegaTime + float(id)*0.01);
    float r_dim = pow(r, float((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;

    int half_super = NUM_SUPER / 2;
    for(int s=0; s<half_super; s++){
        int idx = (s + int(cycle)) % (NUM_INSTANCES/2);
        val += compute_harmonics(idx,r);
    }
    for(int s=half_super; s<NUM_SUPER; s++){
        int idx = (s + int(cycle) + NUM_INSTANCES/2) % 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);
}
"""

# ---------- Globals ----------
window = None
shader = None
vao = None
cycle = 0.0
omega_time = 0.0

# Lookup tables
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)

NUM_SUPER = 128
NUM_INSTANCES = 128
VIRT_WIDTH, VIRT_HEIGHT = 8192, 8192

# ---------- Auto-scaling ----------
def auto_scale_gpu():
    global NUM_SUPER, NUM_INSTANCES
    # Use PyOpenCL to query GPU max registers / compute units
    try:
        platforms = cl.get_platforms()
        gpu = None
        for p in platforms:
            devices = p.get_devices(device_type=cl.device_type.GPU)
            if devices:
                gpu = devices[0]
                break
        if gpu:
            max_workgroup = gpu.max_work_group_size
            compute_units = gpu.max_compute_units
            # Heuristic: scale NUM_SUPER × NUM_INSTANCES to saturate ALUs
            NUM_SUPER = min(512, max_workgroup)
            NUM_INSTANCES = min(512, compute_units*NUM_SUPER)
    except:
        # fallback
        NUM_SUPER = 256
        NUM_INSTANCES = 256
    print(f"[AutoScale] NUM_SUPER={NUM_SUPER}, NUM_INSTANCES={NUM_INSTANCES}")

# ---------- OpenGL Init ----------
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)
    glUniform1i(glGetUniformLocation(shader,"NUM_SUPER"), NUM_SUPER)
    glUniform1i(glGetUniformLocation(shader,"NUM_INSTANCES"), NUM_INSTANCES)
    glUniform1fv(glGetUniformLocation(shader,"fibTable"), 128, fib_table)
    glUniform1fv(glGetUniformLocation(shader,"primeTable"), 128, prime_table)

# ---------- Display ----------
def display():
    global cycle, omega_time
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader)
    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():
    auto_scale_gpu()
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE)
    glutInitWindowSize(1920,1080)
    glutCreateWindow(b"HDGL Prismatic Waterfall (RX480 AutoScale)")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__=="__main__":
    main()
