"""

Simulator of pendulum on cart ver 1 with LQ control

The model of the pendulum on cart is from the book
Nonlinear Systems by H K. Khalil (Prentice-Hall, 2000).

Finn Aakre Haugen, finn.haugen@usn.no

Updated 2021 04 25

"""
# %% Import of packages:

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

# %% Def of process model simulator

def process_sim(x1_k, x2_k, x3_k, x4_k, u_k, 
                proc_params, dt):

    (M, m, d, L, I, g) = proc_params

    c3 = np.cos(x3_k)
    s3 = np.sin(x3_k)
       
    Den1 = (I + m*(L**2))*(m + M) - (m*L*c3)**2

    f2 = ((-m*L*c3)*(m*g*L*s3) 
          + (I + m*(L**2))*((m*L*(x4_k**2))*s3-d*x2_k))/Den1
    B2 = (I + m*(L**2))/Den1

    f4 = ((m + M)*(m*g*L*s3) 
          + (-m*L*c3)*((m*L*(x4_k**2))*s3-d*x2_k))/Den1
    B4 = (-m*L*c3)/Den1
    
    dx1_dt_k = x2_k
    dx2_dt_k = f2 + B2*u_k
    dx3_dt_k = x4_k
    dx4_dt_k = f4 + B4*u_k

    x1_kp1 = x1_k + dt*dx1_dt_k
    x2_kp1 = x2_k + dt*dx2_dt_k
    x3_kp1 = x3_k + dt*dx3_dt_k
    x4_kp1 = x4_k + dt*dx4_dt_k

    return (x1_kp1, x2_kp1, x3_kp1, x4_kp1)

# %% Def of LQ controller function:

def fun_lqr(r_x1_k, r_x3_k, x1_k, x2_k, x3_k, x4_k, G):
    
    u_k = -(G[0,0]*(x1_k - r_x1_k)
            +G[0,1]*(x2_k)
            +G[0,2]*(x3_k - r_x3_k)
            +G[0,3]*x4_k)
    
    return u_k

# %% Process parameters:

M = 1  # [kg]
m = 0.1  # [kg]
L = 0.5  # [m]
d = 0.0  # [N/(m/s)]
I = m*((2*L)**2)/12
g = 9.81  # [m/(s**2)]

proc_params = (M, m, d, L, I, g)

# %% Selection of mode:

op_mode = -1  # -1: Downwards. 1: Upwards.

if op_mode==1:
    r_x3_deg_const = 0  # [deg]
    x3_init_deg =  10  # [deg]
    K_cos = 1
    K_sin = 1
elif op_mode==-1:
    r_x3_deg_const = 180  # [deg]
    x3_init_deg =  170  # [deg]
    K_cos = -1
    K_sin = -1

# %% LQR gain:

Den2 = (I + m*(L**2))*(m + M) - (m*L)**2

A = np.array([
    [0, 1, 0, 0],
    [0, -(I+m*L**2)*d/Den2, -((m*L)**2)*g/Den2, 0],
    [0, 0, 0, 1],
    [0, K_cos*m*L*d/Den2, (m+M)*m*g*L*K_sin/Den2, 0]
    ])

B = np.array([[0],
             [(I+m*L**2)/Den2],
             [0],
             [-K_cos*m*L/Den2]])

Q = np.array([[100, 0, 0, 0],
              [0, 0, 0, 0],
              [0, 0, 100, 0],
              [0, 0, 0, 0]])

R = np.array([[1]])

(G, S, E) = control.lqr(A, B, Q, R)
print('LQ controller gain, G = ', G)
print('Eigenvalues of control system: = ', E)

# Plotting eigenvalues of control system:

plt.close('All')
plt.figure(1)
plt.plot(E.real, E.imag, 'xr')
# plt.xlim(-5, 1)
# plt.ylim(-1, 1)
plt.grid()

# %% State limits:

x1_max = np.inf
x1_min = -np.inf
x2_max = np.inf
x2_min = -np.inf
x3_max = np.inf
x3_min = -np.inf
x4_max = np.inf
x4_min = -np.inf

# %% Simulation time settings:

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

# %% Preallocation of arrays for plotting etc:

t_array = np.zeros(N_sim)

u_array = np.zeros(N_sim)

x1_array = np.zeros(N_sim)
x2_array = np.zeros(N_sim)
x3_array = np.zeros(N_sim)
x3_deg_array = np.zeros(N_sim)
x4_array = np.zeros(N_sim)

r_x1_array = np.zeros(N_sim)
r_x3_array = np.zeros(N_sim)
r_x3_deg_array = np.zeros(N_sim)

# %% Initialization:

x1_k = x1_init = 0  # [m]
x2_k = x2_init = 0  # [m/s]
x3_k = x3_init = x3_init_deg*(np.pi/180)  # [rad]
x4_k = x4_init = 0  # [rad/s]

# %% Simulation loop:

for k in range(0, N_sim):

    t_k = k*dt
    
    Ampl = 0.1  # [m] Amplitude of r_x1
    Tp = 20  # [s] Period of r_x1
    if (0 <= t_k < 10):
        r_x1_k = 0  # [m]
        r_x3_deg_k = r_x3_deg_const  # [deg]
        r_x3_k = r_x3_deg_k*(np.pi/180)  # [rad]
    elif (10 <= t_k < 20):
        r_x1_k = 0.1  # [m]
        r_x3_deg_k = r_x3_deg_const  # [deg]
        r_x3_k = r_x3_deg_k*(np.pi/180)  # [rad]
    else:
        r_x1_k = 0.1 + Ampl*np.sin((2*np.pi/Tp)*t_k)  # [m]
        r_x3_deg_k = r_x3_deg_const  # [deg]
        r_x3_k = r_x3_deg_k*(np.pi/180)  # [rad]
            

    # Limiting the state variables:
    x1_k = np.clip(x1_k, x1_min, x1_max)
    x2_k = np.clip(x2_k, x2_min, x2_max)
    x3_k = np.clip(x3_k, x3_min, x3_max)
    x4_k = np.clip(x4_k, x4_min, x4_max)
    
    # LQR:
    u_k = fun_lqr(r_x1_k, r_x3_k, x1_k, x2_k, x3_k, x4_k, G)

    # Process sim (Euler integration):
    (x1_kp1, x2_kp1, x3_kp1, x4_kp1) = process_sim(
        x1_k,
        x2_k,
        x3_k,
        x4_k,
        u_k,
        proc_params,
        dt)

    # Updating arrays for plotting:
    t_array[k] = t_k
    u_array[k] = u_k
    x1_array[k] = x1_k
    x2_array[k] = x2_k
    x3_array[k] = x3_k
    x3_deg_array[k] = x3_k*180/np.pi
    x4_array[k] = x4_k
    r_x1_array[k] = r_x1_k
    r_x3_array[k] = r_x3_k
    r_x3_deg_array[k] = r_x3_deg_k

    # Time index shift:
    x1_k = x1_kp1
    x2_k = x2_kp1
    x3_k = x3_kp1
    x4_k = x4_kp1

# %% Plotting:

# plt.close('all')

plt.figure(2)

plt.subplot(3, 1, 1)
plt.plot(t_array, r_x1_array, 'r', label='r_x1')
plt.plot(t_array, x1_array, 'b', label='x1')
plt.legend()
plt.grid()
plt.xlim(t_start, t_stop)
plt.ylim(-0.1, 0.6)
plt.xlabel('t [s]')
plt.ylabel('[m]')

plt.subplot(3, 1, 2)
plt.plot(t_array, r_x3_deg_array, 'r', label='r_x3')
plt.plot(t_array, x3_deg_array, 'b', label='x3')
plt.legend()
plt.grid()
plt.xlim(t_start, t_stop)
# plt.ylim(-1, 1)
plt.xlabel('t [s]')
plt.ylabel('[deg]')

plt.subplot(3, 1, 3)
plt.plot(t_array, u_array, 'r', label='u')
plt.legend()
plt.grid()
plt.xlim(t_start, t_stop)
# plt.ylim(-1, 1)
plt.xlabel('t [s]')
plt.ylabel('[N]')

# plt.savefig('pendulum_on_cart_lqr_standing.pdf')
# plt.savefig('pendulum_on_cart_lqr_hanging.pdf')
plt.show()