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

# ---------- Shaders ----------
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 sampler2D latticeTex;
uniform float cycle;
uniform float omegaTime;
uniform float phiPowers[72];
uniform float threshold;

float phi = 1.6180339887;

float hdgl_slot(float val, float r_dim, float omega, int x, int y){
    float idx0 = float(x % 4 == 0 ? 1 : 0);
    float resonance = 0.1 * sin(float(cycle)*0.05 + float(y)) * idx0;
    float wave = 0.3 * float(x%3==0) -0.3 * float(x%3==2);
    float omega_inst = 1.0 / phiPowers[y % 72];
    float rec = r_dim*val*0.5 + 0.25*sin(float(cycle)*r_dim + float(x));
    float new_val = val + omega_inst + resonance + wave + rec + omega*0.05;
    return new_val > threshold ? 1.0 : 0.0;
}

void main(){
    int x = int(texCoord.x * 32.0);
    int y = int(texCoord.y * 2304.0);  // test lattice height
    float val = texelFetch(latticeTex, ivec2(x,y), 0).r;
    float r_dim = 0.3 + 0.01*float(y);
    float new_val = hdgl_slot(val, r_dim, sin(omegaTime), x, y);
    fragColor = vec4(new_val, sin(omegaTime), 0.0, 1.0);
}
"""

# ---------- Globals ----------
window = None
shader = None
vao = None
fbo = [0,0]
tex = [0,0]
current = 0
cycle = 0.0
omega_time = 0.0
phi_powers = np.array([pow(1.6180339887,7*(i+1)) for i in range(72)],dtype=np.float32)
threshold = np.sqrt(1.6180339887)

# ---------- OpenGL Init ----------
def init_gl():
    global shader, vao, fbo, tex
    shader = compileProgram(compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
                            compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER))
    # Fullscreen quad
    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)

    # Ping-pong FBOs
    fbo[0], fbo[1] = glGenFramebuffers(2)
    tex[0], tex[1] = glGenTextures(2)
    for i in range(2):
        glBindTexture(GL_TEXTURE_2D, tex[i])
        glTexImage2D(GL_TEXTURE_2D,0,GL_RGBA32F,32,2304,0,GL_RGBA,GL_FLOAT,None)
        glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MIN_FILTER,GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MAG_FILTER,GL_NEAREST)
        glBindFramebuffer(GL_FRAMEBUFFER,fbo[i])
        glFramebufferTexture2D(GL_FRAMEBUFFER,GL_COLOR_ATTACHMENT0,GL_TEXTURE_2D,tex[i],0)
    # Seed first lattice
    init_lattice = np.zeros((2304,32,4),dtype=np.float32)
    init_lattice[:,0,0] = 1.0
    glBindTexture(GL_TEXTURE_2D,tex[0])
    glTexSubImage2D(GL_TEXTURE_2D,0,0,0,32,2304,GL_RGBA,GL_FLOAT,init_lattice)
    glBindFramebuffer(GL_FRAMEBUFFER,0)

# ---------- Display ----------
def display():
    global cycle, omega_time, current
    next_idx = 1-current
    glBindFramebuffer(GL_FRAMEBUFFER,fbo[next_idx])
    glViewport(0,0,32,2304)
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader)
    glUniform1i(glGetUniformLocation(shader,"latticeTex"),0)
    glUniform1f(glGetUniformLocation(shader,"cycle"),cycle)
    glUniform1f(glGetUniformLocation(shader,"omegaTime"),omega_time)
    glUniform1fv(glGetUniformLocation(shader,"phiPowers"),72,phi_powers)
    glUniform1f(glGetUniformLocation(shader,"threshold"),threshold)
    glActiveTexture(GL_TEXTURE0)
    glBindTexture(GL_TEXTURE_2D, tex[current])
    glBindVertexArray(vao)
    glDrawArrays(GL_TRIANGLES,0,6)
    glBindFramebuffer(GL_FRAMEBUFFER,0)
    glViewport(0,0,1280,720)
    glBindTexture(GL_TEXTURE_2D, tex[next_idx])
    glDrawArrays(GL_TRIANGLES,0,6)
    glutSwapBuffers()
    cycle +=1
    omega_time +=0.05
    current = next_idx

def idle():
    glutPostRedisplay()

def main():
    global window
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE)
    glutInitWindowSize(1280,720)
    window = glutCreateWindow(b"HDGL Max-Scale GPU Lattice + PHI Clock")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__=="__main__":
    main()
