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
import pywifi
from pywifi import PyWiFi
import time

# --------------------------
# Parameters
# --------------------------
N = 800           # Number of particles per network
dt = 0.05         # Animation interval in seconds
particle_size = 5 # Size of dots
max_networks = 6  # Max number of Wi-Fi networks to track

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

# --------------------------
# Matplotlib figure
# --------------------------
fig = plt.figure(figsize=(9,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):
    tk.Label(slider_frame, text=label).pack(side=tk.LEFT)
    var = tk.DoubleVar(value=default)
    slider = tk.Scale(slider_frame, from_=minv, to=maxv, resolution=0.01,
                      orient=tk.HORIZONTAL, variable=var)
    slider.pack(side=tk.LEFT, padx=5)
    return var

morph_var = make_slider("Morph (Polar↔Cartesian)", 0, 1, 0)
ampl_var  = make_slider("Amplitude Scale", 0, 2, 1)
size_var  = make_slider("Particle Size", 1, 10, particle_size)

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

def scan_wifi():
    iface.scan()
    time.sleep(0.5)
    return iface.scan_results()

# --------------------------
# Lattice helpers
# --------------------------
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_coords(t, amplitude, morph):
    """Generate jittery particle positions for analog feel"""
    r = radii * (1 + 0.2*np.sin(2*np.pi*t*0.1))
    x_polar = r * np.cos(angles)
    y_polar = r * np.sin(angles)
    z_polar = zs * amplitude

    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

    # Add small random jitter for analog effect
    jitter = 0.05
    x_cart += np.random.uniform(-jitter,jitter,N)
    y_cart += np.random.uniform(-jitter,jitter,N)
    z_cart += np.random.uniform(-jitter,jitter,N)
    return x_cart, y_cart, z_cart

# --------------------------
# Initialize scatter plots
# --------------------------
colors = plt.cm.get_cmap('hsv', max_networks)
scatters = []

for i in range(max_networks):
    scat = ax.scatter([],[],[], s=size_var.get(), color=colors(i), alpha=0.7)
    scatters.append(scat)

ax.set_xlim(-5,5)
ax.set_ylim(-5,5)
ax.set_zlim(-2,2)

# --------------------------
# Animation update
# --------------------------
def update(frame):
    morph = morph_var.get()
    ampl = ampl_var.get()
    size = size_var.get()

    networks = scan_wifi()[:max_networks]  # Only track first few networks
    for i, net in enumerate(networks):
        x, y, z = get_coords(frame*dt, ampl, morph)
        scatters[i]._offsets3d = (x, y, z)
        scatters[i].set_sizes([size]*N)

    return scatters

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

# --------------------------
# Run
# --------------------------
root.mainloop()
