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

# 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)

# string parameters
E = 1 # youngs modulus
mu = 1 # linear mass density
kappa = 1 # radius of gyration
L = 1 # string length
M_S = mu*L # string mass
S = 1 # string cross sectional area
T = 1 # string tension
c = math.sqrt(T/mu) # transverse wave velocity
stiffness = 1 # string stiffness parameter
sigma = 1 # decay rate
tau = 1/sigma # decay time
omega = 1 # angular frequency

# hammer parameters
M_H = 1 # hammer mass
HSMR = M_H/M_S # hammer-mass string ratio
V_H_0 = 1 # initial hammer velocity at t=0
x_0 = 1 # distance of hammer from agraffe
alpha = x_0 / L # relative hammer striking position

# simulation parameters
f1 = 440 # fundamental frequency
f_e = 44100 # sampling frequency
N = 100 # number of string segments
delta_t = 1/f_e # time step
delta_x = L/N # spatial step
H = f_e * 10 # length of simulation in time

# empirical constants
b_1 = 1 # some constant
b_3 = 1 # 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_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


x = [0] * N # current string position
last_x1 = [0] * N # string position from last timestep
last_x2 = [0] * N # string position from two timesteps ago
x_next = [0] * N # buffer for next string position

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

for n in range(H): # time

    for i in range(N): # space
        x_out[n] = math.sin(n/10000)

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