import numpy as np
import matplotlib.pyplot as plt

# discrete system matrices
A = np.array([
    [1,   0.1],
    [0.1, 0.99]
])

B = np.array([
    [0],
    [0.1]
])

# weight matrices
Q = np.eye(2)
R = 1

# initialize infinite cost-to-go
P0 = Q
Pk = Q

# Ricatti recursion
k = 0
tol = 1e-8
while k < 1 or (np.linalg.norm(P0 - Pk) > tol):

    # update cost-to-go matrix
    P0 = Pk

    # perform Ricatti recursion
    Pk = A.T@P0@A - (A.T@P0@B)@np.linalg.inv(R + B.T@P0@B)@(B.T@P0@A) + Q

    # 
    k = k + 1

print('Infinite-horizon cost-to-go P: \n {}'.format(Pk))

# LQR feedback matrix
K = - np.linalg.inv(B.T@Pk@B +R)@B.T@Pk@A

print('LQR feedback matrix K: \n {}'.format(K))

# Simulate closed-loop (with linear system dynamics)
A_CL = A + B@K

Nsim = 50

x0 = np.array([
    [np.pi/6],
    [0]
])

x = [x0]
u = []
for k in range(Nsim):
    u.append(K@x[-1].squeeze())
    x.append(A_CL@x[-1])
x1 = [state[0] for state in x]
x2 = [state[1] for state in x]

plt.plot(x1, label = '$x_{1,LQR}$', color = 'tab:blue')
plt.plot(x2, label = '$x_{2, LQR}$', color = 'tab:orange')
plt.step(range(len(u)),u, where = 'post', label = '$u_{LQR}$', color = 'tab:green')
plt.legend()
plt.xlabel('$k$')
plt.grid()
#plt.show()


# define NLP
N = 10 # MPC horizon
nx = 2
nu = 1

# create empty optimization problem
import casadi as ca
opti = ca.Opti()

# decision variables
u_all = opti.variable(nu, N)       # all controls
x_all = opti.variable(nx, N+1)     # all states
x0hat = opti.parameter(nx, 1)         # initial state estimate

# objective function
obj = 0
for i in range(N):
    obj += R*u_all[:,i]**2
    obj += x_all[:,i].T@Q@x_all[:,i]

obj += x_all[:,N].T@Pk@x_all[:,N]
opti.minimize(obj)

# constraints
constr = []
opti.subject_to( x_all[:,0] - x0hat == 0)           # initial condition
for i in range(N):                          # dynamics        
    opti.subject_to(x_all[:,i+1] - A@x_all[:,i] - B@u_all[:,i] == 0)
    opti.subject_to(u_all <= 1)
    opti.subject_to(- u_all <= 1)

# Setting solver to ipopt and solving the problem:
opti.solver('ipopt')

# closed-loop response
x_MPC = [x0]
u_MPC = []
for k in range(Nsim):

    # update initial condition and solve OCP
    opti.set_value(x0hat, x_MPC[-1])
    sol = opti.solve()

    # retrieve solution
    X_all = sol.value(x_all)
    U_all = sol.value(u_all)

    # store first controls and next state (perfect model)
    u_MPC.append(U_all[0])
    x_MPC.append(X_all[:,1])

x1_MPC = [state[0] for state in x_MPC]
x2_MPC = [state[1] for state in x_MPC]

plt.plot(x1_MPC, label = '$x_{1,MPC}$', linestyle = '--', color = 'tab:blue')
plt.plot(x2_MPC, label = '$x_{2,MPC}$', linestyle = '--', color = 'tab:orange')
plt.step(range(len(u_MPC)),u_MPC, where = 'post', label = '$u_{MPC}$', linestyle = '--', color = 'tab:green')
plt.legend()
plt.xlabel('$k$')
plt.grid(True)
plt.show()