"""
NOTE: This example shows how to define a nonlinear programming problem,
where all derivative information is automatically generated by
using the AD tool CasADi.

You may want to directly pass C functions to FORCES, for example when you
are using other AD tools, or have custom C functions for derivate
evaluation etc. See the example file "NLPexample_ownFevals.m" for how to
pass custom C functions to FORCES.

This example solves an optimization problem for a robotic manipulator
with the following continuous-time, nonlinear dynamics

   ddtheta1/dt    = FILL IN
   ddtheta2/dt    = FILL IN
   dtheta1/dt     = dtheta1
   dtheta2/dt     = dtheta2
   dtau1/dt       = dtau1
   dtau2/dt       = dtau2

where (theta1,theta2) are the angles describing the configuration of the
manipulator, (dtau1,dtau2) are input torque rates, and a and b are constants
depending on geometric parameters.

The robotic manipulator must track a reference for the angles and their
velocities while minimizing the input torques.

There are bounds on all variables.

Variables are collected stage-wise into z = [dtau1 dtau2 theta1 dtheta1 theta2 dtheta2 tau1 tau2].

(c) Embotech AG, Zurich, Switzerland, 2013-2021
"""

import numpy as np
import scipy.integrate
import matplotlib.pyplot as plt
import forcespro
import forcespro.nlp
from robot_sym_dynamics import dynamics
import casadi

# Model Definition
# ----------------

# Problem dimensions
model = forcespro.nlp.SymbolicModel(21)   # horizon length
model.nvar  = 8    # number of variables
model.neq   = 6    # number of equality constraints
model.nh    = 0    # number of inequality constraint functions
model.npar  = 1    # reference for state 1

nx = 6
nu = 2
Tf = 2

Tsim = 20
Ns = int(Tsim/(Tf/(model.N-1)))   # simulation length?

timing_iter = 1

# Objective function
Q = np.diag([1000, 0.1, 1000, 0.1, 0.01, 0.01])
R = np.diag([0.01, 0.01])
#model.objective = lambda z, p: (1000*(z[2]-p[0]*1.2)**2 + 0.1*z[3]**2 + 1000*(z[4]+p[0]*1.2)**2 + 0.10*z[5]**2 + 0.01*z[6]**2 + 0.01*z[7]**2 +
#                                0.01*z[0]**2 + 0.01*z[1]**2)
model.LSobjective = lambda z, p: casadi.vertcat(np.sqrt(1000) * (z[2] - p[0]*1.2),
                                                np.sqrt(0.1) * (z[3]),
                                                np.sqrt(1000) * (z[4] + p[0]*1.2),
                                                np.sqrt(0.1) * z[5],
                                                np.sqrt(0.01) * z[6],
                                                np.sqrt(0.01) * z[7],
                                                np.sqrt(0.01) * z[0],
                                                np.sqrt(0.01) * z[1])

# Dynamics, i.e. equality constraints 

# We use an explicit RK4 integrator here to discretize continuous dynamics
integrator_stepsize = Tf/(model.N-1)
model.continuous_dynamics = dynamics

# Indices on LHS of dynamical constraint - for efficiency reasons, make
# sure the matrix E has structure [0 I] where I is the identity matrix.
model.E = np.concatenate([np.zeros((nx, nu)), np.identity(nx)], axis=1)

# Inequality constraints
# upper/lower variable bounds lb <= x <= ub
#                       inputs  |                states
#                    dtau1 dtau2  theta1  dtheta1  theta2  dtheta2  tau1  tau2
model.lb = np.array([-200,  -200, -np.pi,    -100, -np.pi,    -100, -100, -100])
model.ub = np.array([ 200,   200,  np.pi,     100,  np.pi,     100,   70,   70])

# Initial and final conditions
xinit = np.array([-0.4, 0, 0.4, 0, 0, 0])
model.xinitidx = range(2, 8)


# Generate solver
# ---------------

# Define solver options
codeoptions = forcespro.CodeOptions()
codeoptions.maxit = 200  # Maximum number of iterations
codeoptions.printlevel = 0  # Use printlevel = 2 to print progress (but not for timings)
codeoptions.optlevel = 3  # 0 no optimization, 1 optimize for size, 2 optimize for speed, 3 optimize for size & speed
codeoptions.nlp.integrator.Ts = integrator_stepsize
codeoptions.nlp.integrator.nodes = 5
codeoptions.nlp.integrator.type = 'ERK4'
codeoptions.solvemethod = 'SQP_NLP'
codeoptions.sqp_nlp.rti = 1
codeoptions.sqp_nlp.maxSQPit = 1
codeoptions.nlp.hessian_approximation = 'gauss-newton'

# Generate FORCESPRO solver
solver = model.generate_solver(codeoptions)


# Simulation
# ----------

# Call solver
# Set initial guess to start solver from:
x0i = model.lb + (model.ub - model.lb) / 2
x0 = np.transpose(np.tile(x0i, (1, model.N)))

problem = {"x0": x0}

x = np.zeros((nx, Ns + 1))
x[:, 0] = xinit
u = np.zeros((nu, Ns))
iter = np.zeros(Ns)
solvetime = np.zeros(Ns)
fevalstime = np.zeros(Ns)

# Set reference
problem["all_parameters"] = np.ones((model.N, 1))

# Cost
cost = np.zeros((Ns, 1))
ode45_intermediate_steps = 10
cost_integration_step_size = integrator_stepsize / ode45_intermediate_steps
cost_integration_grid = np.arange(0, integrator_stepsize, cost_integration_step_size)

# Simulation results
sim_states = []
sim_inputs = []
sim_solvetime = []
sim_qptime = []
sim_rsnorm = []
sim_res_eq = []
sim_fevalstime = []
sim_closed_loop_obj = []

for i in range(0, int(Ns / 2)):
    sim_states.append(x[:, i])

    # Set initial condition
    problem["xinit"] = x[:, i]
    
    # Time to solve the NLP!
    temp_time = +np.inf 
    temp_time_fevals = +np.inf 
    temp_time_qp = +np.inf

    for j in range(0, timing_iter):

        output, exitflag, info = solver.solve(problem)
        if exitflag != 1:
            raise RuntimeError('Solver failed at {}.'.format(i))

        if info.solvetime < temp_time:
            temp_time = info.solvetime

        if info.fevalstime < temp_time_fevals:
            temp_time_fevals = info.fevalstime

        # TODO QP time

    u[:, i] = output["x01"][0:nu]
    iter[i] = info.it
    solvetime[i] = temp_time
    fevalstime[i] = temp_time_fevals
    
    sim_inputs.append(u[:, i])
    sim_solvetime.append(temp_time)
    sim_fevalstime.append(temp_time_fevals)
    sim_qptime.append(temp_time_qp)
    sim_rsnorm.append(info.rsnorm)
    sim_res_eq.append(info.res_eq)
    
    # Simulate dynamics
    xtemp = scipy.integrate.odeint(lambda states, time: dynamics(states, u[:, i]).full().reshape(6), x[:, i], cost_integration_grid)
    
    # Compute cost
    for j in range(0, np.shape(xtemp)[0]):
        cost[i] += cost_integration_step_size * (xtemp[j, :] @ Q @ xtemp[j, :] + u[:, i] @ R @ u[:, i])
    sim_closed_loop_obj.append(cost[i])

    x[:, i + 1] = xtemp[-1, :]

# Set reference
problem["all_parameters"] = -np.ones((model.N, 1))

for i in range(int(Ns/2), Ns):
    
    sim_states.append(x[:, i])
    
    # Set initial condition
    problem["xinit"] = x[:, i]
    
    # Time to solve the NLP!
    temp_time = +np.inf
    temp_time_fevals = +np.inf
    for j in range(0, timing_iter):
        output, exitflag, info = solver.solve(problem)
        if info.solvetime < temp_time:
            temp_time = info.solvetime
        if info.fevalstime < temp_time_fevals:
            temp_time_fevals = info.fevalstime

    assert exitflag == 1, 'Some problem in FORCESPRO solver'

    u[:, i] = output["x01"][0:nu]
    iter[i] = info.it
    solvetime[i] = temp_time
    fevalstime[i] = temp_time_fevals
    
    sim_inputs.append(u[:, i])
    sim_solvetime.append(temp_time)
    sim_fevalstime.append(temp_time_fevals)
    sim_rsnorm.append(info.rsnorm)
    sim_res_eq.append(info.res_eq)
    
    # Simulate dynamics
    xtemp = scipy.integrate.odeint(lambda states, time: dynamics(states, u[:, i]).full().reshape(6), x[:, i], cost_integration_grid)

    # Compute cost
    for j in range(0, np.shape(xtemp)[0]):
        cost[i] += cost_integration_step_size*(xtemp[j, :] @ Q @ xtemp[j, :] + u[:, i] @ R @ u[:, i])
    sim_closed_loop_obj.append(cost[i])

    x[:, i + 1] = xtemp[-1, :]


# Plot results
# ------------

plt.style.use('seaborn')

h1 = plt.gcf()
plt.subplot(2, 2, 1)
plt.grid('both')
plt.plot(np.arange(0, Ns) * integrator_stepsize, x[2, :-1], "C1")
plt.plot(np.arange(0, Ns) * integrator_stepsize, np.concatenate([-1.2*np.ones(int(Ns/2)), 1.2*np.ones(int(Ns/2))], axis=0), "k:")
plt.ylabel('Joint angle [rad]')
plt.gca().legend(["Joint 1", "Reference 1"], loc="lower right", facecolor="white", frameon=True)
plt.ylim([-3, 3])

plt.subplot(2, 2, 2)
plt.grid('both')
plt.plot(np.arange(0, Ns) * integrator_stepsize, x[0, :-1], "C0")
plt.plot(np.arange(0, Ns) * integrator_stepsize, np.concatenate([1.2*np.ones(int(Ns/2)), -1.2*np.ones(int(Ns/2))], axis=0), "k:")
plt.ylabel('Joint angle [rad]')
plt.gca().legend(["Joint 2", "Reference 2"], loc="lower right", facecolor="white", frameon=True)
plt.ylim([-3, 3])

plt.subplot(2, 2, 3)
plt.grid('both')
plt.plot(np.arange(0, Ns) * integrator_stepsize, x[5, :-1], "C0")
plt.plot(np.array([0, Ns]) * integrator_stepsize, [70, 70], 'k:')
plt.gca().legend(["Joint 1", "Limit"], loc="lower right", facecolor="white", frameon=True)
plt.ylabel('Torque [Nm]')
plt.ylim([-20, 80])

plt.subplot(2, 2, 4)
plt.grid('both')
plt.plot(np.arange(0, Ns) * integrator_stepsize, x[4, :-1], "C1")
plt.plot(np.array([0, Ns]) * integrator_stepsize, [70, 70], 'k:')
plt.gca().legend(["Joint 2", "Limit"], loc="lower right", facecolor="white", frameon=True)
plt.ylabel('Torque [Nm]')
plt.ylim([-20, 80])
plt.tight_layout()

plt.figure()
plt.grid('both')
plt.semilogy(np.arange(0, Ns) * integrator_stepsize, solvetime, np.arange(0, Ns) * integrator_stepsize, fevalstime)
plt.legend(["Total solve time", "Ext. function evaluation time"])
plt.xlabel('Time [s]')
plt.ylabel('CPU time [s]')
plt.tight_layout()

# Compute and plot closed loop cost
h3 = plt.figure()
plt.grid('both')
plt.plot(np.arange(0, Ns) * integrator_stepsize, cost)
plt.xlabel('time [s]')
plt.ylabel('cost')
plt.tight_layout()

plt.show()