import tkinter as tk
from tkinter import ttk
import numpy as np
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import sounddevice as sd

# --------------------------
# Parameters
# --------------------------
N = 500            # Number of points in lattice
dt = 0.05          # Animation interval (s)
audio_scale = 0.01 # Scaling for audio modulation
carrier_scale = 0.05  # Environmental signal scale
fm_freq = 0.5      # Frequency for FM modulation
am_freq = 0.2      # Frequency for AM modulation

# --------------------------
# Tkinter setup
# --------------------------
root = tk.Tk()
root.title("HDGL Multi-band 3D Lattice with Audio + Analog Carriers")

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

# --------------------------
# Matplotlib toolbar
# --------------------------
toolbar = NavigationToolbar2Tk(canvas, root)
toolbar.update()
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)
carrier_var = make_slider("Carrier Strength", 0, 0.2, 0.05, 2)
fm_var = make_slider("FM Frequency", 0, 2, fm_freq, 3)
am_var = make_slider("AM Frequency", 0, 2, am_freq, 4)

# --------------------------
# Lattice setup (phyllotaxis)
# --------------------------
phi = (1 + np.sqrt(5)) / 2  # Golden ratio
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(t=0, audio_mod=0):
    """Return 3D coordinates of lattice, with audio modulation, morphing, and carriers"""
    r = radii + audio_mod * audio_scale
    
    # Polar coordinates
    x_polar = r * np.cos(angles)
    y_polar = r * np.sin(angles)
    z_polar = zs * ampl_var.get()
    
    # Analog carriers (FM + AM)
    carrier = carrier_var.get()
    x_carrier = x_polar + carrier * np.sin(2*np.pi*fm_var.get()*t + angles)
    y_carrier = y_polar + carrier * np.sin(2*np.pi*fm_var.get()*t + angles)
    z_carrier = z_polar + carrier * np.sin(2*np.pi*am_var.get()*t) * np.cos(angles)

    # Morph between polar and cartesian
    morph = morph_var.get()
    x_final = x_carrier * (1-morph) + np.linspace(-1,1,N) * morph
    y_final = y_carrier * (1-morph) + np.linspace(-1,1,N) * morph
    z_final = z_carrier * (1-morph) + np.linspace(-1,1,N) * morph

    return x_final, y_final, z_final

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

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

def update(frame):
    audio_mod = np.abs(audio_buffer).mean()
    x, y, z = get_lattice(t=frame*dt, audio_mod=audio_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()
