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

It also shows how to simulate the system directly using the FORCESPRO solver 
object methods.

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-2023
"""

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
    z = np.concatenate((u[:, i], x[:, i]))
    p = np.array(problem['all_parameters'][0:model.npar])
    c, jacc = solver.dynamics(z,p)
    
    # Compute cost
    cost[i] += codeoptions.nlp.integrator.Ts * (float(np.transpose(c) @ Q @ c) + u[:, i] @ R @ u[:, i])
    sim_closed_loop_obj.append(cost[i])

    x[:, i + 1] = np.reshape(c, newshape=(6))

# 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
    z = np.concatenate((u[:, i], x[:, i]))
    p = np.array(problem['all_parameters'][0:model.npar])
    c, jacc = solver.dynamics(z,p)
    
    # Compute cost
    cost[i] += codeoptions.nlp.integrator.Ts * (float(np.transpose(c) @ Q @ c) + u[:, i] @ R @ u[:, i])
    sim_closed_loop_obj.append(cost[i])

    x[:, i + 1] = np.reshape(c, newshape=(6))


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