direct_tests.jl

# Rotations might be better, test TrajectoryOptimization
using ReferenceFrameRotations, DifferentialEquations, Plots
##
function qdot(X, u)
    collect(dquat(Quaternion(X), u))
end
##
q0 = [1.0, 0, 0, 0]
omegab = [1.0; 0; 0]
tmax = 2π
##
ode = ODEProblem((x, p, t) -> qdot(x, p), q0, [0 tmax], omegab)
##
sol = solve(ode, Tsit5())
##
plot(sol)
##
using JuMP
using Ipopt

function euler(f, X, u, dt)
    X + f(X, u)*dt
end

function rk4(f, X, u, dt)
    k1 = dt*f(X     , u)
    k2 = dt*f(X+k1/2, u)
    k3 = dt*f(X+k2/2, u)
    k4 = dt*f(X+k3  , u)
    X + (k1+2k2+2k3+k4)/6
end
## jump works
model = Model(Ipopt.Optimizer)

DCMf = [
    -1 0 0
    0 0 1
    0 1 0
]

# two quaternions represent the same orientation
qf = -convert(Quaternion, DCM(DCMf)) |> collect

N = 50

T = 1
dt = T / (N-1)

tabq = @variable(model, [1:4, 1:N], base_name="q")

tabu = @variable(model, [1:3, 1:N-1], base_name="ω_b")

cost = 0

for i = 1:N-1
    F = rk4((X, u) -> [qdot(X[1:4], u); 1/2 * u' * u], [tabq[:, i]; cost], tabu[:, i], dt)
    @constraint(model, tabq[:, i+1] == F[1:4])
    cost = F[end]
end

#traj constraints
@constraint(model, tabq[:, 1] == q0)
@constraint(model, tabq[:, N] == qf)

# @constraint(model, c[i=1:N], tabq[:, i]' * tabq[:, i] == 1)

@objective(model, Min, cost)

#start values
set_start_value.(tabq, q0)
model
##
optimize!(model)
##
value.(tabu)
## jump can handle multistep rk4
model = Model(Ipopt.Optimizer)
n = 4
m = 3
DCMf = [
    -1 0 0
    0 0 1
    0 1 0
]

# two quaternions represent the same orientation
qf = -convert(Quaternion, DCM(DCMf)) |> collect

function integrate(x, u, dt)
    F = x
    for _ = 1:4
        F = rk4((X, u) -> [qdot(X[1:4], u); 1/2 * u' * u], F, u, dt/4)
    end
    return F
end

"""
    memoize(foo::Function, n_outputs::Int)

Take a function `foo` and return a vector of length `n_outputs`, where element
`i` is a function that returns the equivalent of `foo(x...)[i]`.

To avoid duplication of work, cache the most-recent evaluations of `foo`.
Because `foo_i` is auto-differentiated with ForwardDiff, our cache needs to
work when `x` is a `Float64` and a `ForwardDiff.Dual`.
"""
function memoize(foo::Function, n_outputs::Int)
    last_x, last_f = nothing, nothing
    last_dx, last_dfdx = nothing, nothing
    function foo_i(i, x::T...) where {T<:Real}
        if T == Float64
            if x !== last_x
                last_x, last_f = x, foo(x...)
            end
            return last_f[i]::T
        else
            if x !== last_dx
                last_dx, last_dfdx = x, foo(x...)
            end
            return last_dfdx[i]::T
        end
    end
    return [(x...) -> foo_i(i, x...) for i in 1:n_outputs]
end

int_splat( x1, x2, x3, x4, cost, u1, u2, u3, dt) = integrate([x1, x2, x3, x4, cost], [u1, u2, u3], dt)

mem_int = memoize(int_splat, 5)

@operator(model, op_int1, 9, mem_int[1])
@operator(model, op_int2, 9, mem_int[2])
@operator(model, op_int3, 9, mem_int[3])
@operator(model, op_int4, 9, mem_int[4])
@operator(model, op_int5, 9, mem_int[5])

op_inti = [op_int1, op_int2, op_int3, op_int4, op_int5]

function op_integrate(op_vec, x, u, dt)
    [
        op(x..., u..., dt) for op in op_vec
    ]
end

N = 10

T = 1
dt = T / (N-1)

tabq = @variable(model, [1:4, 1:N], base_name="q")

tabu = @variable(model, [1:3, 1:N-1], base_name="ω_b")

cost = 0
for i = 1:N-1
    F = op_integrate(op_inti, [tabq[:, i]; cost], tabu[:, i], dt)
    @constraint(model, tabq[:, i+1] == F[1:4])
    cost = F[5]
end

#traj constraints
@constraint(model, tabq[:, 1] .== q0) # Commented out otherwise Ipopt says `TOO_FEW_DEGREES_OF_FREEDOM`
@constraint(model, tabq[:, N] .== qf)

# @constraint(model, c[i=1:N], tabq[:, i]' * tabq[:, i] == 1)

# set_start_value.(tabq, q0)

@objective(model, Min, cost)
##
#start values
optimize!(model)
##
value.(tabq)
value.(tabu)