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

# 1) Задаём параметры
A = np.array([[0,      0,       1,      0],
              [0,      0,       0,      1],
              [0, 5.8773, -1.9436, -0.3484],
              [0, 24.5313, -0.3823, -1.4543]])
B = np.array([[0],
              [0],
              [6.2491],
              [1.2311]])
Q = np.diag([100, 0, 0, 0])
R = np.array([[1]])

# 2) Решаем LQR — получаем K, решение Риккати X и собственные числа E
K, X, E = control.lqr(A, B, Q, R)

# 5) Начальные условия и временная сетка
x0 = [0, np.pi/17, 0, 0]
t = np.linspace(0, 10, 1001)

#
control_coefs = -K

theta_ddot_func, phi_ddot_func = non_lin_theta_phi_funcs(lhs_1, lhs_2)
theta_ddot_func = theta_ddot_func.subs(parameters)
phi_ddot_func = phi_ddot_func.subs(parameters)

controled_theta_ddot_func = add_control(theta_ddot_func,control_coefs)
controled_phi_ddot_func = add_control(phi_ddot_func,control_coefs)

theta_ddot_lamb = sp.lambdify(vars_vector, controled_theta_ddot_func)
phi_ddot_lamb = sp.lambdify(vars_vector, controled_phi_ddot_func)

non_lin_sys = non_lin_sys_creator(theta_ddot_lamb, phi_ddot_lamb)

non_lin_sol = solve_ivp(fun = non_lin_sys,y0 = x0[:4],t_span = (t[0],t[-1]), t_eval = t, rtol = 1e-2)
#


# 3) Замкнутая матрица системы
A_cl = A - B @ K

# 4) Формируем модель состояния-выхода для замкнутой системы
#    Здесь C=I, D=0 чтобы на выходе видеть сами состояния
C = np.eye(4)
D = np.zeros((4, 1))
sys_cl = control.ss(A_cl, B, C, D)

# 6) Моделируем свободную (нулевой вход) динамику системы от x0
t_out, y_out = control.initial_response(sys_cl, T=t, X0=x0)

# 7) Рисуем 4 графика переменных состояния
labels = [r'$\theta$', r'$\varphi$', r'$\dot{\theta}$', r'$\dot{\varphi}$']
fig, axes = plt.subplots(4, 1, figsize=(8, 10), sharex=True)

for i in range(4):
    axes[i].plot(t_out, y_out[i, :], 'b', label="линейная")
    axes[i].plot(non_lin_sol["t"], non_lin_sol["y"][i], 'r', label="нелинейная")
    axes[i].set_ylabel(labels[i], fontweight='bold')
    axes[i].legend()
    axes[i].set_xlabel('t')

#plt.tight_layout()
plt.show()