# Import of packages:

import matplotlib.pyplot as plt
import numpy as np

# Model parameters:

m = 70  # [kg]
A = 0.7  # [m2]
g = 9.81  # [m/s^2]
Cd = 1  # [1]
rho = 1.29  # [kg/m3]

# Simulation time settings:

dt = 0.01  # [s]
t_start = 0  # [s]
t_stop = 20  # [s]
N_sim = int((t_stop - t_start)/dt) + 1  # Num time-steps

# Preallocation of arrays for plotting:

t_array = np.zeros(N_sim)
s_array = np.zeros(N_sim)
v_array = np.zeros(N_sim)
F_g_array = np.zeros(N_sim)
F_f_array = np.zeros(N_sim)

# Initialization:

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

# Simulation loop:

for k in range(0, N_sim):

    t_k = k*dt

    # Time-derivatives:
    ds_dt_k = v_k

    F_g_k = m*g
    F_f_k = 0.5*rho*Cd*A*v_k**2
    dv_dt_k = (1/m)*(F_g_k - F_f_k)

    # State predictions:
    s_kp1 = s_k + dt*ds_dt_k
    v_kp1 = v_k + dt*dv_dt_k

    # Arrays for plotting:
    t_array[k] = t_k
    s_array[k] = s_k
    v_array[k] = v_k
    F_g_array[k] = F_g_k
    F_f_array[k] = F_f_k

    # Time shift:
    s_k = s_kp1
    v_k = v_kp1

# Plotting:
    
plt.close('all')  # Close figures
plt.figure(1, figsize=(12, 9))

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

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


plt.subplot(3, 1, 3)
plt.plot(t_array, F_g_array, 'm', label='F_g')
plt.plot(t_array, F_f_array, 'g', label='F_f')
plt.legend()
plt.grid()
plt.xlabel('t [s]')
plt.ylabel('[N]')

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

# Analysis regarding 95% static speed:

v_s = np.sqrt(2*m*g/(rho*Cd*A))  # Static speed
v_95 = 0.95*v_s
x = np.argwhere(v_array >= v_95)
k_95 = np.amin(x)  # Time index of 95% v_s
t_95 = t_array[k_95]
s_95 = s_array[k_95]
print('t_95 [s] =', f'{t_95:.2f}')
print('s_95 [s] =', f'{s_95:.2f}')