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 = ...

    # 
    k = k + 1

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

# LQR feedback matrix
K = ...

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

# Simulate closed-loop (with linear system dynamics)
A_CL = ...

Nsim = 50

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

x = [x0]
u = []
for k in range(Nsim):
    u.append(...)
    x.append(...)


# plot results
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()


# Linear MPC: define NLP
N = ... # 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 += ...
obj += ...
opti.minimize(obj)

# constraints
constr = []
opti.subject_to(... == 0)           # initial condition
for i in range(N):
    opti.subject_to(... == 0)       # system dynamics    
    opti.subject_to(...)   # input constraints
    opti.subject_to(...)

# 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
    ... # hint: opti.set_value(...)

    # solve OCP
    sol = ...

    # 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()