using OptimalControl
using NLPModelsIpopt

ocp = @def begin
    t ∈ [0, 1], time
    x ∈ R², state
    u ∈ R, control
    x(0) == [ -1, 0 ]
    x(1) == [ 0, 0 ]
    ẋ(t) == [ x₂(t), 0 ] + [ 0, u(t) ] # this compiles: ẋ(t) == [ x₂(t), u(t) ]
    ∫( 0.5u(t)^2 ) → min
end

sol = solve(ocp)

ẋ(t) == [ x₂(t), u(t) ]
# or 
ẋ(t) == [ x₂(t), 0 ] + [ x₂(t), u(t) ]
# or
ẋ(t) == [ x₂(t), 0 ] + [ x₁(t), u(t) ]
# or 
ẋ(t) == [ x₂(t), 0 ] + u(t) * [ 0, 1 ]

module Goddard3D

using OptimalControl
using NLPModelsIpopt
#using OrdinaryDiffEq  # to get the Flow function from OptimalControl
#using NonlinearSolve  # interface to NLE solvers
#using MINPACK         # NLE solver: use to solve the shooting equation
using Plots           # to plot the solution

t0 = 0                              # initial time
r0 = [0.999949994, 0.0001, 0.01]    # initial altitude
v0 = [0, 0, 0]                      # initial speed, v = 1 ==> 7910 m/s (orbital velocity at sea level)
m0 = 1                              # initial mass
x0 = [r0..., v0..., m0]
rf = [1.01, 0, 0]                   # r[1] = 1 ==> 6378 km

ocp = @def begin # definition of the optimal control problem

    tf ∈ R, variable
    t ∈ [ t0, tf ], time
    x ∈ R⁷, state
    u ∈ R³, control

    r = x[1:3]
    v = x[4:6]
    m = x[7]

    x(t0) == x0
    r(tf) == rf

    u(t)' * u(t) ≤ 1
    r(t)' * r(t) ≥ 1

    ẋ(t) == F(x(t), u(t))

    m(tf) → max

end;

# Dynamics
const Cd = 310
const Tmax = 3.5
const β = 500
const b = 7

mysqrt(x) = sqrt(x + 1e-6)

F(x, u) = begin

    r = x[1:3]
    v = x[4:6]
    m = x[7]

    rnorm2 = r' * r
    rnorm = mysqrt(rnorm2)
    rnormalized = (1/rnorm) * r
    vnorm = mysqrt(v' * v)
    D = ( Cd * vnorm * exp(-β*(rnorm - 1)) ) * v # Drag force

    ṙ = v
    v̇ = -D/m - (1/rnorm2)*rnormalized + (Tmax / m) * u
    ṁ = -b*Tmax*mysqrt(u' * u)

    return [ṙ..., v̇..., ṁ]

end    

# DIRECT METHOD

function run_direct()
    direct_sol = solve(ocp; grid_size=100)
    plt = plot(direct_sol, size=(900, 900))
end

# INDIRECT SHOOTING METHOD

end # module

Goddard3D.run_direct()

julia> include("goddard.jl");
WARNING: replacing module Goddard3D.
This is Ipopt version 3.14.16, running with linear solver MUMPS 5.7.3.

Number of nonzeros in equality constraint Jacobian...:     7510
Number of nonzeros in inequality constraint Jacobian.:      606
Number of nonzeros in Lagrangian Hessian.............:     4747

Total number of variables............................:     1011
                     variables with only lower bounds:        0
                variables with lower and upper bounds:        0
                     variables with only upper bounds:        0
Total number of equality constraints.................:      710
Total number of inequality constraints...............:      202
        inequality constraints with only lower bounds:      101
   inequality constraints with lower and upper bounds:        0
        inequality constraints with only upper bounds:      101

iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls
   0 -1.0000000e-01 7.98e+137 1.00e+00   0.0 0.00e+00    -  0.00e+00 0.00e+00   0
MUMPS returned INFO(1) = -9 and requires more memory, reallocating.  Attempt 1
  Increasing icntl[13] from 1000 to 2000.
MUMPS returned INFO(1) = -9 and requires more memory, reallocating.  Attempt 2
  Increasing icntl[13] from 2000 to 4000.
MUMPS returned INFO(1) = -9 and requires more memory, reallocating.  Attempt 3
  Increasing icntl[13] from 4000 to 8000.
WARNING: Problem in step computation; switching to emergency mode.
   1r-1.0000000e-01 7.98e+137 9.99e+02 129.9 0.00e+00  20.0 0.00e+00 0.00e+00R  1
Restoration phase failed with unexpected solverreturn status 9

Number of Iterations....: 1

                                   (scaled)                 (unscaled)
Objective...............:  -1.0000000000000001e-01   -1.0000000000000001e-01
Dual infeasibility......:   1.0000000000000000e+00    1.0000000000000000e+00
Constraint violation....:  7.9803725102132858e+129   7.9803725102132855e+137
Variable bound violation:   0.0000000000000000e+00    0.0000000000000000e+00
Complementarity.........:   9.7000000999999991e-01    9.7000000999999991e-01
Overall NLP error.......:  7.9803725102132858e+129   7.9803725102132855e+137


Number of objective function evaluations             = 2
Number of objective gradient evaluations             = 2
Number of equality constraint evaluations            = 2
Number of inequality constraint evaluations          = 2
Number of equality constraint Jacobian evaluations   = 2
Number of inequality constraint Jacobian evaluations = 2
Number of Lagrangian Hessian evaluations             = 1
Total seconds in IPOPT                               = 1.988

EXIT: Restoration Failed!