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

# Optional GPU utilities
try:
    import GPUtil
except ImportError:
    GPUtil = None
    print("[GPU] GPUtil not installed. Using safe fallback for VRAM.")

# ---------- 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;
uniform int latticeHeight;
uniform int yOffset;

float hdgl_slot(float val, float r_dim, float omega, int x, int y){
    float resonance = (x % 4 == 0 ? 0.1 * sin(cycle*0.05 + float(y)) : 0.0);
    float wave = (x % 3 == 0 ? 0.3 : (x % 3 == 1 ? 0.0 : -0.3));
    float omega_inst = phiPowers[y % 72];
    float rec = r_dim*val*0.5 + 0.25*sin(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 * float(latticeHeight)) + yOffset;
    float val = texelFetch(latticeTex, ivec2(x,y - yOffset), 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
textures = []
fbos = []
current = 0
cycle = 0.0
omega_time = 0.0

# ---------- Lattice params ----------
lattice_width = 32
num_instances_base = 50_000
scale_interval = 2.0
sub_tile_height = 256
max_tex_height = 2048
threshold = np.sqrt(1.6180339887)
phi = 1.6180339887
phi_powers = np.array([1.0 / pow(phi, 7*(i+1)) for i in range(72)], dtype=np.float32)

tile_heights = []
tile_count = 0
last_scale_time = time.time()
frame_times = []

# ---------- GPU Detection with fallback ----------
def compute_max_instances(lattice_width, bytes_per_pixel=16):
    if GPUtil:
        gpus = GPUtil.getGPUs()
        if gpus:
            gpu = gpus[0]
            free_vram = gpu.memoryFree * 1024**2
            instance_bytes = lattice_width * 4 * 4
            max_instances = int(free_vram * 0.9 / instance_bytes)
            print(f"[GPU] Free VRAM: {gpu.memoryFree} MB, Max safe instances: {max_instances}")
            return max_instances
    print("[GPU] No GPU detected or GPUtil failed. Using safe fallback for max instances.")
    return 50_000

num_instances_max = compute_max_instances(lattice_width)

# ---------- OpenGL Init ----------
def init_gl():
    global shader, vao, textures, fbos, tile_heights, tile_count
    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)

    reinit_lattice(num_instances_base)

# ---------- Lattice reinit ----------
def reinit_lattice(new_num_instances):
    global textures, fbos, tile_heights, tile_count
    for tex_pair in textures:
        glDeleteTextures(tex_pair)
    for fbo_pair in fbos:
        glDeleteFramebuffers(2, fbo_pair)
    textures.clear()
    fbos.clear()

    tile_count = (new_num_instances + max_tex_height - 1) // max_tex_height
    tile_heights[:] = [min(max_tex_height, new_num_instances - i*max_tex_height) for i in range(tile_count)]

    for t, th in enumerate(tile_heights):
        tex_pair = glGenTextures(2)
        fbo_pair = glGenFramebuffers(2)
        for i in range(2):
            glBindTexture(GL_TEXTURE_2D, tex_pair[i])
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, lattice_width, th, 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_pair[i])
            glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, tex_pair[i],0)
        textures.append(tex_pair)
        fbos.append(fbo_pair)

        for start in range(0, th, sub_tile_height):
            end = min(start + sub_tile_height, th)
            init_chunk = np.zeros((end-start, lattice_width,4), dtype=np.float32)
            glBindTexture(GL_TEXTURE_2D, tex_pair[0])
            glTexSubImage2D(GL_TEXTURE_2D,0,0,start,lattice_width,end-start,GL_RGBA,GL_FLOAT,init_chunk)
    glBindFramebuffer(GL_FRAMEBUFFER,0)

# ---------- Display ----------
def display():
    global cycle, omega_time, current
    next_idx = 1-current
    for t, th in enumerate(tile_heights):
        for y_start in range(0, th, sub_tile_height):
            h = min(sub_tile_height, th - y_start)
            glBindFramebuffer(GL_FRAMEBUFFER, fbos[t][next_idx])
            glViewport(0,0,lattice_width,h)
            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)
            glUniform1i(glGetUniformLocation(shader,"latticeHeight"), h)
            glUniform1i(glGetUniformLocation(shader,"yOffset"), y_start)
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_2D, textures[t][current])
            glBindVertexArray(vao)
            glDrawArrays(GL_TRIANGLES,0,6)

    glBindFramebuffer(GL_FRAMEBUFFER,0)
    glViewport(0,0,1280,720)
    for t, th in enumerate(tile_heights):
        glBindTexture(GL_TEXTURE_2D, textures[t][next_idx])
        glDrawArrays(GL_TRIANGLES,0,6)

    glutSwapBuffers()
    start = time.time()
    cycle += 1
    omega_time += 0.05
    current = next_idx
    frame_times.append(time.time()-start)
    if len(frame_times) > 100: frame_times.pop(0)

# ---------- Idle ----------
def idle():
    global num_instances_base, last_scale_time
    glutPostRedisplay()
    now = time.time()

    if now - last_scale_time > scale_interval and num_instances_base < num_instances_max:
        gpu_headroom = 1.0
        if GPUtil:
            gpus = GPUtil.getGPUs()
            if gpus:
                gpu_headroom = max(0.1, 1.0 - gpus[0].load)
        step = int(max(1_000, (1.0 - num_instances_base/num_instances_max) * 50_000 * gpu_headroom))
        num_instances_base += step
        if num_instances_base > num_instances_max:
            num_instances_base = num_instances_max
        last_scale_time = now
        print(f"[Scaling] Instances: {num_instances_base}, Step: {step}, GPU headroom: {gpu_headroom:.2f}")
        reinit_lattice(num_instances_base)

    if frame_times:
        avg = sum(frame_times)/len(frame_times)
        print(f"[Perf] Avg frame: {avg*1000:.2f} ms")

# ---------- Main ----------
def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE)
    glutInitWindowSize(1280,720)
    glutCreateWindow(b"HDGL Full GPU Saturation + PHI Clock")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__=="__main__":
    main()
