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

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

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

FRAGMENT_SRC = """
#version 430
in vec2 texCoord;
out vec4 fragColor;

uniform sampler2D latticeTex;
uniform float cycle;
uniform float omegaTime;
uniform float threshold;

// φ powers + Fibonacci + binary + prime factors pre-packed
layout(std140, binding = 0) uniform PrismaticBlock {
    float phiPowers[72];
    float fibPowers[72];
    float dyadic[72];
    float primes[72];
};

uniform int latticeHeight;
uniform int yOffset;

// ------------------
// Prismatic Recursion Core
// ------------------
float prismatic_slot(float val, int x, int y) {
    float phi   = phiPowers[y % 72];
    float fib   = fibPowers[y % 72];
    float dy    = dyadic[y % 72];
    float prime = primes[y % 72];

    // Resonance from lattice position
    float resonance = 0.05 * sin(cycle * 0.05 + float(x + y));

    // Analog hybrid wave overlay
    float wave = (x % 3 == 0 ? 0.3 :
                 (x % 3 == 1 ? 0.0 : -0.3));

    // Field tension (Ω analog)
    float omega = sin(omegaTime + float(y) * 0.01);

    // Radial exponent scaling
    float r_dim = 0.3 + 0.01 * float(y);

    // Core analog hybrid operator
    float superposed = sqrt(phi * fib * dy * prime * abs(omega)) * r_dim;

    float rec = val * 0.5 + 0.25 * sin(cycle * r_dim + float(x));
    float new_val = val + superposed + resonance + wave + rec;

    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 new_val = prismatic_slot(val, x, y);

    fragColor = vec4(new_val, sin(omegaTime), cos(cycle*0.1), 1.0);
}
"""

# ---------- Globals ----------
window = None
shader = None
vao = None
textures = []
fbos = []
current = 0
cycle = 0.0
omega_time = 0.0

lattice_width = 32
num_instances_base = 50_000_000  # start huge
scale_interval = 2.0
sub_tile_height = 256
max_tex_height = 2048
threshold = np.sqrt(1.6180339887)

# Precompute recursion tables
phi = 1.6180339887
phi_powers = np.array([1.0 / pow(phi, 7*(i+1)) for i in range(72)], dtype=np.float32)
fib_powers = np.array([1.0/float(i+1) for i in range(72)], dtype=np.float32)
dyadic = np.array([2.0**(-i) for i in range(72)], dtype=np.float32)
primes = np.array([1.0/float(p) for p in list(range(2, 379))[:72]], dtype=np.float32)

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

# ---------- GPU Detection ----------
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.")
    return 100_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)

    # Upload recursion buffers
    block_index = glGetUniformBlockIndex(shader, "PrismaticBlock")
    glUniformBlockBinding(shader, block_index, 0)
    ubo = glGenBuffers(1)
    glBindBuffer(GL_UNIFORM_BUFFER, ubo)
    data = np.concatenate([phi_powers, fib_powers, dyadic, primes])
    glBufferData(GL_UNIFORM_BUFFER, data.nbytes, data, GL_STATIC_DRAW)
    glBindBufferBase(GL_UNIFORM_BUFFER, 0, ubo)

    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)
            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 Prismatic Recursive Analog Model - Full GPU Saturation")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutMainLoop()

if __name__=="__main__":
    main()
