import matplotlib.pyplot as plt
import numpy as np

m = 1  # [kg]
K = 10  # [N/m]
D = 0.5  # [N/(m/s)]

# For state limitation. Not effective here.
s_min = -np.inf; s_max = np.inf,
v_min = -np.inf; v_max = np.inf

dt = 0.001  # [s]
t_start = 0  # [s]
t_stop = 10  # [s]
N_sim = int((t_stop - t_start)/dt) + 1

t_array = np.zeros(N_sim)
s_array = np.zeros(N_sim)
v_array = np.zeros(N_sim)
F_array = np.zeros(N_sim)

s_k = s_init = 0  # [m]
v_k = s_init = 0  # [m/s]

for k in range(0, N_sim):

    t_k = k*dt

    # State limitation:
    s_k = np.clip(s_k, s_min, s_max)
    v_k = np.clip(v_k, v_min, v_max) 

    if 0 <= t_k <= 1: F_k = 0  # [N]
    else: F_k = 1  # [N]

    ds_dt_k = v_k
    dv_dt_k = (1/m)*(F_k - D*v_k - K*s_k)

    s_kp1 = s_k + dt*ds_dt_k
    v_kp1 = v_k + dt*dv_dt_k
    
    t_array[k] = t_k
    s_array[k] = s_k
    v_array[k] = v_k
    F_array[k] = F_k
    
    s_k = s_kp1
    v_k = v_kp1
    
plt.close('all')
plt.figure(1, figsize=(12, 9))

plt.subplot(3, 1, 1)
plt.plot(t_array, s_array, 'b', label='s')
plt.legend()
plt.xlabel('t [s]')
plt.ylabel('[m]')
plt.grid()

plt.subplot(3, 1, 2)
plt.plot(t_array, v_array, 'g', label='v')
plt.legend()
plt.grid()
plt.xlabel('t [s]')
plt.ylabel('[m/s]')

plt.subplot(3, 1, 3)
plt.plot(t_array, F_array, 'r', label='F')
plt.legend()
plt.grid()
plt.xlabel('t [s]')
plt.ylabel('[N]')

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