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

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

num_particles = 5000  # reduced for smoother GPU update
particle_positions = None
particle_velocities = None
particle_colors = None
radii = None

polar_mode = False
room_mesh = None
frame_times = []

# Wi-Fi
wifi = pywifi.PyWiFi()
iface = wifi.interfaces()[0]
networks = []

# Orbit camera
cam_angle_x, cam_angle_y = 30, 30
cam_distance = 3.0
mouse_last = None
orbiting = False

# ---------- Vertex & Fragment Shaders ----------
VERTEX_SRC = """
#version 330
layout(location = 0) in vec3 pos;
layout(location = 1) in vec3 col;
out vec3 vColor;
uniform float morph;
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 = 4.0;
    vColor = col;
}
"""

FRAGMENT_SRC = """
#version 330
in vec3 vColor;
out vec4 fragColor;
void main(){
    fragColor = vec4(vColor,1.0);
}
"""

# ---------- Initialize Particles ----------
def init_particles():
    global particle_positions, particle_velocities, particle_colors, radii
    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)
    particle_colors = np.ones((num_particles,3),dtype=np.float32)*0.0
    radii = np.zeros(num_particles,dtype=np.float32)

# ---------- Initialize Room Mesh ----------
def init_room():
    global room_mesh
    room_mesh = np.array([
        [-1,-1,0],[1,-1,0],[1,1,0],[-1,1,0],  # floor
        [-1,-1,2],[1,-1,2],[1,1,2],[-1,1,2]   # ceiling
    ], dtype=np.float32)

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

    # Position buffer
    vbo = glGenBuffers(2)
    glBindBuffer(GL_ARRAY_BUFFER, vbo[0])
    glBufferData(GL_ARRAY_BUFFER, particle_positions.nbytes, particle_positions, GL_DYNAMIC_DRAW)
    glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,0,None)
    glEnableVertexAttribArray(0)

    # Color buffer
    glBindBuffer(GL_ARRAY_BUFFER, vbo[1])
    glBufferData(GL_ARRAY_BUFFER, particle_colors.nbytes, particle_colors, GL_DYNAMIC_DRAW)
    glVertexAttribPointer(1,3,GL_FLOAT,GL_FALSE,0,None)
    glEnableVertexAttribArray(1)

    glEnable(GL_PROGRAM_POINT_SIZE)
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    glClearColor(0,0,0,1)

# ---------- Wi-Fi Scan Thread ----------
def scan_loop():
    global networks
    while True:
        iface.scan()
        time.sleep(1)
        results = iface.scan_results()
        networks = [(r.ssid,r.signal) for r in results]

threading.Thread(target=scan_loop, daemon=True).start()

# ---------- Display ----------
def display():
    global particle_positions, particle_velocities, particle_colors, radii
    start = time.time()
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glUseProgram(shader)

    # --- Update particle positions ---
    particle_positions += particle_velocities
    np.clip(particle_positions, -1,1, out=particle_positions)

    # Map Wi-Fi signals to some particles
    for i, (ssid, signal) in enumerate(networks[:num_particles]):
        strength = np.clip((signal+100)/50.0,0,1)
        particle_colors[i] = [1-strength,strength,0.2]

    # Send updated positions/colors to GPU
    glBindBuffer(GL_ARRAY_BUFFER, vbo[0])
    glBufferSubData(GL_ARRAY_BUFFER, 0, particle_positions.nbytes, particle_positions)
    glBindBuffer(GL_ARRAY_BUFFER, vbo[1])
    glBufferSubData(GL_ARRAY_BUFFER, 0, particle_colors.nbytes, particle_colors)

    # Morph uniform
    morph_loc = glGetUniformLocation(shader,"morph")
    glUniform1f(morph_loc, 1.0 if polar_mode else 0.0)

    # --- Orbit camera ---
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    glTranslatef(0,0,-cam_distance)
    glRotatef(cam_angle_y,1,0,0)
    glRotatef(cam_angle_x,0,1,0)

    # Draw particles
    glBindVertexArray(vao)
    glDrawArrays(GL_POINTS,0,num_particles)

    # Draw room mesh as lines
    glColor3f(1,0,0)
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[:4]:
        glVertex3f(*v)
    glEnd()
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[4:]:
        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}")

# ---------- Mouse ----------
def mouse(button,state,x,y):
    global orbiting, mouse_last
    if button==GLUT_LEFT_BUTTON:
        orbiting = state==GLUT_DOWN
        mouse_last = (x,y)

def motion(x,y):
    global cam_angle_x, cam_angle_y, mouse_last
    if orbiting and mouse_last:
        dx = x-mouse_last[0]
        dy = y-mouse_last[1]
        cam_angle_x += dx*0.5
        cam_angle_y += dy*0.5
        cam_angle_y = np.clip(cam_angle_y,-90,90)
        mouse_last = (x,y)

def mouse_wheel(button, dir, x, y):
    global cam_distance
    cam_distance *= 0.9 if dir>0 else 1.1
    cam_distance = np.clip(cam_distance,1,10)

# ---------- Main ----------
def main():
    init_particles()
    init_room()
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(1280,720)
    glutCreateWindow(b"Wi-Fi Room HDGL Visualizer")
    init_gl()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutKeyboardFunc(keyboard)
    glutMouseFunc(mouse)
    glutMotionFunc(motion)
    # Wheel support (freeglut)
    try: glutMouseWheelFunc(mouse_wheel)
    except: pass
    glEnable(GL_DEPTH_TEST)
    glutMainLoop()

if __name__=="__main__":
    main()
