import sys
import time
import numpy as np
from OpenGL.GL import *
from OpenGL.GL.shaders import compileProgram, compileShader
from OpenGL.GLUT import glutInit, glutCreateWindow, glutInitDisplayMode, GLUT_RGB, GLUT_DOUBLE
from OpenGL.GLUT import glutInitWindowSize, glutHideWindow, GLUT_DEPTH

# ------------------------------------------------------
# Config
# ------------------------------------------------------
VEC_SIZE = 16  # bytes per vec (4 floats)
TOTAL_ACTIVE_BYTES = int(7.5 * (1024 ** 3))  # ~7.5 GB
MODE = "virtual" if "--mode" not in sys.argv else sys.argv[sys.argv.index("--mode") + 1]

# ------------------------------------------------------
# Shader Sources
# ------------------------------------------------------
VERT_SHADER = """
#version 430
layout(location = 0) in uint idx;
void main() {
    gl_Position = vec4(float(idx % 2u) * 2.0 - 1.0,
                       float((idx/2u) % 2u) * 2.0 - 1.0,
                       0.0, 1.0);
}
"""

FRAG_SHADER_VIRTUAL = """
#version 430
out vec4 fragColor;
uniform uint start_index;
void main() {
    // Pure compute – churn numbers, no memory
    uint v = start_index + uint(gl_FragCoord.x) + uint(gl_FragCoord.y)*1024u;
    float f = float((v * 1664525u + 1013904223u) & 0xFFFFFFFFu);
    fragColor = vec4(f*0.0000001, f*0.00000001, f*0.000000001, 1.0);
}
"""

FRAG_SHADER_ACTIVE = """
#version 430
layout(std430, binding=0) buffer Data {
    vec4 data[];
};
out vec4 fragColor;
uniform uint start_index;
void main() {
    uint idx = start_index + uint(gl_FragCoord.x) + uint(gl_FragCoord.y)*1024u;
    data[idx % data.length()] *= 1.000001; // force memory read/write
    fragColor = vec4(data[idx % data.length()].xyz, 1.0);
}
"""

# ------------------------------------------------------
# Benchmark Executor
# ------------------------------------------------------
class HDGLExecutor:
    def __init__(self, mode):
        self.mode = mode
        glutInit()
        glutInitDisplayMode(GLUT_RGB | GLUT_DOUBLE | GLUT_DEPTH)
        glutInitWindowSize(64, 64)
        glutCreateWindow(b"HDGL Benchmark")
        glutHideWindow()

        frag_src = FRAG_SHADER_VIRTUAL if mode == "virtual" else FRAG_SHADER_ACTIVE
        self.shader = compileProgram(
            compileShader(VERT_SHADER, GL_VERTEX_SHADER),
            compileShader(frag_src, GL_FRAGMENT_SHADER)
        )

        if mode == "active":
            # Allocate big buffer (~7.5GB)
            nvecs = TOTAL_ACTIVE_BYTES // VEC_SIZE
            print(f"Allocating {nvecs:,} active vectors (~{TOTAL_ACTIVE_BYTES/1024**3:.2f} GB)")
            self.ssbo = glGenBuffers(1)
            glBindBuffer(GL_SHADER_STORAGE_BUFFER, self.ssbo)
            glBufferData(GL_SHADER_STORAGE_BUFFER, TOTAL_ACTIVE_BYTES, None, GL_DYNAMIC_DRAW)
            glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, self.ssbo)

    def run(self, iterations=50):
        glUseProgram(self.shader)
        start_time = time.time()
        processed = 0

        for i in range(iterations):
            glUniform1ui(glGetUniformLocation(self.shader, "start_index"), i * 1024 * 1024)
            glDrawArrays(GL_POINTS, 0, 1024 * 1024)  # 1M points per dispatch
            processed += 1024 * 1024
            glFinish()

        elapsed = time.time() - start_time
        rate = processed / elapsed
        print(f"[{self.mode.upper()} MODE]")
        print(f"Processed {processed:,} vectors in {elapsed:.2f}s "
              f"({rate:,.0f} vec/s)")

# ------------------------------------------------------
# Entry
# ------------------------------------------------------
if __name__ == "__main__":
    executor = HDGLExecutor(MODE)
    executor.run()
