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 sounddevice as sd
import socket
import json

# --------------------------
# Parameters
# --------------------------
N = 500            # Number of points in lattice
dt = 0.05          # Animation interval (s)
audio_scale = 0.01 # Scaling for audio modulation
env_k = 2.0        # Soft envelope scaling for EMF/audio combination

# --------------------------
# Tkinter setup
# --------------------------
root = tk.Tk()
root.title("HDGL Multi-band 3D Lattice (Audio + Wi-Fi EMF)")

# --------------------------
# 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):
    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)
    return var

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

# --------------------------
# Lattice setup (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 soft_envelope(signal, k=2.0):
    return np.tanh(k * signal)

def get_lattice(t=0, audio_mod=0, env_mod=0):
    """Return 3D coordinates of lattice, with audio + EMF modulation"""
    combined_mod = soft_envelope(audio_mod + env_mod, k=env_k)
    r = radii + combined_mod * audio_scale
    x_polar = r * np.cos(angles)
    y_polar = r * np.sin(angles)
    z_polar = zs * ampl_var.get()

    morph = morph_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

# --------------------------
# Audio input callback
# --------------------------
audio_buffer = np.zeros(1024)
def audio_callback(indata, frames, time, status):
    global audio_buffer
    audio_buffer = indata[:,0]

stream = sd.InputStream(callback=audio_callback, channels=1, samplerate=44100)
stream.start()

# --------------------------
# Wi-Fi EMF socket setup
# --------------------------
UDP_IP = "0.0.0.0"
UDP_PORT = 5005
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind((UDP_IP, UDP_PORT))
sock.setblocking(False)

def get_wifi_emf():
    """Non-blocking read of UDP EMF data from Wi-Fi sensor"""
    try:
        data, _ = sock.recvfrom(1024)
        payload = json.loads(data.decode())
        return float(payload.get("emf", 0.0))
    except:
        return 0.0

# --------------------------
# Animation setup
# --------------------------
scat = ax.scatter([], [], [], c=[], cmap='viridis', s=20)

def update(frame):
    audio_mod = np.abs(audio_buffer).mean()
    env_mod = get_wifi_emf()
    x, y, z = get_lattice(t=frame*dt, audio_mod=audio_mod, env_mod=env_mod)
    scat._offsets3d = (x, y, z)
    scat.set_array(z)
    ax.set_xlim(-5,5)
    ax.set_ylim(-5,5)
    ax.set_zlim(-2,2)
    return scat,

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

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