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

# ---------- Globals ----------
window = None
shader = None
vao = None

num_particles = 50_000
particle_positions = None
particle_velocities = None
polar_mode = False
room_mesh = None
frame_times = []

# ---------- Vertex & Fragment Shaders ----------
VERTEX_SRC = """
#version 330
layout(location = 0) in vec3 pos;
uniform float morph; // 0 = Cartesian, 1 = Polar
void main(){
    vec3 p = pos;
    if(morph > 0.5){
        float r = length(p.xy);
        float theta = atan(p.y,p.x);
        p.x = r*cos(theta);
        p.y = r*sin(theta);
    }
    gl_Position = vec4(p,1.0);
    gl_PointSize = 2.0;
}
"""

FRAGMENT_SRC = """
#version 330
out vec4 fragColor;
void main(){
    fragColor = vec4(0.0,1.0,0.0,1.0); // particle green
}
"""

# ---------- Initialize Particles ----------
def init_particles():
    global particle_positions, particle_velocities
    particle_positions = np.random.uniform(-1,1,(num_particles,3)).astype(np.float32)
    particle_velocities = np.random.uniform(-0.002,0.002,(num_particles,3)).astype(np.float32)

# ---------- Initialize Room Geometry Placeholder ----------
def init_room():
    global room_mesh
    # Simple room: cube walls + table placeholder
    room_mesh = np.array([
        [-0.9,-0.9,0.0],[0.9,-0.9,0.0],[0.9,0.9,0.0],[-0.9,0.9,0.0], # floor
        [-0.9,-0.9,1.0],[0.9,-0.9,1.0],[0.9,0.9,1.0],[-0.9,0.9,1.0]   # ceiling
    ], dtype=np.float32)

# ---------- OpenGL Init ----------
def init_gl():
    global shader, vao
    shader = compileProgram(
        compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )

    vao = glGenVertexArrays(1)
    glBindVertexArray(vao)

    vbo = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, vbo)
    glBufferData(GL_ARRAY_BUFFER, particle_positions.nbytes, particle_positions, GL_DYNAMIC_DRAW)
    glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,0,None)
    glEnableVertexAttribArray(0)

    glEnable(GL_PROGRAM_POINT_SIZE)
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    glClearColor(0.0,0.0,0.0,1.0)

# ---------- Display ----------
def display():
    global particle_positions, particle_velocities
    start = time.time()
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glUseProgram(shader)
    
    # Update particle positions (analog vibration)
    particle_positions += particle_velocities
    np.clip(particle_positions, -1,1, out=particle_positions)
    
    # Send updated positions to GPU
    glBindBuffer(GL_ARRAY_BUFFER, vao)
    glBufferSubData(GL_ARRAY_BUFFER, 0, particle_positions.nbytes, particle_positions)
    
    # Morph uniform
    morph_loc = glGetUniformLocation(shader,"morph")
    glUniform1f(morph_loc, 1.0 if polar_mode else 0.0)
    
    # Draw particles
    glBindVertexArray(vao)
    glDrawArrays(GL_POINTS, 0, num_particles)
    
    # Draw room mesh as simple lines
    glColor3f(1.0,0.0,0.0)
    glBegin(GL_LINE_LOOP)
    for v in room_mesh:
        glVertex3f(*v)
    glEnd()
    
    glutSwapBuffers()
    
    frame_times.append(time.time()-start)
    if len(frame_times) > 100:
        frame_times.pop(0)
        print(f"[Perf] Avg frame: {np.mean(frame_times)*1000:.2f} ms")

# ---------- Idle ----------
def idle():
    glutPostRedisplay()

# ---------- Keyboard ----------
def keyboard(key, x, y):
    global polar_mode
    if key == b'p':
        polar_mode = not polar_mode
        print(f"[Mode] Polar: {polar_mode}")

# ---------- Main ----------
def main():
    init_particles()
    init_room()
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(1280,720)
    glutCreateWindow(b"HDGL Analog Particle Room Visualizer")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutKeyboardFunc(keyboard)
    glutMainLoop()

if __name__=="__main__":
    main()
