import matplotlib.pyplot as plt
import numpy as np

m = 250
K = 49050
D = 7004
g = 9.81
S = 0.55

dt = 0.001
t_start = 0
t_stop = 4
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)
z_array = np.zeros(N_sim)

s_k = s_init = 0.5
v_k = v_init = 0

for k in range(0, N_sim):

    t_k = k*dt
    
    if (0 <= t_k <= 1): z_k = 0
    else: z_k = 0.05

    ds_dt_k = v_k
    dv_dt_k = (1/m)*(-m*g - D*v_k + K*(z_k - (s_k - S)))

    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
    z_array[k] = z_k

    s_k = s_kp1
    v_k = v_kp1
    
plt.close('all')
plt.figure(1, figsize=(12, 9))

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

plt.subplot(2, 1, 2)
plt.plot(t_array, z_array, 'r', label='z')
plt.legend()
plt.xlabel('t [s]')
plt.ylabel('[m]')
plt.grid()

# plt.savefig('plot_sim_hjuloppheng.pdf')
# plt.savefig('plot_sim_hjuloppheng_D_liten.pdf')
plt.show()