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

# ---------------------- Configuration ----------------------
NUM_SLOTS = 4096  # full GPU-friendly
NUM_STRANDS = 32  # multi-strand prismatic saturation
WINDOW_W, WINDOW_H = 1280, 720

# ---------------------- Shader Sources --------------------
VERTEX_SHADER = """
#version 330
layout(location = 0) in float a_index;
uniform float u_time;
uniform int u_numSlots;
uniform float u_numStrands;
out float v_slotIndex;

void main(){
    v_slotIndex = a_index;
    float x = mod(a_index, sqrt(float(u_numSlots))) / sqrt(float(u_numSlots)) * 2.0 - 1.0;
    float y = floor(a_index / sqrt(float(u_numSlots))) / sqrt(float(u_numSlots)) * 2.0 - 1.0;
    gl_Position = vec4(x, y, 0.0, 1.0);
    gl_PointSize = 2.0;
}
"""

FRAGMENT_SHADER = """
#version 330
in float v_slotIndex;
uniform sampler1D u_slotData;
uniform float u_time;
uniform int u_numSlots;
uniform int u_numStrands;
out vec4 fragColor;

void main(){
    int idx = int(v_slotIndex) % u_numSlots;
    float slot = texelFetch(u_slotData, idx, 0).r;

    // --- GPU-side dynamic rescaling ---
    float scale = 1.0;
    if(slot > 1e6) scale = 1e-6 * slot;
    slot /= scale;

    // --- Prismatic RGB/Alpha mapping ---
    float phase = fract(u_time*0.05 + slot*0.618);
    float r = sin(phase*6.2831);
    float g = cos(phase*6.2831);
    float b = sin(phase*3.1415);
    float a = smoothstep(0.0, 1.0, slot);

    fragColor = vec4(r*slot, g*slot, b*slot, a);
}
"""

# ---------------------- GLFW & OpenGL Init -----------------
if not glfw.init():
    sys.exit("Failed to initialize GLFW")

window = glfw.create_window(WINDOW_W, WINDOW_H, "HDGL Superglyphs18", None, None)
if not window:
    glfw.terminate()
    sys.exit("Failed to create window")

glfw.make_context_current(window)

# ---------------------- Shader Compilation -----------------
shader = compileProgram(
    compileShader(VERTEX_SHADER, GL_VERTEX_SHADER),
    compileShader(FRAGMENT_SHADER, GL_FRAGMENT_SHADER)
)

# ---------------------- Vertex Array -----------------------
indices = np.arange(NUM_SLOTS, dtype=np.float32)
vao = glGenVertexArrays(1)
vbo = glGenBuffers(1)
glBindVertexArray(vao)
glBindBuffer(GL_ARRAY_BUFFER, vbo)
glBufferData(GL_ARRAY_BUFFER, indices.nbytes, indices, GL_STATIC_DRAW)
glEnableVertexAttribArray(0)
glVertexAttribPointer(0, 1, GL_FLOAT, GL_FALSE, 0, ctypes.c_void_p(0))
glBindBuffer(GL_ARRAY_BUFFER, 0)
glBindVertexArray(0)

# ---------------------- GPU Slot Data ----------------------
slot_data = np.zeros(NUM_SLOTS, dtype=np.float32)
for i in range(NUM_SLOTS):
    slot_data[i] = float(i+1)  # simple deterministic initialization

slot_tex = glGenTextures(1)
glBindTexture(GL_TEXTURE_1D, slot_tex)
glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
glTexImage1D(GL_TEXTURE_1D, 0, GL_R32F, NUM_SLOTS, 0, GL_RED, GL_FLOAT, slot_data)
glBindTexture(GL_TEXTURE_1D, 0)

# ---------------------- Uniform Locations ------------------
glUseProgram(shader)
loc_time = glGetUniformLocation(shader, "u_time")
loc_slots = glGetUniformLocation(shader, "u_numSlots")
loc_strands = glGetUniformLocation(shader, "u_numStrands")
loc_tex = glGetUniformLocation(shader, "u_slotData")
glUniform1i(loc_slots, NUM_SLOTS)
glUniform1i(loc_strands, NUM_STRANDS)
glUniform1i(loc_tex, 0)

# ---------------------- Main Loop --------------------------
t = 0.0
while not glfw.window_should_close(window):
    glfw.poll_events()
    glClearColor(0.0, 0.0, 0.0, 1.0)
    glClear(GL_COLOR_BUFFER_BIT)

    t += 0.016
    glUseProgram(shader)
    glUniform1f(loc_time, t)

    glActiveTexture(GL_TEXTURE0)
    glBindTexture(GL_TEXTURE_1D, slot_tex)

    glBindVertexArray(vao)
    glDrawArrays(GL_POINTS, 0, NUM_SLOTS)
    glBindVertexArray(0)

    glfw.swap_buffers(window)

glfw.terminate()
