'''
Calculating poles and simulating impulse response of a mass-spring-damper
Finn Aakre Haugen (finn@techteach.no)
Updated 2025 06 29
'''

# %% Imports:

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

#%% Model parameters

m = 20  # [kg] 
d = 4  # [N/(m/s)] 
k = 2  # [N/m] 

#%% Creating the transfer function

s = control.tf('s')

H = 1/(m*s**2 + d*s + k)

#%% Calculation and plotting of poles

poles = control.poles(H)

plt.close('all')
plt.figure(1, figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.plot(poles.real, poles.imag, 'or', label='Poles')
plt.plot([0, 0], [-1, 1], 'k--')  # Vertical blue line
plt.plot([-1, 1], [0, 0], 'k--')  # Horizontal blue line
plt.legend()
plt.grid()
plt.xlim(-1, 1)
plt.ylim(-1, 1)
plt.xlabel('Real')
plt.ylabel('Imag')
plt.title(f'Poles = {np.around(poles, decimals=3)}')

print (f'Poles of H(s) = {np.around(poles, decimals=3)}')

#%% Simulation and plotting of impulse response

dt = 0.01
t_start = 0
t_stop = 60
t_array = np.arange(t_start, t_stop+dt, dt)    

(t, y) = control.impulse_response(H, t_array)

plt.subplot(1, 2, 2)
plt.plot(t, y, 'b', label='Impulse response, h')
plt.legend()
plt.grid()
plt.xlabel('t [s]')

# plt.savefig('poles_and_impulse_resp_mfd.pdf')
plt.show()