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

# https://people.cs.uchicago.edu/~ridg/stabil/pianostring.pdf
# a differential equations model for a piano string

# eq 1: governing wave equation
# d2y/dt2 = c^2*d2y/dx2 - stiffness*c^2*L^2*d4y/dx4 - 2*b_1*dy/dt + 2*b_3*d3y/dt3 + f(x, x_0, t)
# where y = string's transverse displacement, b_1, b_3 = damping coefficients, f = force density
# c = sqrt(T/mu), transverse wave velocity , T = string tension, mu = string linear mass density)

# eq 2: string stiffness
# stiffness = K^2*(E*S/(T*L^2))
# where K = string's radius of gyration (r/2 for a circular string), E = string's Young's modulus,
# S = cross sectional string area, T = string tension, L = string length

# eq 3: decay
# sigma = 1/tau = b_1 + b_3*omega^2
# where sigma = decay rate, tau = decay time, omega = angular frequency

# eq 4: excitation
# f(x, x_0, t) = f_H(t) * g(x, x_0)
# where f_H(t) = hammer force, g(x, x_0) = hammer dimensional effect

# eq 5: hammer time history
# read the paper if you care, useful only for derivation

# eq 6: power law
# F_H(t) = K*|eta(t) - y(x_0, t)|^p
# where eta(t) is the transverse displacement of the hammer head
# and p = stiffness nonlinear exponent

# eq 7: hammer displacement
# M_H*d2eta/dt2 = -F_H(t)
# where M_H is the mass of the hamemr head

# eq 8: boundary conditions
# y(0, t) = y(L, t) = 0, fixed ends dont move
# d2y/dx2(0, t) = d2y/dx2(L, t) = 0, displacement along the string approaching the ends is continuous

# discrete time implementation

# eq 9: continuous to discrete
# y(x, t) -> y(x_i, t_n) -> y(i, n)
# divide the string into i segments and iterate over n timesteps
# where x_i = delta_x * i and t_n = delta_t * n

# eq 10: recurrence derivation
# y(i, n+1) = a_1*y(i, n) + a_2*y(i, n-1) + a_3*[y(i+1, n) + y(i-1, n)] + a_4*[y(i+2,n) + y(i-2,n)]
#               + a_5*[y(i+1, n-1) + y(i-1, n-1) + y(i, n-2)]
#               + [delta_t^2 * N*F_H(n) * g(i, i_0)]/M_S
# where a_1 through a_5 are defined as follows:
# a_1 = [2 - 2*r^2 + b_1/delta_t - 6*stiffness*N^2*r^2]/D
# a_2 = [-1 + b_1*delta_t + 2*b_3/delta_t]/D
# a_3 = [r^2*(1 + 4*stiffness*N^2)]/D
# a_4 = [b_3/delta_t - stiffness*N^2*r^2]/D
# a_5 = [-b_3/delta_t]/D
# where D = 1 + b_1*delta_t + 2*b_3/delta_t
# and r = c*delta_t/delta_x

# eq 11: stability condition
# N_max = sqrt{[-1 + sqrt(1+16*stiffness*gamma^2)]/(8*stiffness)}
# where
# eq 12: idk what gamma represents
# gamma = f_e/(2*f_1), f_e = sampling frequency and f_1 = fundamental frequency
# or
# eq 13: if neglecting stiffness
# N_max = gamma

# eq 14: rest condition
# y(i, 0) = 0

# eq 15: discrete hammer displacement
# at t = delta_t (n = 1)
# eta(1) = V_H_0 * delta_t

# eq 16: discrete truncated taylor series
# y(i, 1) = [y(i + 1, 0) + y(i - 1, 0)]/2

# eq 17: hammer force exertion
# F_H(1) = K*|eta(1) - y(i_0, 1)|^p

# eq 18: string displacement iteration
# y(i, 2) = y(i-1, 1)] + y(i + 1, 1) - y(i, 0) + [delta_t^2 * N*F_H(1) * g(i, i_0)]/M_S

# eq 19: hammer displacement iteration
# eta(2) = 2*eta(1) - eta(0) - [delta_t^2 * F_H(1)]/M_H

# eq 20: hammer force iteration
# F_H(2) = K*|eta(2) - y(i_0, 2)|^p

# eq 21: hammer rest condition
# eta(n + 1) < y(i_0, n + 1)

# eq 22: spacial boundary conditions
# y(0, n) = y(N, n) = 0

# eq 23: temporal boundary conditions
# y(-1, n) = -y(1, n)
# y(N + 1, n) = -y(N - 1, n)

# simulation parameters
f1 = 50 # fundamental frequency
f_e = 44100 # sampling frequency
N = 200 # number of string segments
delta_t = 1/f_e # time step
L = 0.5 # string length, meters
delta_x = L/N # spatial step
H = int(f_e * 2) # length of simulation in time

# string parameters
E = 100 * 10**9 # youngs modulus, steel = 200GPa
rho = 8000 # density, steel, kg/m^3
radius = 0.005 # meters
kappa = radius/2 # radius of gyration, r/2 for a circular string
S = math.pi*radius**2 # string cross sectional area, assuming circular
mu = S*rho # linear mass density
M_S = mu*L # string mass
T = 500 # string tension, N
c = math.sqrt(T/mu) # transverse wave velocity
I = math.pi*radius**4 / 4
k = math.sqrt(E*I/(rho*S))
# stiffness = k**2 * delta_t**2 / delta_x**4 # string stiffness parameter
stiffness = 0
sigma = 1 # decay rate
tau = 1/sigma # decay time
omega = 1/f1 # angular frequency

# hammer parameters
M_H = 0.5 # hammer mass, kg
HSMR = M_H/M_S # hammer-mass string ratio
V_H_0 = 10000000000 # initial hammer velocity at t=0, m/s
x_0 = L/2 # distance of hammer from agraffe
alpha = x_0 / L # relative hammer striking position

# empirical constants
b_1 = 0 # some constant
b_3 = 0 # some constant 
K = 1 # hammer stiffness
p = 1 # stiffness nonlinear exponent

# derived components
D = 1 + b_1*delta_t + 2*b_3/delta_t
r = c*delta_t/delta_x
a_1 = (2 - 2*r**2 + b_3/delta_t - 6*stiffness*N**2*r**2)/D
a_2 = (-1 + b_1*delta_t + 2*b_3/delta_t)/D
a_3 = (r**2*(1 + 4*stiffness*N**2))/D
a_4 = (b_3/delta_t - stiffness*N**2*r**2)/D
a_5 = (-b_3/delta_t)/D

# string
x = [0] * N # current string position
x_initial = [0] * N # assuming string at rest
last_x1 = [0] * N # string position from last timestep
last_x2 = [0] * N # string position from two timesteps ago
x_sample = int(0.8*N) # location where we sample the string position for signal

# hammer
F_H_current = 0 # current force exterted by hammer
eta_current = 0 # current hammer position
eta_last = 0 # last hammer position
i_H = 16 # hammer width
x_hammer = int(alpha * N) # location where hammer strikes
eta_0 = 0 # initial hammer displacement

x_out = np.zeros(H) # taking this as the sound output at some arbitrary point along the string
t = np.arange(0, H/f_e, delta_t)

def plot_current():
    plt.plot(np.arange(0, N, 1), x)
    plt.title("Response")
    plt.xlabel("t")
    plt.ylabel("x(t, x_sample)")
    plt.grid()
    plt.show()

def spatial_window(i):
    if(i < x_hammer+i_H/2 and i > x_hammer-i_H/2):
        return abs(i-x_hammer-i_H/2)
    else:
        return 0

# first iteration, n = 1
last_x1 = x_initial.copy()
eta_current = V_H_0 * delta_t + eta_0

for i in range(1, N-1):
    x[i] = (last_x1[i+1] + last_x1[i-1])/2 # irrelavant if string is at rest at hammer strike, but relevant if a repeat hit
last_x1 = x.copy()
F_H_current = K * abs(eta_current - x[x_hammer])**p

# second iteration, n = 2
eta_next = 2*eta_current - eta_0 - (delta_t**2*N*F_H_current)/M_S
eta_last = eta_current
eta_current = eta_next

for i in range(1, N-1):
    x[i] = last_x1[i+1] + last_x1[i-1] - x_initial[i] + ((delta_t**2)*N*F_H_current*spatial_window(i))/M_S

last_x2 = last_x1.copy()
last_x1 = x.copy()
F_H_current = K * abs(eta_current - x[x_hammer])**p

# rest of the iterations
for n in range(H): # time

    if(n <= 2):
        continue

    eta_next = 2*eta_current - eta_last - (delta_t**2*N*F_H_current)/M_S
    eta_last = eta_current
    eta_current = eta_next

    push = eta_current - x[x_hammer]
    if(push > 0):
        F_H_current = K * abs(push)**p
    else:
        F_H_current = 0
    F_H_current = 0

    next_x = [0] * N

    for i in range(2, N-2): # space

        # term_1 = a_1*x[i] + a_2*last_x1[i]
        # term_2 = a_3*(x[i+1] + x[i-1]) + a_4*(x[i+2] + x[i-2])
        # term_3 = a_5*(last_x1[i+1] + last_x1[i-1] + last_x2[i])
        # term_4 = (delta_t**2 * N*F_H_current * spatial_window(i))/M_S

        # next_x[i] = term_1 + term_2 + term_3 + term_4
        next_x[i] = 2*(1-r**2)*x[i] + r**2 * (x[i+1] + x[i-1]) - last_x1[i]

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

    last_x2 = last_x1.copy()
    last_x1 = x.copy()
    x = next_x.copy()

    x_out[n] = x[x_sample]

# plotting
plt.plot(t, x_out)
plt.title("Response")
plt.xlabel("t")
plt.ylabel("x(t, x_sample)")
plt.grid()
plt.show()

sd.play(x_out, f_e)
sd.wait()