import math
import numpy as np
import matplotlib.pyplot as plt
import sounddevice as sd
import time

sample_rate = 44100
seconds = 10

N = 100 # number of string segments
I = int(sample_rate * seconds) # number of samples to simulate

# string parameters
rho = 8000 # density, steel, kg/m^3
radius = 0.001 # meters
S = math.pi*radius**2 # string cross sectional area, assuming circular
mu = S*rho # linear mass density
T = 1200 # string tension, N
c = math.sqrt(T/mu) # transverse wave velocity
kappa = 0.001 # stiffness coefficient
sigma = 0.5 # damping coefficient
L = 1.0 # length of string
strike_position = 0.2 # x of impulse location
impulse_width = 0.02 # x of impulse width
impulse_velocity = 1000.0 # x/t of impulse magnitude
sample_position = 0.1 # percentage along L of sampling for audio

f_0 = c / (2*L) # fundamental frequency of a non-stiff string
f_1 = f_0 * math.sqrt(1 + kappa) # fundamental frequency of the stiff string
print("fundamental frequency =", f_1)

dx = L / N # delta x
dt = 1/sample_rate
#dt = 0.2 * dx / c
if(dx**2 < (c*dt)**2 + 4*(kappa*dt/(dx**2))**2):
    print("warning: possibly unstable, increase sample rate or decrease string segments")

# derived constants
r1 = c * dt/dx 
r2 = (c * dt/dx) ** 2
s1 = kappa * dt/dx**2
s2 = (kappa * dt/dx**2) ** 2
a1 = 2 - 2*sigma*dt
a2 = 2*sigma*dt - 1

# string grid
x = np.linspace(0, L, N + 1) # linspace my beloved

# state vectors
y_last = np.zeros(N + 1)
y_current = np.zeros(N + 1)
y_next = np.zeros(N + 1)
# output
y_sample = np.zeros(I)

# vector of initial velocity across the string as a result of the impulse
v0 = impulse_velocity * np.exp(-((x - strike_position) ** 2)/(2 * impulse_width ** 2))

# initial conditions:
# y(x, 0) = 0 and dy/dt(x, 0) = v0(x)
# boundary conditions: y(0, t) = y(L, t) = 0 as well as all dy/dts thereafter

def show_plot():
    plt.plot(np.arange(0, N+1, 1), y_current)
    plt.grid()
    plt.show()

start_time = time.perf_counter()

def applyImpulse():
    for n in range(2, N-2):
        y_current[n] = y_current[n] + dt*v0[n]

# first iteration
for n in range(2, N-2):
    y_current[n] = 0.5 * r1**2 * (y_last[n-1] - 2*y_last[n] + y_last[n+1])
applyImpulse()

n_sample = int(sample_position * L * N)
y_sample[0] = y_last[n_sample]
y_sample[1] = y_current[n_sample]

# rest of the iterations
for i in range(2, I):

    for n in range(2, N-2):
        # stiff wave equation
        y_xx = y_current[n+1] - 2*y_current[n] + y_current[n-1]
        y_xxxx = y_current[n-2] - 4*y_current[n-1] + 6*y_current[n] - 4*y_current[n+1] + y_current[n+2]

        y_next[n] = a1 * y_current[n] + a2 * y_last[n] + r2 * y_xx - s2 * y_xxxx

    y_next[0] = 0
    y_next[1] = 0
    y_next[N-1] = 0
    y_next[N-2] = 0

    # y_sample[i] = math.tanh(y_next[n_sample])
    y_sample[i] = y_next[n_sample]

    y_last = y_current.copy()
    y_current = y_next.copy()

    if(i == 120000):
        applyImpulse()

    if(i % 10000*seconds == 0):
        print(f"{i/I * 100:4.4}% complete")

end_time = time.perf_counter()
elapsed = end_time - start_time
print(f"Executed in {elapsed:.3f} seconds. {elapsed/seconds*100:.2f}% overshoot")

plt.plot(np.arange(0, I, 1), y_sample)
plt.grid()
plt.show()

sd.play(y_sample, sample_rate)
sd.wait()