# %% Import:

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

# %% Model parameters:

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

# %% System matrices as 2D arrays:

A = np.array([[0, 1], [-k/m, -d/m]])
B = np.array([[0], [1/m]])
C = np.array([[1, 0]])
D = np.array([[0]])

# %% Creating the state space model:

S = control.ss(A, B, C, D)

# %% Defining signals:

t0 = 0  # [s]
t1 = 50  # [s]
dt = 0.01  # [s]
nt = int(t1/dt) + 1 # Number of points of sim time
t = np.linspace(t0, t1, nt)
F = 10*np.ones(nt)  # [N]

# %% Initial state:

x1_0 = 1  # [m]
x2_0 = 0  # [m/s]
x0 = np.array([x1_0, x2_0])

# %% Simulation:

(t, y, x) = control.forced_response(S, t, F, x0,
                                    return_x=True)

# %% Extracting individual states:

x1 = x[0,:]
x2 = x[1,:]

# %% Plotting:

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

plt.subplot(3, 1, 1)
plt.plot(t, x1, 'blue')
plt.grid()
plt.legend(labels=('x1 [m]',))

plt.subplot(3, 1, 2)
plt.plot(t, x2, 'green')
plt.grid()
plt.legend(labels=('x2 [m/s]',))

plt.subplot(3, 1, 3)
plt.plot(t, F, 'red')
plt.grid()
plt.legend(labels=('F [N]',))
plt.xlabel('t [s]')

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