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 = 1000
        self.camera = {'x': 30, 'y': 30, 'dist': 5, 'orbiting': False, 'last_pos': None}
        
        # Data
        self.particles = np.random.uniform(-0.8, 0.8, (self.max_particles, 3)).astype(np.float32)
        self.velocities = np.random.uniform(-0.005, 0.005, (self.max_particles, 3)).astype(np.float32)
        self.colors = np.random.uniform(0.2, 1.0, (self.max_particles, 3)).astype(np.float32)  # Start with some color
        self.scaffold = {'pos': np.zeros((200, 3), dtype=np.float32), 'count': 0}
        
        # State
        self.networks = []
        self.ssid_map = {}
        self.params = {'amp': 1.0, 'noise': 0.02, 'morph': 0.0, 'persist': 0.9}
        self.paused = False
        self.time = 0
        self.active_particles = self.max_particles  # Start with all particles visible
        
        # Setup
        self._init_wifi()
        self._init_gui()
        
    def _init_wifi(self):
        def scan():
            try:
                wifi = pywifi.PyWiFi()
                iface = wifi.interfaces()[0]
                print("[WiFi] Starting scan...")
                while True:
                    try:
                        iface.scan()
                        time.sleep(2)
                        results = iface.scan_results()
                        new_networks = [(r.ssid, r.signal) for r in results if r.ssid]
                        if new_networks != self.networks:
                            self.networks = new_networks
                            print(f"[WiFi] Found {len(self.networks)} networks")
                    except Exception as e:
                        print(f"[WiFi] Scan error: {e}")
                        time.sleep(5)
            except Exception as e:
                print(f"[WiFi] Init failed: {e}")
                # Create fake data for testing
                self.networks = [("TestNet1", -40), ("TestNet2", -60), ("TestNet3", -80)]
                
        threading.Thread(target=scan, daemon=True).start()
    
    def _init_gui(self):
        def setup():
            try:
                root = tk.Tk()
                root.title("WiFi Viz Controls")
                root.geometry("300x250")
                
                # Sliders
                self.gui_vars = {}
                for name, (min_val, max_val, init) in [
                    ('amp', (0.1, 2.0, 1.0)), 
                    ('noise', (0, 0.1, 0.02)), 
                    ('morph', (0, 1, 0)), 
                    ('persist', (0.8, 0.99, 0.9))
                ]:
                    frame = ttk.Frame(root)
                    frame.pack(fill='x', padx=5, pady=2)
                    ttk.Label(frame, text=name.title(), width=8).pack(side='left')
                    var = tk.DoubleVar(value=init)
                    scale = ttk.Scale(frame, from_=min_val, to=max_val, variable=var, length=200)
                    scale.pack(side='right')
                    self.gui_vars[name] = var
                
                # Buttons
                btn_frame = ttk.Frame(root)
                btn_frame.pack(fill='x', padx=5, pady=10)
                ttk.Button(btn_frame, text="Pause", command=self._toggle_pause).pack(side='left', padx=2)
                ttk.Button(btn_frame, text="Clear", command=self._clear_scaffold).pack(side='left', padx=2)
                ttk.Button(btn_frame, text="Reset", command=self._reset_view).pack(side='left', padx=2)
                
                # Status
                self.status_var = tk.StringVar(value="Starting...")
                ttk.Label(root, textvariable=self.status_var).pack()
                
                # Update loop
                def update():
                    try:
                        for k, v in self.gui_vars.items():
                            self.params[k] = v.get()
                        self.status_var.set(f"Networks: {len(self.networks)} | Particles: {self.active_particles}")
                    except: pass
                    root.after(100, update)
                update()
                
                print("[GUI] Controls ready")
                root.mainloop()
            except Exception as e:
                print(f"[GUI] Error: {e}")
                
        threading.Thread(target=setup, daemon=True).start()
        
    def _toggle_pause(self):
        self.paused = not self.paused
        print(f"[GUI] {'Paused' if self.paused else 'Resumed'}")
        
    def _clear_scaffold(self):
        self.scaffold['count'] = 0
        print("[GUI] Scaffold cleared")
        
    def _reset_view(self):
        self.camera = {'x': 30, 'y': 30, 'dist': 5, 'orbiting': False, 'last_pos': None}
        print("[GUI] View reset")
    
    def init_gl(self):
        # Simple shaders that should work
        vertex_shader = """
        #version 120
        attribute vec3 position;
        attribute vec3 color;
        varying vec3 v_color;
        uniform float u_morph;
        uniform float u_time;
        
        void main() {
            vec3 pos = position;
            
            // Apply morph transformation
            if (u_morph > 0.1) {
                float r = length(pos.xy);
                float theta = atan(pos.y, pos.x);
                pos.x = mix(pos.x, r * cos(theta + u_time * 0.1), u_morph);
                pos.y = mix(pos.y, r * sin(theta + u_time * 0.1), u_morph);
            }
            
            gl_Position = gl_ModelViewProjectionMatrix * vec4(pos, 1.0);
            gl_PointSize = 6.0;
            v_color = color;
        }
        """
        
        fragment_shader = """
        #version 120
        varying vec3 v_color;
        
        void main() {
            // Make round points
            vec2 coord = gl_PointCoord - vec2(0.5);
            if (length(coord) > 0.5) discard;
            
            gl_FragColor = vec4(v_color, 0.8);
        }
        """
        
        try:
            self.shader = compileProgram(
                compileShader(vertex_shader, GL_VERTEX_SHADER),
                compileShader(fragment_shader, GL_FRAGMENT_SHADER)
            )
            print("[GL] Shaders compiled successfully")
        except Exception as e:
            print(f"[GL] Shader error: {e}")
            self.shader = None
            return
        
        # Create VAO and VBOs
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.vao)
        
        # Position buffer
        self.pos_vbo = glGenBuffers(1)
        glBindBuffer(GL_ARRAY_BUFFER, self.pos_vbo)
        glBufferData(GL_ARRAY_BUFFER, self.particles.nbytes, self.particles, GL_DYNAMIC_DRAW)
        pos_loc = glGetAttribLocation(self.shader, "position")
        glVertexAttribPointer(pos_loc, 3, GL_FLOAT, GL_FALSE, 0, None)
        glEnableVertexAttribArray(pos_loc)
        
        # Color buffer
        self.color_vbo = glGenBuffers(1)
        glBindBuffer(GL_ARRAY_BUFFER, self.color_vbo)
        glBufferData(GL_ARRAY_BUFFER, self.colors.nbytes, self.colors, GL_DYNAMIC_DRAW)
        color_loc = glGetAttribLocation(self.shader, "color")
        glVertexAttribPointer(color_loc, 3, GL_FLOAT, GL_FALSE, 0, None)
        glEnableVertexAttribArray(color_loc)
        
        # 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.05, 0.05, 0.1, 1.0)
        
        print("[GL] OpenGL initialized")
        
    def update(self):
        if self.paused: 
            return
        
        self.time += 0.02
        
        # Always move particles
        noise = np.random.normal(0, self.params['noise'], self.particles.shape).astype(np.float32)
        self.particles += self.velocities + noise
        
        # Keep particles in bounds
        np.clip(self.particles, -1.5, 1.5, out=self.particles)
        
        # Bounce off walls
        for i in range(3):
            mask = (self.particles[:, i] >= 1.4) | (self.particles[:, i] <= -1.4)
            self.velocities[mask, i] *= -1
        
        # Scale by amplitude
        display_particles = self.particles * self.params['amp']
        
        # Update colors from WiFi or use cycling colors
        if self.networks:
            for i, (ssid, signal) in enumerate(self.networks):
                if i >= self.max_particles:
                    break
                    
                if ssid not in self.ssid_map:
                    self.ssid_map[ssid] = i
                
                idx = self.ssid_map[ssid] % self.max_particles
                strength = np.clip((signal + 100) / 50.0, 0, 1)
                
                # Color based on signal strength
                self.colors[idx] = [1.0 - strength, strength, 0.5]
                
                # Occasionally add to scaffold
                if np.random.random() < 0.005 and self.scaffold['count'] < 200:
                    pos = display_particles[idx] * (0.3 + strength * 0.7)
                    self.scaffold['pos'][self.scaffold['count']] = pos
                    self.scaffold['count'] += 1
            
            self.active_particles = min(len(self.networks), self.max_particles)
        else:
            # Cycle colors when no WiFi
            for i in range(self.max_particles):
                hue = (self.time + i * 0.1) % 1.0
                self.colors[i] = [
                    0.5 + 0.5 * np.sin(hue * 6.28),
                    0.5 + 0.5 * np.sin((hue + 0.33) * 6.28),
                    0.5 + 0.5 * np.sin((hue + 0.66) * 6.28)
                ]
            self.active_particles = self.max_particles
    
    def render(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        if not self.shader:
            return
            
        # Set up 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)
        
        # Use shader
        glUseProgram(self.shader)
        
        # Set uniforms
        morph_loc = glGetUniformLocation(self.shader, "u_morph")
        time_loc = glGetUniformLocation(self.shader, "u_time")
        if morph_loc >= 0:
            glUniform1f(morph_loc, self.params['morph'])
        if time_loc >= 0:
            glUniform1f(time_loc, self.time)
        
        # Update buffers
        glBindBuffer(GL_ARRAY_BUFFER, self.pos_vbo)
        glBufferSubData(GL_ARRAY_BUFFER, 0, self.particles.nbytes, self.particles)
        
        glBindBuffer(GL_ARRAY_BUFFER, self.color_vbo)
        glBufferSubData(GL_ARRAY_BUFFER, 0, self.colors.nbytes, self.colors)
        
        # Draw particles
        glBindVertexArray(self.vao)
        glDrawArrays(GL_POINTS, 0, self.active_particles)
        
        # Draw scaffold points
        if self.scaffold['count'] > 0:
            glUseProgram(0)  # Use fixed pipeline for scaffold
            glColor3f(0.8, 0.8, 1.0)
            glPointSize(12.0)
            glBegin(GL_POINTS)
            for i in range(self.scaffold['count']):
                glVertex3fv(self.scaffold['pos'][i])
            glEnd()
        
        glutSwapBuffers()
    
    def keyboard(self, key, x, y):
        if key == b'p': 
            self.paused = not self.paused
            print(f"{'Paused' if self.paused else 'Resumed'}")
        elif key == b'r': 
            self._reset_view()
        elif key == b's': 
            self._clear_scaffold()
        elif key in b'+=': 
            self.camera['dist'] = max(1, self.camera['dist'] * 0.9)
        elif key == b'-': 
            self.camera['dist'] = min(15, self.camera['dist'] * 1.1)
        elif key == b'\x1b':  # Escape
            sys.exit(0)
    
    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
        elif button == GLUT_RIGHT_BUTTON and state == GLUT_DOWN:
            self._reset_view()
    
    def motion(self, x, y):
        if self.camera['orbiting'] and self.camera['last_pos']:
            dx = x - self.camera['last_pos'][0]
            dy = 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):
        factor = 0.9 if direction > 0 else 1.1
        self.camera['dist'] = np.clip(self.camera['dist'] * factor, 1, 15)
    
    def idle(self):
        self.update()
        glutPostRedisplay()

def main():
    print("WiFi Room Visualizer")
    print("Controls: Mouse=orbit, P=pause, S=clear, R=reset, +/-=zoom, ESC=exit")
    
    viz = WiFiVisualizer()
    
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(1280, 720)
    glutCreateWindow(b"WiFi Visualizer")
    
    # Set perspective
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(60.0, 1280.0/720.0, 0.1, 100.0)
    glMatrixMode(GL_MODELVIEW)
    
    viz.init_gl()
    
    glutDisplayFunc(viz.render)
    glutIdleFunc(viz.idle)
    glutKeyboardFunc(viz.keyboard)
    glutMouseFunc(viz.mouse)
    glutMotionFunc(viz.motion)
    
    try:
        glutMouseWheelFunc(viz.wheel)
    except:
        print("[Warning] Mouse wheel not supported")
    
    print("[Main] Starting visualization...")
    glutMainLoop()

if __name__ == "__main__":
    main()