import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

# --- PARAMETERS ---
N = 150               # number of lattice points
dt = 0.05             # animation timestep
morph_speed = 0.002   # morphing speed 0-1
phi = (1 + np.sqrt(5)) / 2

# Environmental carrier signals (example: FM/AM)
env_freqs = [0.2, 0.5, 1.0]  
env_amps = [0.5, 0.3, 0.4]

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

# Initial Cartesian coordinates
X = r * np.cos(theta)
Y = r * np.sin(theta)
Z = np.zeros_like(X)

# --- FIGURE ---
fig = plt.figure(figsize=(10,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(-10,10)
ax.set_box_aspect([1,1,0.6])

# Morph variable
morph = 0.0
morph_direction = 1

# --- UPDATE FUNCTION ---
def update(frame):
    global morph, morph_direction
    
    t = frame*dt
    morph += morph_direction * morph_speed
    if morph >= 1.0: morph_direction = -1
    if morph <= 0.0: morph_direction = 1

    # AM/FM environmental modulations
    r_mod = r + env_amps[0]*np.sin(2*np.pi*env_freqs[0]*t + r)
    theta_mod = theta + env_amps[1]*np.sin(2*np.pi*env_freqs[1]*t + theta)
    z_mod = env_amps[2]*np.sin(2*np.pi*env_freqs[2]*t + r)

    # Polar coordinates
    Xp = r_mod * np.cos(theta_mod)
    Yp = r_mod * np.sin(theta_mod)
    Zp = z_mod

    # Cartesian coordinates (simply stretch along axes for demo)
    Xc = r_mod * np.cos(theta_mod) + 0.3*Xp
    Yc = r_mod * np.sin(theta_mod) + 0.3*Yp
    Zc = Zp + 0.2*Yp

    # Morph between polar and cartesian
    Xf = (1-morph)*Xp + morph*Xc
    Yf = (1-morph)*Yp + morph*Yc
    Zf = (1-morph)*Zp + morph*Zc

    # Size and color modulation (analog feel)
    sizes = 30 + 20*np.sin(2*np.pi*env_freqs[0]*t + r)
    colors = Zf

    # Update scatter
    points._offsets3d = (Xf.flatten(), Yf.flatten(), Zf.flatten())
    points.set_sizes(sizes)
    points.set_array(colors)

    return points,

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