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
import json
import os

class RoomMapper:
    def __init__(self):
        # Room structure - persistent
        self.room = {
            'floor': np.array([[-2,-2,0], [2,-2,0], [2,2,0], [-2,2,0]], dtype=np.float32),
            'ceiling': np.array([[-2,-2,2.5], [2,-2,2.5], [2,2,2.5], [-2,2,2.5]], dtype=np.float32),
            'walls': [],
            'furniture': [],
            'fixtures': []  # lights, outlets, etc.
        }
        self.room_file = "room_layout.json"
        self._load_room()
        
        # EMF scanning strategy
        self.emf_scanner = {
            'scan_interval': 3.0,  # Respect WiFi card limits
            'last_scan': 0,
            'networks': {},
            'history': {},  # Track signal over time
            'spatial_map': {}  # Map signals to 3D positions
        }
        
        # Visualization layers
        self.static_layer = StaticRoomRenderer()
        self.emf_layer = EMFRenderer()
        
        # Camera and controls
        self.camera = {'x': 45, 'y': 20, 'dist': 8, 'orbiting': False, 'last_pos': None}
        self.params = {
            'show_room': True,
            'show_emf': True, 
            'emf_intensity': 1.0,
            'time_decay': 0.95,
            'spatial_res': 0.2  # meters per grid cell
        }
        self.time = 0
        
        self._init_wifi()
        self._init_gui()
        
    def _load_room(self):
        """Load persistent room layout"""
        if os.path.exists(self.room_file):
            try:
                with open(self.room_file, 'r') as f:
                    data = json.load(f)
                    # Convert lists back to numpy arrays
                    for key in ['furniture', 'fixtures']:
                        if key in data:
                            self.room[key] = [np.array(item, dtype=np.float32) for item in data[key]]
                print(f"[Room] Loaded {len(self.room['furniture'])} furniture, {len(self.room['fixtures'])} fixtures")
            except Exception as e:
                print(f"[Room] Load error: {e}")
    
    def _save_room(self):
        """Save persistent room layout"""
        try:
            data = {
                'furniture': [item.tolist() for item in self.room['furniture']],
                'fixtures': [item.tolist() for item in self.room['fixtures']]
            }
            with open(self.room_file, 'w') as f:
                json.dump(data, f)
            print("[Room] Layout saved")
        except Exception as e:
            print(f"[Room] Save error: {e}")
    
    def _init_wifi(self):
        """Initialize WiFi with respectful scanning"""
        def smart_scan():
            try:
                wifi = pywifi.PyWiFi()
                iface = wifi.interfaces()[0]
                print("[WiFi] Starting smart scanning (3s intervals)")
                
                while True:
                    current_time = time.time()
                    if current_time - self.emf_scanner['last_scan'] >= self.emf_scanner['scan_interval']:
                        try:
                            iface.scan()
                            time.sleep(1)  # Wait for scan to complete
                            results = iface.scan_results()
                            
                            # Process results
                            new_networks = {}
                            for r in results:
                                if r.ssid:
                                    new_networks[r.ssid] = {
                                        'signal': r.signal,
                                        'freq': getattr(r, 'freq', 2400),
                                        'timestamp': current_time
                                    }
                            
                            self.emf_scanner['networks'] = new_networks
                            self.emf_scanner['last_scan'] = current_time
                            
                            # Update history for stability analysis
                            for ssid, data in new_networks.items():
                                if ssid not in self.emf_scanner['history']:
                                    self.emf_scanner['history'][ssid] = []
                                self.emf_scanner['history'][ssid].append(data['signal'])
                                if len(self.emf_scanner['history'][ssid]) > 20:
                                    self.emf_scanner['history'][ssid].pop(0)
                            
                            print(f"[WiFi] Scanned {len(new_networks)} networks")
                            
                        except Exception as e:
                            print(f"[WiFi] Scan error: {e}")
                    
                    time.sleep(0.5)  # Check interval
                    
            except Exception as e:
                print(f"[WiFi] Init failed: {e}")
                # Create test data
                self._create_test_emf()
        
        threading.Thread(target=smart_scan, daemon=True).start()
    
    def _create_test_emf(self):
        """Create test EMF data when WiFi unavailable"""
        def test_loop():
            while True:
                test_networks = {
                    "Router_5G": {"signal": -35 + np.random.normal(0, 3), "freq": 5000},
                    "Router_2.4G": {"signal": -45 + np.random.normal(0, 5), "freq": 2400},
                    "Neighbor_WiFi": {"signal": -65 + np.random.normal(0, 8), "freq": 2400},
                    "Phone_Hotspot": {"signal": -50 + np.random.normal(0, 10), "freq": 2400}
                }
                self.emf_scanner['networks'] = test_networks
                time.sleep(3)
        
        threading.Thread(target=test_loop, daemon=True).start()
    
    def _init_gui(self):
        def setup():
            try:
                root = tk.Tk()
                root.title("Room Mapper Controls")
                root.geometry("350x400")
                
                # Room controls
                room_frame = ttk.LabelFrame(root, text="Room Mapping", padding=10)
                room_frame.pack(fill='x', padx=5, pady=5)
                
                ttk.Button(room_frame, text="Add Furniture", command=self._add_furniture).pack(side='left', padx=2)
                ttk.Button(room_frame, text="Add Fixture", command=self._add_fixture).pack(side='left', padx=2)
                ttk.Button(room_frame, text="Save Room", command=self._save_room).pack(side='left', padx=2)
                
                # EMF controls
                emf_frame = ttk.LabelFrame(root, text="EMF Visualization", padding=10)
                emf_frame.pack(fill='x', padx=5, pady=5)
                
                self.gui_vars = {}
                for name, (min_val, max_val, init) in [
                    ('emf_intensity', (0.1, 3.0, 1.0)),
                    ('time_decay', (0.8, 0.99, 0.95)),
                    ('spatial_res', (0.1, 0.5, 0.2))
                ]:
                    frame = ttk.Frame(emf_frame)
                    frame.pack(fill='x', pady=2)
                    ttk.Label(frame, text=name.replace('_', ' ').title(), width=12).pack(side='left')
                    var = tk.DoubleVar(value=init)
                    ttk.Scale(frame, from_=min_val, to=max_val, variable=var, length=150).pack(side='right')
                    self.gui_vars[name] = var
                
                # Toggle buttons
                toggle_frame = ttk.Frame(emf_frame)
                toggle_frame.pack(fill='x', pady=5)
                
                self.show_room_var = tk.BooleanVar(value=True)
                ttk.Checkbutton(toggle_frame, text="Show Room", variable=self.show_room_var).pack(side='left')
                
                self.show_emf_var = tk.BooleanVar(value=True)
                ttk.Checkbutton(toggle_frame, text="Show EMF", variable=self.show_emf_var).pack(side='left')
                
                # Status
                status_frame = ttk.LabelFrame(root, text="Status", padding=5)
                status_frame.pack(fill='x', padx=5, pady=5)
                
                self.status_var = tk.StringVar(value="Initializing...")
                ttk.Label(status_frame, textvariable=self.status_var).pack()
                
                # Update loop
                def update():
                    try:
                        for k, v in self.gui_vars.items():
                            self.params[k] = v.get()
                        self.params['show_room'] = self.show_room_var.get()
                        self.params['show_emf'] = self.show_emf_var.get()
                        
                        networks = len(self.emf_scanner['networks'])
                        furniture = len(self.room['furniture'])
                        fixtures = len(self.room['fixtures'])
                        self.status_var.set(f"Networks: {networks} | Furniture: {furniture} | Fixtures: {fixtures}")
                    except: pass
                    root.after(100, update)
                update()
                
                root.mainloop()
            except Exception as e:
                print(f"[GUI] Error: {e}")
        
        threading.Thread(target=setup, daemon=True).start()
    
    def _add_furniture(self):
        """Add furniture at current camera focus"""
        # Simple box furniture at origin for now
        furniture = np.array([
            [-0.5, -0.3, 0], [0.5, -0.3, 0], [0.5, 0.3, 0], [-0.5, 0.3, 0],  # base
            [-0.5, -0.3, 0.8], [0.5, -0.3, 0.8], [0.5, 0.3, 0.8], [-0.5, 0.3, 0.8]  # top
        ], dtype=np.float32)
        
        # Offset by some randomness so they don't overlap
        offset = np.random.uniform(-1, 1, 3)
        offset[2] = 0  # Keep on floor
        furniture += offset
        
        self.room['furniture'].append(furniture)
        print(f"[Room] Added furniture #{len(self.room['furniture'])}")
    
    def _add_fixture(self):
        """Add electrical fixture (light, outlet, etc.)"""
        # Random wall position
        wall_pos = np.array([
            np.random.choice([-2, 2]),  # On wall
            np.random.uniform(-1.5, 1.5),
            np.random.uniform(0.5, 2.0)
        ], dtype=np.float32)
        
        self.room['fixtures'].append(wall_pos)
        print(f"[Room] Added fixture #{len(self.room['fixtures'])}")
    
    def init_gl(self):
        """Initialize OpenGL for dual-layer rendering"""
        # Simple vertex/fragment shaders
        vs = """
        #version 120
        attribute vec3 position;
        attribute vec3 color;
        varying vec3 v_color;
        uniform float u_alpha;
        
        void main() {
            gl_Position = gl_ModelViewProjectionMatrix * vec4(position, 1.0);
            gl_PointSize = 8.0;
            v_color = color;
        }
        """
        
        fs = """
        #version 120
        varying vec3 v_color;
        uniform float u_alpha;
        
        void main() {
            gl_FragColor = vec4(v_color, u_alpha);
        }
        """
        
        try:
            self.shader = compileProgram(compileShader(vs, GL_VERTEX_SHADER), compileShader(fs, GL_FRAGMENT_SHADER))
            print("[GL] Shaders ready")
        except Exception as e:
            print(f"[GL] Shader error: {e}")
            self.shader = None
        
        # OpenGL state
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glClearColor(0.02, 0.02, 0.08, 1.0)
        
        # Initialize sub-renderers
        self.static_layer.init()
        self.emf_layer.init()
    
    def update(self):
        """Update both layers"""
        self.time += 0.02
        
        # Static layer: Room structure (slow updates)
        self.static_layer.update(self.room)
        
        # Dynamic layer: EMF visualization (fast updates)
        self.emf_layer.update(self.emf_scanner, self.params, self.time)
    
    def render(self):
        """Render both layers"""
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        # Set 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)
        
        # Render static room first (background)
        if self.params['show_room']:
            self.static_layer.render()
        
        # Render dynamic EMF overlay
        if self.params['show_emf']:
            self.emf_layer.render(self.params)
        
        glutSwapBuffers()
    
    # Input handlers (same as before)
    def keyboard(self, key, x, y):
        if key == b'r': self.camera = {'x': 45, 'y': 20, 'dist': 8, 'orbiting': False, 'last_pos': None}
        elif key == b'f': self._add_furniture()
        elif key == b'x': self._add_fixture()
        elif key == b's': self._save_room()
        elif key in b'+=': self.camera['dist'] = max(2, self.camera['dist'] * 0.9)
        elif key == b'-': self.camera['dist'] = min(20, self.camera['dist'] * 1.1)
        elif key == b'\x1b': 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
    
    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.3
            self.camera['y'] = np.clip(self.camera['y'] + dy * 0.3, -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), 2, 20)
    
    def idle(self):
        self.update()
        glutPostRedisplay()

class StaticRoomRenderer:
    """Renders persistent room structure"""
    def init(self):
        self.room = None
    
    def update(self, room):
        self.room = room
    
    def render(self):
        if not self.room:
            return
            
        # Draw room wireframe
        glColor3f(0.3, 0.3, 0.4)
        glLineWidth(2.0)
        
        # Floor
        glBegin(GL_LINE_LOOP)
        for vertex in self.room['floor']:
            glVertex3fv(vertex)
        glEnd()
        
        # Ceiling
        glBegin(GL_LINE_LOOP)
        for vertex in self.room['ceiling']:
            glVertex3fv(vertex)
        glEnd()
        
        # Walls (vertical lines)
        glBegin(GL_LINES)
        for i in range(4):
            glVertex3fv(self.room['floor'][i])
            glVertex3fv(self.room['ceiling'][i])
        glEnd()
        
        # Furniture (solid)
        glColor3f(0.6, 0.4, 0.2)
        for furniture in self.room['furniture']:
            glBegin(GL_QUADS)
            # Draw simple box (just top face for now)
            for i in [4, 5, 6, 7]:  # Top face indices
                glVertex3fv(furniture[i])
            glEnd()
        
        # Fixtures (points)
        glColor3f(1.0, 1.0, 0.3)
        glPointSize(12.0)
        glBegin(GL_POINTS)
        for fixture in self.room['fixtures']:
            glVertex3fv(fixture)
        glEnd()

class EMFRenderer:
    """Renders dynamic EMF overlay"""
    def init(self):
        self.emf_grid = {}  # 3D grid of EMF intensities
        self.particles = np.random.uniform(-2, 2, (500, 3)).astype(np.float32)
        self.colors = np.zeros((500, 3), dtype=np.float32)
        self.active_count = 0
    
    def update(self, emf_scanner, params, time):
        # Update EMF particles based on network data
        active_networks = list(emf_scanner['networks'].items())[:500]
        
        for i, (ssid, data) in enumerate(active_networks):
            if i >= len(self.particles):
                break
            
            # Signal strength affects position and color
            signal = data['signal']
            strength = np.clip((signal + 100) / 50.0, 0, 1)
            
            # Move particles based on frequency
            freq_factor = 1.0 if data.get('freq', 2400) < 3000 else 1.5
            
            # Orbital motion based on signal strength
            radius = 1.0 + strength * 2.0
            angle = time * freq_factor + i * 0.5
            
            self.particles[i] = [
                radius * np.cos(angle) * params['emf_intensity'],
                radius * np.sin(angle) * params['emf_intensity'],
                strength * 2.0  # Height based on signal strength
            ]
            
            # Color: red for weak, green for strong
            self.colors[i] = [1.0 - strength, strength, 0.3 + strength * 0.4]
        
        self.active_count = len(active_networks)
    
    def render(self, params):
        # Render EMF particles
        glPointSize(6.0)
        glBegin(GL_POINTS)
        for i in range(min(self.active_count, len(self.particles))):
            glColor3fv(self.colors[i])
            glVertex3fv(self.particles[i])
        glEnd()

def main():
    print("Room Mapper with EMF Overlay")
    print("Controls: F=add furniture, X=add fixture, S=save room")
    print("Mouse=orbit, +/-=zoom, R=reset, ESC=exit")
    
    mapper = RoomMapper()
    
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(1400, 800)
    glutCreateWindow(b"Room Mapper - Static + EMF")
    
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(50.0, 1400.0/800.0, 0.1, 100.0)
    
    mapper.init_gl()
    
    glutDisplayFunc(mapper.render)
    glutIdleFunc(mapper.idle)
    glutKeyboardFunc(mapper.keyboard)
    glutMouseFunc(mapper.mouse)
    glutMotionFunc(mapper.motion)
    try: glutMouseWheelFunc(mapper.wheel)
    except: pass
    
    glutMainLoop()

if __name__ == "__main__":
    main()