#%%
import os
import shutil
import numpy as np
import casadi as ca
import matplotlib.pyplot as plt

# Runge Kutta order 4 integration method
def rk4step(ode, h, x, u):

    k1 = ode(x, u)
    k2 = ode(x + h/2 * k1, u)
    k3 = ode(x + h/2 * k2, u)
    k4 = ode(x + h * k3, u)
    x_next = x + h/6 * (k1 + 2*k2 + 2*k3 + k4)
    return x_next

# choose different options for ODE
# linear dynamics (nonlinear_dyn = False): mass-spring systems
# nonlinear dynamics (nonlinear_dyn = True): pendulum
nonlinear_dyn = True            # set to True for nonlinear dynamics
control_penalty_steep = False   # set to True for steep control penalty (a^4)


# problem parameters
T = 4               # continous time horizon
N = 30              # discrete time horizon
dt = T/N            # time step / control discretization
N_int = 10          # number of integrator steps per time step
ns = 2              # number of states
na = 1              # number of actions
k_spring = 1        # spring constant
a_grav = 1          # gravity acceleration
L = 1               # pendulum length

s0bar = np.array([2, 0])        # initial state
sNbar = np.array([0, 0])        # target state

# define ode
s = ca.SX.sym('x', ns)          # symbol for state
a = ca.SX.sym('u', na)          # symbol for action
p = s[0]                        # position
v = s[1]                        # velocity

pdot = v                        # time derivative of position is velocity
if nonlinear_dyn:
    vdot = a_grav * ca.sin((a-p / L))
else:
    vdot = k_spring * (a - p)

sdot = ca.veccat(pdot, vdot )               # time derivative of state / ode
ode = ca.Function('ode', [s, a], [ sdot ])  # as function

# define integrator of dynamics
# Runge-Kutta 4th order integrator
s_next = rk4step(ode, dt/N_int, s, a)                   # as casadi expression
rk4_step = ca.Function('rk4_step', [s, a], [s_next])    # as casadi function
# apply function rk4_step() N_int times on itself
integrator = rk4_step.mapaccum(N_int)       

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

# decision variables
a_all = opti.variable(na, N)       # all controls
s_all = opti.variable(ns, N+1)     # all states

# objective function
obj = 0
for i in range(N):
    obj += a_all[:,i]**2
    obj += .01 * s_all[0,i]**2
if control_penalty_steep:
    for i in range(N):
        obj += a_all[:,i]**4
# passing the objective to opti
opti.minimize(obj)


# constraints
constr = []
opti.subject_to( s_all[:,0] - s0bar == 0)           # initial condition
for i in range(N):                          # dynamics        
    opti.subject_to(s_all[:,i+1] - integrator(s_all[:,i], a_all[:,i])[:,-1] == 0)
opti.subject_to( s_all[:,N] - sNbar == 0)           # terminal constraint


# Setting solver to ipopt and solving the problem:
opti.solver('ipopt')
# opti.solver('sqpmethod')
sol = opti.solve()

A_all = sol.value(a_all)
S_all = sol.value(s_all)

# plot the result
fig, ax = plt.subplots()
ax.plot(A_all, '--or')
ax.plot(S_all[0,:], 'b')
ax.plot(S_all[1,:], 'b')
ax.set_title(r"Optimal solution hanging chain (without extra constraints)")