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

class WiFiVisualizer:
    def __init__(self):
        # Config
        self.max_particles = 5000
        self.camera = {'x': 30, 'y': 30, 'dist': 5, 'orbiting': False, 'last_pos': None}
        
        # Data
        self.particles = np.random.uniform(-1, 1, (self.max_particles, 3)).astype(np.float32)
        self.velocities = np.random.uniform(-0.001, 0.001, (self.max_particles, 3)).astype(np.float32)
        self.colors = np.zeros((self.max_particles, 3), dtype=np.float32)
        self.scaffold = {'pos': np.zeros((1000, 3), dtype=np.float32), 'count': 0}
        
        # State
        self.networks = []
        self.ssid_map = {}
        self.params = {'amp': 1.0, 'noise': 0.2, 'morph': 0.0, 'persist': 0.9}
        self.paused = False
        self.time = 0
        
        # Setup
        self._init_wifi()
        self._init_gui()
        
    def _init_wifi(self):
        def scan():
            try:
                wifi = pywifi.PyWiFi()
                iface = wifi.interfaces()[0]
                while True:
                    iface.scan()
                    time.sleep(1)
                    results = iface.scan_results()
                    self.networks = [(r.ssid, r.signal) for r in results if r.ssid]
            except: pass
        threading.Thread(target=scan, daemon=True).start()
    
    def _init_gui(self):
        def setup():
            try:
                root = tk.Tk()
                root.title("WiFi Viz")
                root.geometry("300x200")
                
                # Sliders
                vars = {}
                for name, (min_val, max_val, init) in [
                    ('amp', (0.1, 2.0, 1.0)), ('noise', (0, 0.5, 0.2)), 
                    ('morph', (0, 1, 0)), ('persist', (0.5, 0.99, 0.9))
                ]:
                    ttk.Label(root, text=name.title()).pack()
                    var = tk.DoubleVar(value=init)
                    ttk.Scale(root, from_=min_val, to=max_val, variable=var).pack(fill='x')
                    vars[name] = var
                
                # Update params
                def update():
                    for k, v in vars.items():
                        self.params[k] = v.get()
                    root.after(100, update)
                update()
                
                # Controls
                ttk.Button(root, text="Pause", command=lambda: setattr(self, 'paused', not self.paused)).pack()
                ttk.Button(root, text="Clear", command=lambda: setattr(self.scaffold, 'count', 0)).pack()
                root.mainloop()
            except: pass
        threading.Thread(target=setup, daemon=True).start()
    
    def init_gl(self):
        # Shaders
        vs = """
        #version 330
        in vec3 pos; in vec3 col;
        uniform float morph, time;
        out vec3 vCol;
        void main(){
            vec3 p = pos;
            if(morph > 0.5){
                float r = length(p.xy) * (1.0 + sin(time) * 0.3);
                p.xy = vec2(r * cos(atan(p.y, p.x)), r * sin(atan(p.y, p.x)));
            }
            gl_Position = vec4(p, 1.0);
            gl_PointSize = mix(4.0, 8.0, length(col));
            vCol = col;
        }"""
        
        fs = """
        #version 330
        in vec3 vCol; out vec4 fragCol;
        void main(){
            vec2 c = gl_PointCoord - 0.5;
            if(length(c) > 0.5) discard;
            fragCol = vec4(vCol, 0.8);
        }"""
        
        self.shader = compileProgram(compileShader(vs, GL_VERTEX_SHADER), compileShader(fs, GL_FRAGMENT_SHADER))
        
        # VAO setup
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.vao)
        
        self.vbos = glGenBuffers(2)
        glBindBuffer(GL_ARRAY_BUFFER, self.vbos[0])
        glBufferData(GL_ARRAY_BUFFER, self.particles.nbytes, self.particles, GL_DYNAMIC_DRAW)
        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, None)
        glEnableVertexAttribArray(0)
        
        glBindBuffer(GL_ARRAY_BUFFER, self.vbos[1])
        glBufferData(GL_ARRAY_BUFFER, self.colors.nbytes, self.colors, GL_DYNAMIC_DRAW)
        glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, None)
        glEnableVertexAttribArray(1)
        
        # OpenGL state
        glEnable(GL_PROGRAM_POINT_SIZE)
        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glEnable(GL_DEPTH_TEST)
        glClearColor(0.02, 0.02, 0.1, 1)
        
    def update(self):
        if self.paused: return
        
        self.time += 0.05
        
        # Update particles
        noise = np.random.normal(0, self.params['noise'], self.particles.shape).astype(np.float32)
        self.particles += self.velocities + noise
        np.clip(self.particles, -1, 1, out=self.particles)
        self.particles *= self.params['amp']
        
        # Update colors from WiFi
        active = 0
        for ssid, signal in self.networks[:self.max_particles]:
            if ssid not in self.ssid_map:
                self.ssid_map[ssid] = len(self.ssid_map)
            
            idx = self.ssid_map[ssid]
            if idx >= self.max_particles: continue
            
            strength = np.clip((signal + 100) / 50, 0, 1)
            self.colors[idx] = [1-strength, strength, 0.3]
            active = max(active, idx + 1)
            
            # Add to scaffold occasionally
            if np.random.random() < 0.001 and self.scaffold['count'] < 1000:
                self.scaffold['pos'][self.scaffold['count']] = self.particles[idx] * strength
                self.scaffold['count'] += 1
        
        self.active_particles = active
    
    def render(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glUseProgram(self.shader)
        
        # Camera
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        glTranslatef(0, 0, -self.camera['dist'])
        glRotatef(self.camera['y'], 1, 0, 0)
        glRotatef(self.camera['x'], 0, 1, 0)
        
        # Update uniforms
        glUniform1f(glGetUniformLocation(self.shader, "morph"), self.params['morph'])
        glUniform1f(glGetUniformLocation(self.shader, "time"), self.time)
        
        # Update buffers
        if hasattr(self, 'active_particles'):
            glBindBuffer(GL_ARRAY_BUFFER, self.vbos[0])
            glBufferSubData(GL_ARRAY_BUFFER, 0, self.particles.nbytes, self.particles)
            glBindBuffer(GL_ARRAY_BUFFER, self.vbos[1])
            glBufferSubData(GL_ARRAY_BUFFER, 0, self.colors.nbytes, self.colors)
            
            # Draw particles
            glBindVertexArray(self.vao)
            glDrawArrays(GL_POINTS, 0, self.active_particles)
            
            # Draw scaffold
            if self.scaffold['count'] > 0:
                glColor3f(0.8, 0.8, 1.0)
                glPointSize(10)
                glBegin(GL_POINTS)
                for i in range(self.scaffold['count']):
                    glVertex3fv(self.scaffold['pos'][i])
                glEnd()
                glPointSize(4)
        
        glutSwapBuffers()
    
    def keyboard(self, key, x, y):
        if key == b'p': self.paused = not self.paused
        elif key == b'r': self.camera.update({'x': 30, 'y': 30, 'dist': 5})
        elif key == b's': self.scaffold['count'] = 0
        elif key in b'+=': self.camera['dist'] = max(1, self.camera['dist'] * 0.9)
        elif key == b'-': self.camera['dist'] = min(10, self.camera['dist'] * 1.1)
    
    def mouse(self, button, state, x, y):
        if button == GLUT_LEFT_BUTTON:
            self.camera['orbiting'] = (state == GLUT_DOWN)
            self.camera['last_pos'] = (x, y) if self.camera['orbiting'] else None
    
    def motion(self, x, y):
        if self.camera['orbiting'] and self.camera['last_pos']:
            dx, dy = x - self.camera['last_pos'][0], y - self.camera['last_pos'][1]
            self.camera['x'] += dx * 0.5
            self.camera['y'] = np.clip(self.camera['y'] + dy * 0.5, -90, 90)
            self.camera['last_pos'] = (x, y)
    
    def wheel(self, button, direction, x, y):
        self.camera['dist'] = np.clip(self.camera['dist'] * (0.9 if direction > 0 else 1.1), 1, 10)
    
    def idle(self):
        self.update()
        glutPostRedisplay()

def main():
    viz = WiFiVisualizer()
    
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(1280, 720)
    glutCreateWindow(b"WiFi Visualizer")
    
    viz.init_gl()
    
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(60, 16/9, 0.1, 100)
    
    glutDisplayFunc(viz.render)
    glutIdleFunc(viz.idle)
    glutKeyboardFunc(viz.keyboard)
    glutMouseFunc(viz.mouse)
    glutMotionFunc(viz.motion)
    try: glutMouseWheelFunc(viz.wheel)
    except: pass
    
    print("WiFi Visualizer | P=pause S=clear R=reset +/-=zoom")
    glutMainLoop()

if __name__ == "__main__":
    main()