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

# Question 1a)

# Set up states and controls
nx = 2
nu = 1
x  = ca.SX.sym('x', nx)
u  = ca.SX.sym('u', nu)

# ODE
x0dot = ...
x1dot = ...
f = ca.Function('f', [x,u], [...])

# ODE data in dict
ode = {}
ode['x'] = x
ode['p'] = u
ode['ode'] = ...

# integrator options
Ts = ... # sampling time
int_opts = {'tf': ..., 'number_of_finite_elements': ...} #  
F = ca.integrator('F','rk', ode, int_opts)

# Question 1b)
x0 = ...
u0 = ...
T = ... # simulation horizon

xsim = [x0]
for k in range(int(T/Ts)):
    xnext = F(x0=..., p=...)['xf']
    xsim.append(xnext)
xsim = ca.horzcat(*xsim).full().squeeze()

plt.figure()
plt.plot(xsim[0,:], color = 'tab:blue', label = r'$x_1$ ')
plt.plot(xsim[1,:], color = 'tab:orange', label = r'$x_2$ ')
plt.legend()
plt.grid(True)
plt.title('Forward simulation')
plt.show()


# Question 2a)

# set initial state
x0 = ...

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

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


# construct objective of sequential formulation
xk = x0hat
obj = 0
for k in range(N):
    obj += ... # stage cost
    xk = ... # symbolic forward simulation
obj += ... # terminal cost
opti.minimize(obj)

# create solver, set initial condition and solve
opti.solver('ipopt')
opti.set_value(x0hat, x0)
sol = opti.solve()


# optimal state trajectory via forward simulation
U_sequential = sol.value(u_all)
xsim = [x0]
for k in range(N):
    xsim.append(F(x0=xsim[-1], p=U_sequential[k])['xf'])
X_sequential = ca.horzcat(*xsim).full().squeeze()

plt.figure()
plt.plot(X_sequential[0,:], color = 'tab:blue', label = r'$x_1$ ')
plt.plot(X_sequential[1,:], color = 'tab:orange', label = r'$x_2$')
plt.legend()
plt.grid(True)
plt.title('Sequential solution')
plt.show()


# Question 2b)

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

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

# construct objective of simultaneous formulation
obj = 0
for k in range(N):
    obj += ... # stage cost
obj += ... # terminal cost
opti.minimize(obj)

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

# create solver, set initial condition and solve
opti.solver('ipopt')
opti.set_value(x0hat, x0)
sol = opti.solve()

# retrieve solution
U_simultaneous = sol.value(u_all)
X_simultaneous = sol.value(x_all)

plt.figure()
plt.plot(X_sequential[0,:], color = 'tab:blue', label = r'$x_1$ seq.')
plt.plot(X_sequential[1,:], color = 'tab:orange', label = r'$x_2$ seq.')
plt.plot(X_simultaneous[0,:], linestyle = '--', color = 'tab:blue', label = r'$x_1$ sim.')
plt.plot(X_simultaneous[1,:], linestyle = '--', color = 'tab:orange', label = r'$x_2$ sim.')
plt.legend()
plt.grid(True)
plt.title('Simultaneous vs. sequential solution')
plt.show()

# Question 2c)

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

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

# construct objective of simultaneous formulation
obj = ... # copy from unconstrained simultaneous case
opti.minimize(obj)

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

# create solver, set initial condition and solve
opti.solver('ipopt')
opti.set_value(x0hat, x0)
sol = opti.solve()

# retrieve solution
U_constrained = sol.value(u_all)
X_constrained = sol.value(x_all)

plt.figure()
plt.plot(X_simultaneous[0,:], color = 'tab:blue', label = r'$x_1$ ')
plt.plot(X_simultaneous[1,:], color = 'tab:orange', label = r'$x_2$')
plt.plot(X_constrained[0,:], linestyle = '--', color = 'tab:blue', label = r'$x_1$ constr')
plt.plot(X_constrained[1,:], linestyle = '--', color = 'tab:orange', label = r'$x_2$ constr')
plt.legend()
plt.grid(True)
plt.title('Constrained vs. unconstrained inputs')
plt.show()


# Question 3a)
x_clsim = [x0]
u_clsim = []
cpu_time = []
for k in range(100):
    opti.set_value(x0hat, ...)
    sol = opti.solve()

    U_sol = sol.value(u_all)
    X_sol = sol.value(x_all)
    x_clsim.append(...)
    u_clsim.append(...)

    cpu_time.append(opti.stats()['t_wall_total'])


x_clsim = ca.horzcat(*x_clsim).full().squeeze()
plt.figure()
plt.plot(x_clsim[0,:], color = 'tab:blue', label = r'$x_1$')
plt.plot(x_clsim[1,:], color = 'tab:orange', label = r'$x_2$')
plt.legend()
plt.grid(True)
plt.title('Closed-loop simulation')
plt.show()

# Question 3b)
sqp_opts = {'qpsol': ..., 'max_iter': ...}
opti.solver(..., sqp_opts)
x_rti = [x0]
u_rti = []
cpu_time_rti = []
for k in range(100):
    opti.set_value(x0hat, ...)
    try:
        sol = opti.solve() # will fail because of max_iter
    except:
        # do nothing
        32.0
    U_sol = opti.debug.value(u_all)
    X_sol = opti.debug.value(x_all)

    cpu_time_rti.append(opti.stats()['t_wall_total'])

    x_rti.append(...) # forward simulation with first optimal control input!
    u_rti.append(...)

    opti.set_initial(x_all, X_sol)
    opti.set_initial(u_all, U_sol)


x_rti = ca.horzcat(*x_rti).full().squeeze()
plt.figure()
plt.plot(x_clsim[0,:], color = 'tab:blue', label = r'$x_1$ full')
plt.plot(x_clsim[1,:], color = 'tab:orange', label = r'$x_2$_full')
plt.plot(x_rti[0,:], linestyle = '--', color = 'tab:blue', label = r'$x_1$ RTI')
plt.plot(x_rti[1,:], linestyle = '--', color = 'tab:orange', label = r'$x_2$ RTI')
plt.legend()
plt.grid(True)
plt.title('RTI vs. full MPC')
plt.show()

print('Full MPC, mean CPU time per iteration: {} s'.format(np.mean(cpu_time)))
print('RTI MPC, mean CPU time per iteration: {} s'.format(np.mean(cpu_time_rti)))