import tkinter as tk
import numpy as np
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import pywifi
import time

# --------------------------
# Parameters
# --------------------------
N = 200            # Points per network
dt = 0.05          # Animation interval (s)
k_envelope = 0.5   # Softening factor for tanh envelope

# --------------------------
# Tkinter setup
# --------------------------
root = tk.Tk()
root.title("Wi-Fi EMF Organic Lattice Visualizer")

# --------------------------
# Matplotlib figure
# --------------------------
fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(111, projection='3d')
canvas = FigureCanvasTkAgg(fig, master=root)
canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)

# --------------------------
# Slider frame
# --------------------------
slider_frame = tk.Frame(root)
slider_frame.pack(side=tk.BOTTOM, fill=tk.X)

def make_slider(label, minv, maxv, default, row):
    tk.Label(slider_frame, text=label).grid(row=row, column=0, sticky='w')
    var = tk.DoubleVar(value=default)
    slider = tk.Scale(slider_frame, from_=minv, to=maxv, resolution=0.01,
                      orient=tk.HORIZONTAL, variable=var)
    slider.grid(row=row, column=1, sticky='we')
    return var

morph_var = make_slider("Morph (Polar→Cartesian)", 0, 1, 0, 0)
ampl_var  = make_slider("Amplitude Scale", 0, 2, 1, 1)

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

def scan_networks():
    iface.scan()
    time.sleep(1)  # small delay for scan results
    return iface.scan_results()

# --------------------------
# Lattice function (phyllotaxis)
# --------------------------
phi = (1 + np.sqrt(5)) / 2
theta = 2 * np.pi / phi
radii = np.sqrt(np.arange(N))
angles = np.arange(N) * theta
zs = np.linspace(-1,1,N)

def get_lattice(morph, ampl, signal_mod):
    r = radii * ampl + signal_mod
    x_polar = r * np.cos(angles)
    y_polar = r * np.sin(angles)
    z_polar = zs * ampl

    # Morph between polar and cartesian
    x = x_polar * (1-morph) + np.linspace(-1,1,N) * morph
    y = y_polar * (1-morph) + np.linspace(-1,1,N) * morph
    z = z_polar * (1-morph) + np.linspace(-1,1,N) * morph
    return x, y, z

# --------------------------
# Animation setup
# --------------------------
scatters = {}      # dict: ssid -> scatter object
colors = plt.colormaps['hsv']  # Matplotlib 3.11+ for unique colors

def update(frame):
    networks = scan_networks()
    active_ssids = [net.ssid for net in networks if net.ssid.strip() != '']
    max_networks = max(len(active_ssids), 1)

    # Create new scatter objects for new networks
    for idx, ssid in enumerate(active_ssids):
        if ssid not in scatters:
            x, y, z = get_lattice(morph_var.get(), ampl_var.get(), 0)
            scatters[ssid] = ax.scatter(x, y, z,
                                        s=20,
                                        c=[colors(idx/max_networks)]*N,
                                        alpha=0.8)

    # Update each scatter
    for idx, ssid in enumerate(active_ssids):
        signal_strength = next((net.signal for net in networks if net.ssid == ssid), -100)
        # Normalize signal and apply soft envelope
        normalized = np.tanh(k_envelope * (signal_strength+100)/100)
        x, y, z = get_lattice(morph_var.get(), ampl_var.get(), normalized)
        scatters[ssid]._offsets3d = (x, y, z)
        scatters[ssid].set_array(z)  # For color mapping if needed

    # Remove any disconnected networks
    for ssid in list(scatters.keys()):
        if ssid not in active_ssids:
            scatters[ssid].remove()
            del scatters[ssid]

    ax.set_xlim(-5,5)
    ax.set_ylim(-5,5)
    ax.set_zlim(-2,2)
    return list(scatters.values())

ani = FuncAnimation(fig, update, interval=dt*1000, blit=False)

# --------------------------
# Start Tkinter
# --------------------------
root.mainloop()
