import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D  # needed for 3D plotting

# --- PARAMETERS ---
N = 100          # number of lattice points
dt = 0.05        # animation timestep
morph = 0.0      # 0=polar, 1=cartesian

# Phyllotaxis constants
phi = (1 + np.sqrt(5)) / 2

# Carrier signals (FM/AM)
env_freqs = [0.5, 1.0, 2.0]  # example environmental frequencies

# --- LATTICE ---
theta = np.array([n * 2*np.pi/phi for n in range(N)])
r = np.sqrt(np.arange(N))

# Convert to Cartesian initially
X = r * np.cos(theta)
Y = r * np.sin(theta)
Z = np.zeros_like(X)

# --- SETUP FIGURE ---
fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection='3d')
points = ax.scatter(X, Y, Z, c=Z, cmap='viridis', s=40)
ax.set_xlim(-15,15)
ax.set_ylim(-15,15)
ax.set_zlim(-5,5)

# --- UPDATE FUNCTION ---
def update(frame):
    global morph

    # Phyllotaxis lattice evolution
    t = frame*dt
    r_mod = r + 0.5*np.sin(2*np.pi*env_freqs[0]*t)  # AM-like radial modulation
    theta_mod = theta + 0.5*np.sin(2*np.pi*env_freqs[1]*t)  # FM-like angular modulation

    # Morph between polar and Cartesian
    Xp = r_mod * np.cos(theta_mod)
    Yp = r_mod * np.sin(theta_mod)
    Zp = 0.5*np.sin(2*np.pi*env_freqs[2]*t + r_mod)  # vertical modulation

    Xf = (1-morph)*Xp + morph*Xp       # placeholder if you want Cartesian morph later
    Yf = (1-morph)*Yp + morph*Yp
    Zf = Zp

    # Flatten for scatter
    points._offsets3d = (Xf.flatten(), Yf.flatten(), Zf.flatten())
    points.set_array(Zf.flatten())

    return points,

# --- ANIMATION ---
ani = FuncAnimation(fig, update, interval=dt*1000, blit=False)
plt.show()
