import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import pywifi
import time
import random
import tkinter as tk
from tkinter import ttk

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

def scan_networks():
    iface.scan()
    time.sleep(0.5)
    return [(res.ssid, res.signal) for res in iface.scan_results()]

# ---------------- Parameters ----------------
dt = 0.05
n_particles = 500
radius_speed = 0.02
stable_threshold = 3.0
stable_frames = 30
persistence = 0.92   # fade factor for trails

# ---------------- State ----------------
particle_positions = np.random.rand(n_particles, 3) * 2 - 1
particle_colors = np.random.rand(n_particles, 3)
radii = np.zeros(n_particles)

scaffold_points = []
scaffold_colors = []
stability_tracker = {}
show_scaffold = True

# Persistent trail buffer
trail_positions = np.zeros((n_particles, 3))
trail_colors = np.zeros((n_particles, 3))

# ---------------- GUI Sliders ----------------
root = tk.Tk()
root.title("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)

# ---------------- Matplotlib Setup ----------------
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.set_facecolor("black")
fig.patch.set_facecolor("black")

scat_cloud = ax.scatter([], [], [], s=6, c=[], alpha=0.7)
scat_trails = ax.scatter([], [], [], s=4, c=[], alpha=0.3)
scat_scaffold = ax.scatter([], [], [], s=20, c=[], alpha=0.9, marker='^')

ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
ax.set_zlim(-1, 1)
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])

# ---------------- Toggle ----------------
def on_key(event):
    global show_scaffold
    if event.key == 's':
        show_scaffold = not show_scaffold
fig.canvas.mpl_connect('key_press_event', on_key)

# ---------------- Update ----------------
def update(frame):
    global radii, particle_positions, particle_colors
    global trail_positions, trail_colors
    global scaffold_points, scaffold_colors

    networks = scan_networks()
    amp = amp_var.get()
    noise = noise_var.get()
    morph = morph_var.get()

    # Sonar radial sweep
    radii += radius_speed
    radii = np.where(radii > 1.5, 0, radii)

    theta = np.linspace(0, 2*np.pi, n_particles)
    phi = np.linspace(0, np.pi, n_particles)
    np.random.shuffle(theta)
    np.random.shuffle(phi)

    # Polar coords
    x_polar = radii * np.sin(phi) * np.cos(theta)
    y_polar = radii * np.sin(phi) * np.sin(theta)
    z_polar = radii * np.cos(phi)

    # Cartesian coords
    x_cart = np.linspace(-1, 1, n_particles)
    y_cart = np.linspace(-1, 1, n_particles)
    z_cart = np.linspace(-1, 1, n_particles)

    # Morph between polar & cartesian
    particle_positions[:, 0] = (1-morph) * x_polar + morph * x_cart
    particle_positions[:, 1] = (1-morph) * y_polar + morph * y_cart
    particle_positions[:, 2] = (1-morph) * z_polar + morph * z_cart

    # Noise
    particle_positions += np.random.normal(0, noise, particle_positions.shape)

    # Color by RSSI
    if networks:
        for i, (ssid, signal) in enumerate(networks[:n_particles]):
            strength = max(-100, signal) / -100.0
            particle_colors[i] = [1-strength, strength, random.uniform(0, 0.5)]

            # Stability check
            if ssid not in stability_tracker:
                stability_tracker[ssid] = []
            stability_tracker[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)
                    pos = particle_positions[i] * (0.5 + mean_signal/-100.0)
                    scaffold_points.append(pos)
                    scaffold_colors.append([0.8, 0.8, 1.0])
                    stability_tracker[ssid] = []

    # --- Trails persistence ---
    trail_positions = persistence * trail_positions + (1-persistence) * particle_positions
    trail_colors = persistence * trail_colors + (1-persistence) * particle_colors

    # Update cloud
    scat_cloud._offsets3d = (particle_positions[:, 0]*amp,
                             particle_positions[:, 1]*amp,
                             particle_positions[:, 2]*amp)
    scat_cloud.set_color(particle_colors)

    # Update trails
    scat_trails._offsets3d = (trail_positions[:, 0]*amp,
                              trail_positions[:, 1]*amp,
                              trail_positions[:, 2]*amp)
    scat_trails.set_color(trail_colors)

    # Update scaffold if visible
    if show_scaffold and scaffold_points:
        pts = np.array(scaffold_points)
        scat_scaffold._offsets3d = (pts[:, 0], pts[:, 1], pts[:, 2])
        scat_scaffold.set_color(scaffold_colors)
        scat_scaffold.set_alpha(0.9)
    else:
        scat_scaffold.set_alpha(0.0)

    return scat_cloud, scat_trails, scat_scaffold

# ---------------- Run ----------------
ani = FuncAnimation(fig, update, interval=dt*1000, blit=False)
plt.show()
root.mainloop()
