from casadi import *

# look at: https://github.com/casadi/casadi/blob/master/docs/examples/python/vdp_indirect_multiple_shooting.py

T = 3.0
N = 30

x = SX.sym('x')
y = SX.sym('y')
theta = SX.sym('theta')
state = vertcat(x, y, theta)
r = 0.03
R = 0.05
d = 0.02
#r = SX.sym('r')
#R = SX.sym('R')
#d = SX.sym('d')
omegar = SX.sym('omegar')
omegal = SX.sym('omegal')
control = vertcat(omegar, omegal)
f1 = (r/2 * cos(theta) - r*d/(2*R) * sin(theta)) * omegar + (r/2 * cos(theta) + r*d/(2*R) * sin(theta)) * omegal
f2 = (r/2 * sin(theta) + r*d/(2*R) * cos(theta)) * omegar + (r/2 * sin(theta) - r*d/(2*R) * cos(theta)) * omegal
f3 = r/(2*R) * omegar - r/(2*R) * omegal
xdot = vertcat(f1, f2, f3)
f = Function('f', [x,y,theta, omegar, omegal], [f1, f2, f3]) 
print("f = {}".format(f))

L = x**2 + y**2 + 1e-2 * theta**2 + 1e-4 * (omegar**2 + omegal**2)

# Fixed step Runge-Kutta 4 integrator
M = 4 # RK4 steps per interval
DT = T/N/M
print("DT = {}".format(DT))
f = Function('f', [state, control], [xdot, L])
X0 = MX.sym('X0', 3)
U = MX.sym('U', 2)
X = X0
Q = 0
runge_kutta = True
if runge_kutta:
    for j in range(M):
        k1, k1_q = f(X, U)
        k2, k2_q = f(X + DT/2 * k1, U)
        k3, k3_q = f(X + DT/2 * k2, U)
        k4, k4_q = f(X + DT * k3, U)
        X=X+DT/6*(k1 +2*k2 +2*k3 +k4)
        Q = Q + DT/6*(k1_q + 2*k2_q + 2*k3_q + k4_q)
else:
    DT = T/N
    k1, k1_q = f(X, U)
    X = X + DT * k1
    Q = Q + DT * k1_q
F = Function('F', [X0, U], [X, Q],['x0','p'],['xf','qf'])

#F_euler = Function('F_euler', [X0, U], [Xeuler, Qeuler], ['x0', 'p'], ['xf', 'qf'])

Fk = F(x0=[0.2,0.3, 0.0],p=[-1.1, 1.1])
print(Fk['xf'])
print(Fk['qf'])

# Start with an empty NLP
w=[]
w0 = []
lbw = []
ubw = []
J = 0
g=[]
lbg = []
ubg = []

# Formulate the NLP
Xk = MX([1.1, 1.1, 0.0])
for k in range(N):
    # New NLP variable for the control
    U1k = MX.sym('U1_' + str(k), 2)
    #U2k = MX.sym('U2_' + str(k))
    w += [U1k]
    lbw += [-10, -10]
    ubw += [10, 10]
    w0 += [0, 0]

    # Integrate till the end of the interval
    Fk = F(x0=Xk, p=U1k)
    Xk = Fk['xf']
    J=J+Fk['qf']

    # Add inequality constraint
    #g += [Xk[1]]
    #lbg += [-.0]
    #ubg += [inf]

# Create an NLP solver
prob = {'f': J, 'x': vertcat(*w), 'g': vertcat(*g)}
solver = nlpsol('solver', 'ipopt', prob);

# Solve the NLP
sol = solver(x0=w0, lbx=lbw, ubx=ubw, lbg=lbg, ubg=ubg)
w_opt = sol['x']

# Plot the solution
u_opt = w_opt
x_opt = [[1.1, 1.1, -0.0]]
for k in range(N):
    Fk = F(x0=x_opt[-1], p=u_opt[2*k:2*k+2])
    x_opt += [Fk['xf'].full()]
x1_opt = [r[0] for r in x_opt]
x2_opt = [r[1] for r in x_opt]
x3_opt = [r[2] for r in x_opt]

tgrid = [T/N*k for k in range(N+1)]
import matplotlib.pyplot as plt
plt.figure(1)
plt.clf()
plt.plot(tgrid, x1_opt, '--')
plt.plot(tgrid, x2_opt, '-')
plt.plot(tgrid, x3_opt, '*')
#plt.step(tgrid, vertcat(DM.nan(1), u_opt), '-.')
plt.xlabel('t')
plt.legend(['x1','x2','x3','u'])
plt.grid()
#plt.show()

# alternative solution using multiple shooting (way faster!)
opti = Opti() # Optimization problem

# ---- decision variables ---------
X = opti.variable(3,N+1) # state trajectory
Q = opti.variable(1,N+1) # state trajectory
posx   = X[0,:]
posy   = X[1,:]
angle = X[2,:]
U = opti.variable(2,N)   # control trajectory (throttle)
#T = opti.variable()      # final time

# ---- objective          ---------
#opti.minimize(T) # race in minimal time

# ---- dynamic constraints --------
#f = lambda x,u: vertcat(f1, f2, f3) # dx/dt = f(x,u)

dt = T/N # length of a control interval
for k in range(N): # loop over control intervals
   # Runge-Kutta 4 integration
   k1, k1_q = f(X[:,k],         U[:,k])
   k2, k2_q = f(X[:,k]+dt/2*k1, U[:,k])
   k3, k3_q = f(X[:,k]+dt/2*k2, U[:,k])
   k4, k4_q = f(X[:,k]+dt*k3,   U[:,k])
   x_next = X[:,k] + dt/6*(k1+2*k2+2*k3+k4)
   q_next = Q[:,k] + dt/6*(k1_q + 2 * k2_q + 2 * k3_q + k4_q)
   opti.subject_to(X[:,k+1]==x_next) # close the gaps
   opti.subject_to(Q[:,k+1]==q_next) # close the gaps
opti.minimize(Q[:,N])

# ---- path constraints -----------
#limit = lambda pos: 1-sin(2*pi*pos)/2
#opti.subject_to(speed<=limit(pos))   # track speed limit
opti.subject_to(opti.bounded(-10,U,10)) # control is limited

# ---- boundary conditions --------
opti.subject_to(posx[0]==1.10)   # start at position 0 ...
opti.subject_to(posy[0]==1.10) # ... from stand-still
opti.subject_to(angle[0]==0.0)  # finish line at position 1
#opti.subject_to(speed[-1]==0)  # .. with speed 0
opti.subject_to(Q[:,0]==0.0)

# ---- misc. constraints  ----------
#opti.subject_to(X[1,:]>=0) # Time must be positive
#opti.subject_to(X[2,:]<=4) # Time must be positive
#opti.subject_to(X[2,:]>=-2) # Time must be positive

r = 0.25
p = (0.5, 0.5)
for k in range(N):
    opti.subject_to((X[0,k]-p[0])**2 + (X[1,k]-p[1])**2  > r**2)
    pass


# ---- initial values for solver ---
#opti.set_initial(speed, 1)
#opti.set_initial(T, 1)

# ---- solve NLP              ------
opti.solver("ipopt") # set numerical backend
sol = opti.solve()   # actual solve

#x0 = sol.value(opti.x)
#lam_g0 = sol.value(opti.lam_g)
#opti.set_initial(opti.lam_g, lam_g0)
#opti.set_initial(opti.x, x0)
#opti.solve()

from pylab import plot, step, figure, legend, show, spy

plot(sol.value(posx),label="posx")
plot(sol.value(posy),label="posy")
plot(sol.value(angle),label="angle")

plt.figure()
plot(sol.value(posx), sol.value(posy))
ax = plt.gca()
circle = plt.Circle(p, r)
ax.add_artist(circle)
#plot(limit(sol.value(pos)),'r--',label="speed limit")
#step(range(N),sol.value(U),'k',label="throttle")
legend(loc="upper left")
plt.show()
pass
# linearization
# A = zeros(states.shape[0], states.shape[0])
# for i in range(f.shape[0]):
#     for j in range(states.shape[0]):
#         A[i,j] = diff(f[i,0], states[j])
# Alin = A.subs([(theta,0), (omegar,0), (omegal,0)])
# print("A = {}".format(Alin))
# B = zeros(f.shape[0], controls.shape[0])
# for i in range(f.shape[0]):
#     for j in range(controls.shape[0]):
#         B[i,j] = diff(f[i,0], controls[j])
# print("B = {}".format(B))
# dfdtheta = diff(f, theta)
#print(dfdtheta.doit())

# takeaway: linearization is not helpful, because the linearized system is not stabilizable
# -> alternative: use nonlinear control method