import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

class DecayingSineWave:
    def __init__(self, freq=5, decay=0.05, sample_rate=60):
        self.freq = freq
        self.decay = decay
        self.sample_rate = sample_rate
        self.triggers = []

    def trigger(self, t):
        self.triggers.append(t)

    def sample(self, t):
        value = 0.0
        still_active = []
        for start_time in self.triggers:
            age = t - start_time
            if age >= 0:
                v = np.sin(2 * np.pi * self.freq * age / self.sample_rate) * np.exp(-self.decay * age)
                value += v
                if np.exp(-self.decay * age) > 1e-3:
                    still_active.append(start_time)
        self.triggers = still_active
        return value

# --- Initialize ---
wave = DecayingSineWave(freq=5, decay=0.05, sample_rate=60)
wave_array = np.zeros(512)
time = [0]
max_len = 512

fig, ax = plt.subplots()
line, = ax.plot(np.arange(512), wave_array, lw=2)
trig_dots, = ax.plot([], [], 'ro', markersize=4)

ax.set_xlim(0, 511)
ax.set_ylim(-1.2, 1.2)
ax.set_title("Click to Trigger Decaying Sine Wave")
ax.set_xlabel("Sample Index (0 = current)")
ax.set_ylabel("Amplitude")
ax.grid(True)

trigger_times = []

# --- Click handler ---
def on_click(event):
    current_time = time[0]
    wave.trigger(current_time)
    trigger_times.append(current_time)

fig.canvas.mpl_connect('button_press_event', on_click)

# --- Animation update ---
def update(frame):
    global wave_array
    current_time = time[0]

    # Shift buffer to the right (older samples move toward the end)
    wave_array = wave_array * 0.995
    wave_array = np.roll(wave_array, 1)
    # Insert new sample at index 0
    wave_array[0] = wave.sample(current_time)
    print (wave_array)

    line.set_data(np.arange(512), wave_array)

    # Trigger markers
    visible_triggers = [tt for tt in trigger_times if current_time - 512 < tt <= current_time]
    x = [current_time - tt for tt in visible_triggers]  # 0 = current time
    y = [1.0 for _ in x]
    trig_dots.set_data(x, y)

    time[0] += 1
    return line, trig_dots

ani = FuncAnimation(fig, update, interval=1000 / 60, blit=True)
plt.show()