import scipy.signal as sig
import matplotlib.pyplot as plt
import numpy as np
import math

# simple first order step response simulation
# www.halvorsen.blog/documents/programming/python/resources/powerpoints/State Space Models with Python.pdf

# simulation Parameters

# string pde (wave equation)
# rho*A*d2y(x,t)/dt2 + c*dy(x,t)/dt - T*d2y(x,t)/dx2 = f(x,t)
# where rho*A = 1 dimensional string density, c = string damping, T = string tension, y(x,t) = displacement, f(x,t) = exitation force

# assume fixed ends: y(0, t) = y(L, t) = 0

# displacement y(x,t) can be represented as the sum from n=1 to N of qn(t)*sin(n*pi*x/L) where q_n(t) is the modal coordinate of mode n (displacement is the sum of all resonating modes)
# therefore:
# q..n + 2*zeta_n*omega_n*q.n + omega_n2*q_n = b_n*u(t)
# where omega_n = n*pi/L * sqrt(T/(rho*A)) for an ideal string
# and b_n = sin(n*pi*x_h/L) but assuming the hammer strikes at midpoint, x_h = L/2 ( b_n = sin(n*pi/2) )

# state space representation:
# each mode:
#「 q_n_dot     | =「      0               1        | 「  q_n   | + 「  0  |* u
# L q_n_dot_dot 」= L -omega_n^2  -2*zeta_n*omega_n 」 L q_ndot 」   L b_n 」
# and x_dot = Ax + Bu
# where A = diagonal matrix of A_n and B = [ 0 b_1 0 b_2 ... b_n]^T

# lets start with 3 nodes where the string is tuned to 440hz (we'll get to arbitrary modes evantually)
f_1 = 440 # fundamental frequency
def f_n(n):
    return f_1 * n
def omega_n(n):
    return 2*math.pi*f_n(n) # a cooler option would be omega_n = c*n*omega_1*sqrt(1+B*n^2) to factor in string stiffness to its vibration mode

# x = [ q1, q1dot, q2, q2dot, q3, q3dot ]^T < --state vector
omega_1 = omega_n(1)
omega_2 = omega_n(2)
omega_3 = omega_n(3)
zeta_1 = 0.0001 # i guessed
zeta_2 = 2 * zeta_1
zeta_3 = 3 * zeta_1
A = [
        [          0,                 1,           0,                 0,           0,                 0],
        [-omega_1**2, -2*zeta_1*omega_1,           0,                 0,           0,                 0],
        [          0,                 0,           0,                 1,           0,                 0],
        [          0,                 0, -omega_2**2, -2*zeta_2*omega_2,           0,                 0],
        [          0,                 0,           0,                 0,           0,                 1],
        [          0,                 0,           0,                 0, -omega_3**2, -2*zeta_3*omega_3]
    ] # isnt this formatting gorgeous

B = [ [0], [0.707], [0], [0], [0], [-0.707] ]
c_1 = 0.001
c_2 = c_1 / 2
c_3 = c_1 / 3
C = [[-c_1*omega_1**2, -2*c_1*zeta_1*omega_1, -c_2*omega_2**2, -2*c_2*zeta_2*omega_2, -c_3*omega_3**2, -2*c_3*zeta_3*omega_3]]
D = c_1 * 0.707 + c_2 * 0 + c_3 * (-0.707)

# input
#v_h = 100 # hammer velocity
#u = v_h * impulse(t)

x0 = [[0], [0], [0], [0], [0], [0]]
start = 0
stop = 10
step = 1/44100

t = np.arange(start,stop,step)
sys = sig.StateSpace(A, B, C, D)

# step Response
t, y = sig.impulse(sys, T=t)

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