"""
Simulation of mass-spring-damper with Python Control Package
Finn Aakre Haugen, TechTeach, finn@techteach.no
2025 03 28
"""

# %% Imports

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

#%% Defining Laplace variable for def of transfer function

s = control.tf('s')

#%% Model

m = 20  # [kg]
D = 4  # [N/(m/s)]
K = 2  # [N/m]

H = 1/(m*s**2 + D*s + K)

# %% Simulation of unit step response

t_start = 0
t_stop = 50
ts = 0.01
nt = int((t_stop-t_start)/ts) + 1 # Number of points of sim time
t = np.linspace(t_start, t_stop, nt)

Ampl_F = 5
F = np.ones(nt)*Ampl_F

# %% Simulation
    
(t, y) = control.forced_response(H, t, F)

# %% Plotting

plt.close('all')
plt.figure(1)

plt.subplot(2,1,1)
plt.plot(t, y, color='red', label='y')
plt.legend()
plt.grid()
plt.ylabel('[m]')

plt.subplot(2,1,2)
F_plot = F
F_plot[0] = 0  # To get a nice step in the plot of F
plt.plot(t, F_plot, color='blue', label='F')
plt.legend()
plt.grid()
plt.ylabel('[N]')
plt.xlabel('t [s]')

plt.savefig('plot_sim_msd.pdf')
plt.show()