import numpy as np
import matplotlib.pyplot as plt
import casadi as ca

# Set up problem data:

# discrete system matrices
A = np.array([
    [1,   0.1],
    [0.1, 0.99]
])
B = np.array([
    [0],
    [0.1]
])
nx = 2
nu = 1

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

# MPC horizon
N = 10 

# Question 1a)

# create empty optimization problem
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 += ...
opti.minimize(obj)

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

opti.subject_to(... == 0) # terminal constraint

# Setting solver to ipopt and solving the problem:
opti.solver('ipopt', {}, {'max_iter': 100})

# Question 1b): Investigate feasible set
n_points = 21
theta0 = np.linspace(-np.pi/6, np.pi/6, n_points, endpoint= True)
dtheta0 = np.linspace(-2*np.pi/3, 2*np.pi/3, n_points, endpoint=True)
thetav, dthetav = np.meshgrid(theta0, dtheta0)
u_MPC = np.zeros((n_points, n_points))

for i in range(len(theta0)):
    for j in range(len(dtheta0)):
        x0 = ...
        opti.set_value(...)
        
        try:
            sol = ...
            u_MPC[i,j] = ...
        except:
            u_MPC[i,j] = np.nan
            thetav[i,j] = np.nan
            dthetav[i,j] = np.nan

plt.figure()
plt.plot(thetav, dthetav, marker = 'o', linestyle = 'none', color = 'tab:green')
plt.xlim([-np.pi/3, np.pi/3])
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\dot{\theta}$')
plt.show()


# Question 1c): Feasible set of LQR controller

# Compute LQR feedback matrix and cost-to-go (see Exercise 3)
P0 = Q
Pk = Q
k = 0
tol = 1e-8
while k < 1 or (np.linalg.norm(P0 - Pk) > tol):
    P0 = Pk
    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
K = - np.linalg.inv(B.T@Pk@B +R)@B.T@Pk@A
A_CL = A + B@K

# sample
thetav_lqr, dthetav_lqr = np.meshgrid(theta0, dtheta0)
u_lqr = np.zeros((n_points, n_points))
Nsim = 50 # closed-loop simulation horizon for checking constraint violation
for i in range(len(theta0)):
    for j in range(len(dtheta0)):
        x0 = ...
        u0 = ...

        violate_constraints = False
        for k in range(Nsim):
            uk = ...

            if np.abs(uk) > 1 or np.abs(x0[0]) > np.pi/6:
                violate_constraints = True
                break

            x0 = ... # forward simulation

        if violate_constraints:
            u_lqr[i,j] = np.nan
            thetav_lqr[i,j] = np.nan
            dthetav_lqr[i,j] = np.nan
        else:
            u_lqr[i,j] = u0

plt.figure()
plt.plot(thetav, dthetav, marker = 'o', linestyle = 'none', color = 'tab:green')
plt.plot(thetav_lqr, dthetav_lqr, marker = 's', linestyle = 'none', color = 'tab:red', alpha = 0.6)
plt.xlim([-np.pi/3, np.pi/3])
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\dot{\theta}$')
plt.show()

# Question 1d): Plot feedback law
fig = plt.figure()
plt.plot(...)
plt.grid(True)
plt.xlim([-np.pi/3, np.pi/3])
plt.xlabel(r'$\theta$')
plt.ylabel(r'$u_0$')
plt.show()


# Question 2b): Terminal set formulation - investigate feasible set
Af = np.array([
   [-0.0053,   -0.5294],
    [0.0053,    0.5294],
   [-0.5198,   -0.9400],
    [0.5198,    0.9400],
   [-1.0000,         0],
    [1.0000,         0]
]
)

bf = np.array([
    [1], [1], [1], [1], [3], [3]
])

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) # dynamics    
    opti.subject_to(... <= ...)  # constraints

opti.subject_to(... == 0) # terminal constraint

# Setting solver to ipopt and solving the problem:
opti.solver('ipopt', {}, {'max_iter': 100})

# test initial conditions
n_points = 21
thetaf, dthetaf = np.meshgrid(theta0, dtheta0)
u_MPCf = np.zeros((n_points, n_points))
for i in range(len(theta0)):
    for j in range(len(dtheta0)):
        x0 = ...
        opti.set_value(...)
        
        try:
            sol = ...
            u_MPCf[i,j] = ...
        except:
            u_MPCf[i,j] = np.nan
            thetaf[i,j] = np.nan
            dthetaf[i,j] = np.nan

plt.figure()
plt.plot(thetaf, dthetaf, marker = 'o', linestyle = 'none', color = 'tab:blue', alpha = 0.8)
plt.plot(thetav, dthetav, marker = 'o', linestyle = 'none', color = 'tab:green', alpha = 0.8)
plt.xlim([-np.pi/3, np.pi/3])
plt.xlabel(r'$\theta$')
plt.ylabel(r'$\dot{\theta}$')
plt.show()