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
import tkinter as tk
from tkinter import ttk

# ---------- Globals ----------
window = None
shader = None
particle_vao = None
particle_vbo = None
scaffold_vao = None
scaffold_vbo = None

max_particles = 10000
particle_positions = None
particle_velocities = None
particle_colors = None
particle_trails = None
max_scaffold = 1000
scaffold_positions = None
scaffold_colors = None
num_scaffold = 0

polar_mode = False
room_mesh = None
frame_times = []

# Wi-Fi
wifi = pywifi.PyWiFi()
iface = wifi.interfaces()[0]
networks = []
ssid_to_index = {}
next_index = 0
stability_tracker = {}
stable_threshold = 3.0
stable_frames = 30

# Control variables (to be updated by sliders)
amplitude = 1.0
noise = 0.3
morph = 0.0
persistence = 0.92
radius_speed = 0.02
radii = 0.0

# Orbit camera
cam_angle_x, cam_angle_y = 30, 30
cam_distance = 4.5  # Adjusted for better initial view
mouse_last = None
orbiting = False

# Tkinter GUI
root = None
amp_var = None
noise_var = None
morph_var = None
pers_var = None

# ---------- 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){
        p.x = pos.x * cos(pos.y);
        p.y = pos.x * sin(pos.y);
        p.z = pos.z;
    }
    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, particle_trails, scaffold_positions, scaffold_colors
    particle_positions = np.random.uniform(-1,1,(max_particles,3)).astype(np.float32)
    particle_velocities = np.random.uniform(-0.002,0.002,(max_particles,3)).astype(np.float32)
    particle_colors = np.zeros((max_particles,3),dtype=np.float32)
    particle_trails = np.zeros((max_particles,3),dtype=np.float32)
    scaffold_positions = np.zeros((max_scaffold,3),dtype=np.float32)
    scaffold_colors = np.ones((max_scaffold,3),dtype=np.float32) * 0.5

# ---------- 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
        [-1,-1,0],[-1,-1,2],[-1,1,2],[-1,1,0], # left wall
        [1,-1,0],[1,-1,2],[1,1,2],[1,1,0]     # right wall
    ], dtype=np.float32)

# ---------- OpenGL Init ----------
def init_gl():
    global shader, particle_vao, particle_vbo, scaffold_vao, scaffold_vbo
    shader = compileProgram(
        compileShader(VERTEX_SRC, GL_VERTEX_SHADER),
        compileShader(FRAGMENT_SRC, GL_FRAGMENT_SHADER)
    )
    particle_vao = glGenVertexArrays(1)
    glBindVertexArray(particle_vao)
    particle_vbo = glGenBuffers(2)
    glBindBuffer(GL_ARRAY_BUFFER, particle_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)
    glBindBuffer(GL_ARRAY_BUFFER, particle_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)

    scaffold_vao = glGenVertexArrays(1)
    glBindVertexArray(scaffold_vao)
    scaffold_vbo = glGenBuffers(2)
    glBindBuffer(GL_ARRAY_BUFFER, scaffold_vbo[0])
    glBufferData(GL_ARRAY_BUFFER, scaffold_positions.nbytes, scaffold_positions, GL_DYNAMIC_DRAW)
    glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,0,None)
    glEnableVertexAttribArray(0)
    glBindBuffer(GL_ARRAY_BUFFER, scaffold_vbo[1])
    glBufferData(GL_ARRAY_BUFFER, scaffold_colors.nbytes, scaffold_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)
    glEnable(GL_DEPTH_TEST)

# ---------- Tkinter GUI Setup ----------
def init_gui():
    global root, amp_var, noise_var, morph_var, pers_var
    root = tk.Tk()
    root.title("Wi-Fi Room Visualizer Controls")

    def make_slider(label, from_, to, init):
        frame = ttk.Frame(root)
        frame.pack(fill='x')
        ttk.Label(frame, text=label).pack(side='left')
        var = tk.DoubleVar(value=init)
        slider = ttk.Scale(frame, from_=from_, to=to, orient='horizontal', variable=var)
        slider.pack(side='right', fill='x', expand=True)
        return var

    amp_var = make_slider("Amplitude", 0.1, 2.0, 1.0)
    noise_var = make_slider("Noise", 0.0, 1.0, 0.3)
    morph_var = make_slider("Morph (Polar→Cartesian)", 0.0, 1.0, 0.0)
    pers_var = make_slider("Persistence (trail length)", 0.5, 0.99, 0.92)

    # Non-blocking Tkinter update
    def update_tk():
        root.update()
        glutPostRedisplay()
        root.after(10, update_tk)

    root.after(10, update_tk)

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

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

# ---------- Display ----------
def display():
    global particle_positions, particle_velocities, particle_colors, particle_trails
    global num_scaffold, radii, next_index, amplitude, noise, morph, persistence
    start = time.time()
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glUseProgram(shader)

    # Update from sliders
    try:
        amplitude = amp_var.get()
        noise = noise_var.get()
        morph = morph_var.get()
        persistence = pers_var.get()
        polar_mode = morph > 0.5
    except:
        pass  # GUI not ready yet

    # Radial sweep
    radii += radius_speed
    if radii > 1.5:
        radii = 0.0

    # Update particles
    particle_positions += particle_velocities + np.random.normal(0, noise, particle_positions.shape).astype(np.float32)
    np.clip(particle_positions, -1, 1, out=particle_positions)
    particle_positions *= amplitude

    # Assign networks to particles
    next_index = 0
    for ssid, signal in networks:
        if ssid not in ssid_to_index:
            ssid_to_index[ssid] = next_index
            next_index += 1
        i = ssid_to_index[ssid]
        if i >= max_particles:
            continue
        strength = np.clip((signal + 100) / 50.0, 0, 1)
        particle_colors[i] = [1 - strength, strength, 0.2]
        particle_trails[i] = persistence * particle_trails[i] + (1 - persistence) * particle_positions[i]

        # Stability check
        stability_tracker.setdefault(ssid, []).append(signal)
        if len(stability_tracker[ssid]) > stable_frames:
            hist = stability_tracker[ssid][-stable_frames:]
            if np.var(hist) < stable_threshold:
                mean_signal = np.mean(hist)
                mean_strength = np.clip((mean_signal + 100) / 50.0, 0, 1)
                pos = particle_positions[i] * (0.5 + mean_strength * 1.0)
                if num_scaffold < max_scaffold:
                    scaffold_positions[num_scaffold] = pos
                    scaffold_colors[num_scaffold] = [0.8, 0.8, 1.0]
                    num_scaffold += 1
                stability_tracker[ssid] = []

    # Update particle GPU buffers
    glBindBuffer(GL_ARRAY_BUFFER, particle_vbo[0])
    glBufferSubData(GL_ARRAY_BUFFER, 0, particle_positions.nbytes, particle_positions)
    glBindBuffer(GL_ARRAY_BUFFER, particle_vbo[1])
    glBufferSubData(GL_ARRAY_BUFFER, 0, particle_colors.nbytes, particle_colors)

    # Morph uniform
    morph_loc = glGetUniformLocation(shader, "morph")
    glUniform1f(morph_loc, morph)
    if polar_mode:
        particle_positions[:, 0] *= radii

    # 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(particle_vao)
    glDrawArrays(GL_POINTS, 0, min(next_index, max_particles))

    # Draw scaffold
    if num_scaffold > 0:
        glBindBuffer(GL_ARRAY_BUFFER, scaffold_vbo[0])
        glBufferSubData(GL_ARRAY_BUFFER, 0, scaffold_positions[:num_scaffold].nbytes, scaffold_positions[:num_scaffold])
        glBindBuffer(GL_ARRAY_BUFFER, scaffold_vbo[1])
        glBufferSubData(GL_ARRAY_BUFFER, 0, scaffold_colors[:num_scaffold].nbytes, scaffold_colors[:num_scaffold])
        glPointSize(10.0)
        glBindVertexArray(scaffold_vao)
        glDrawArrays(GL_POINTS, 0, num_scaffold)
        glPointSize(4.0)

    # Draw room mesh
    glDisable(GL_DEPTH_TEST)
    glColor3f(1, 0, 0)
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[0:4]: glVertex3f(*v)
    glEnd()
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[4:8]: glVertex3f(*v)
    glEnd()
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[8:12]: glVertex3f(*v)
    glEnd()
    glBegin(GL_LINE_LOOP)
    for v in room_mesh[12:16]: glVertex3f(*v)
    glEnd()
    glEnable(GL_DEPTH_TEST)

    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):
    if key == b'p':
        global polar_mode, morph
        polar_mode = not polar_mode
        morph = 1.0 if polar_mode else 0.0
        morph_var.set(morph)
        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 Reader Visualizer")
    init_gl()
    init_gui()
    glutDisplayFunc(display)
    glutIdleFunc(idle)
    glutKeyboardFunc(keyboard)
    glutMouseFunc(mouse)
    glutMotionFunc(motion)
    try: glutMouseWheelFunc(mouse_wheel)
    except: pass
    glutMainLoop()

if __name__ == "__main__":
    main()