import tkinter as tk
from tkinter import ttk
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
from collections import deque
import pywifi
from pywifi import PyWiFi, const

# --------------------------
# Parameters
# --------------------------
N_POINTS = 400          # Points per network
dt = 0.05               # Animation interval (s)
soft_k = 0.5            # Envelope scaling
vibration_amp = 0.05    # Analog jitter amplitude
max_networks = 10       # Max SSIDs to track

# --------------------------
# Tkinter setup
# --------------------------
root = tk.Tk()
root.title("HDGL Wi-Fi Analog 3D Lattice")

# --------------------------
# 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)

# --------------------------
# Lattice setup
# --------------------------
phi = (1 + np.sqrt(5)) / 2  # Golden ratio
theta = 2 * np.pi / phi
radii = np.sqrt(np.arange(N_POINTS))
angles = np.arange(N_POINTS) * theta
zs = np.linspace(-1, 1, N_POINTS)

def lattice_coords(N, t, morph=0):
    r = radii
    x_polar = r * np.cos(angles)
    y_polar = r * np.sin(angles)
    z_polar = zs * ampl_var.get()

    x_cart = x_polar * (1-morph) + np.linspace(-1,1,N) * morph
    y_cart = y_polar * (1-morph) + np.linspace(-1,1,N) * morph
    z_cart = z_polar * (1-morph) + np.linspace(-1,1,N) * morph

    return x_cart, y_cart, z_cart

# --------------------------
# Soft envelope filter
# --------------------------
def soft_env(s):
    return np.tanh(soft_k * s)

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

def scan_networks():
    iface.scan()
    import time
    time.sleep(0.2)
    results = iface.scan_results()
    networks = {}
    for r in results[:max_networks]:
        networks[r.ssid] = r.signal
    return networks

# --------------------------
# Network buffers & colors
# --------------------------
network_buffers = {}  # ssid -> deque of points
dots = {}             # ssid -> scatter
colors = plt.cm.get_cmap('hsv', max_networks)

# --------------------------
# Animation update
# --------------------------
def update(frame):
    morph = morph_var.get()
    networks = scan_networks()
    x_base, y_base, z_base = lattice_coords(N_POINTS, frame*dt, morph)

    for i, (ssid, rssi) in enumerate(networks.items()):
        amp_mod = ampl_var.get() * soft_env((rssi+100)/60)
        idx = int(i * N_POINTS / max_networks)
        x = x_base[idx] * amp_mod + np.random.uniform(-vibration_amp,vibration_amp,N_POINTS)
        y = y_base[idx] * amp_mod + np.random.uniform(-vibration_amp,vibration_amp,N_POINTS)
        z = z_base[idx] * amp_mod + np.random.uniform(-vibration_amp,vibration_amp,N_POINTS)

        if ssid not in network_buffers:
            network_buffers[ssid] = deque(maxlen=N_POINTS)
        network_buffers[ssid].append(np.vstack([x,y,z]).T)

        buf = np.array(network_buffers[ssid])
        x_plot = buf[:,0].flatten()
        y_plot = buf[:,1].flatten()
        z_plot = buf[:,2].flatten()

        if ssid not in dots:
            dots[ssid] = ax.scatter(x_plot, y_plot, z_plot, color=colors(i), s=15)
        else:
            dots[ssid]._offsets3d = (x_plot, y_plot, z_plot)

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

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

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