text:
"""
spin21.jl - derivative robustness for the empty problem with simultaneous errors
"""

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const TO = TrajectoryOptimization

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin21"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# Define the optimization.
const CONTROL_COUNT = 1
const PN_STEPS = 2
const MAX_COST_VALUE = 1e12
const ILQR_DJ_TOL = 1e-4

# Define the problem.
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * STATE_SIZE_ISO + 1 * CONTROL_COUNT
const INITIAL_STATE1 = [1., 0, 0, 0]
const INITIAL_STATE2 = [0., 1, 0, 0]
# state indices
const STATE1_IDX = 1:STATE_SIZE_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + STATE_SIZE_ISO
const CONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
# const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const D1STATE1_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + STATE_SIZE_ISO
const D1STATE2_IDX = D1STATE1_IDX[end] + 1:D1STATE1_IDX[end] + STATE_SIZE_ISO
const D2STATE1_IDX = D1STATE2_IDX[end] + 1:D1STATE2_IDX[end] + STATE_SIZE_ISO
const D2STATE2_IDX = D2STATE1_IDX[end] + 1:D2STATE1_IDX[end] + STATE_SIZE_ISO
# control indices
const DCONTROLS_IDX = 1:CONTROL_COUNT

# Specify logging.
const VERBOSE = true
const SAVE = true

# Misc. constants
EMPTY_V = []

# Define the dynamics.
struct Model{DO} <: AbstractModel
    Model(DO::Int64=0) = new{DO}()
end
RobotDynamics.state_dim(::Model{DO}) where DO = (
    ASTATE_SIZE_BASE + DO * STATE_COUNT * STATE_SIZE_ISO
)
RobotDynamics.control_dim(::Model{DO}) where DO = CONTROL_COUNT

function RobotDynamics.dynamics(model::Model{DO}, astate::StaticVector,
                                acontrols::StaticVector, time::Real) where DO
    negi_h = (
        astate[CONTROLS_IDX][1] * NEGI_H1_ISO
    )
    delta_state1 = negi_h * astate[STATE1_IDX]
    delta_state2 = negi_h * astate[STATE2_IDX]
    delta_control = acontrols[DCONTROLS_IDX]

    if DO == 1
        delta_d11state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d11state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_d12state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d12state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
            delta_d11state1;
            delta_d11state2;
            delta_d12state1;
            delta_d12state2;
        ]
    elseif DO == 2
        delta_d1state1 = NEGI_H0_ISO * astate[STATE1_IDX] + negi_h * astate[D1STATE1_IDX]
        delta_d1state2 = NEGI_H0_ISO * astate[STATE2_IDX] + negi_h * astate[D1STATE2_IDX]
        delta_d2state1 = 2 * NEGI_H0_ISO * astate[D1STATE1_IDX] + negi_h * astate[D2STATE1_IDX]
        delta_d2state2 = 2 * NEGI_H0_ISO * astate[D1STATE2_IDX] + negi_h * astate[D2STATE2_IDX]
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
            delta_d1state1;
            delta_d1state2;
            delta_d2state1;
            delta_d2state2;
        ]
    else
        delta_astate = [
            delta_state1;
            delta_state2;
            delta_control;
        ]
    end

    return delta_astate
end

function run_traj(;gate_type=xpi, evolution_time=10., solver_type=alilqr,
                  postsample=false, initial_save_file_path=nothing,
                  initial_save_type=jl, sqrtbp=false, derivative_order=0,
                  integrator_type=rk6, max_penalty=1e11, qs=nothing,
                  smoke_test=false, dt=2.5e-3, al_kickout_tol=1e-9,
                  constraint_tol=1e-10)
    model = Model(derivative_order)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.
    x0 = SVector{n}([
        INITIAL_STATE1;
        INITIAL_STATE2;
        zeros(1 * CONTROL_COUNT);
        zeros(derivative_order * STATE_COUNT * STATE_SIZE_ISO);
    ])
    if gate_type == xpi
        target_state1 = Array(XPI_ISO_1)
        target_state2 = Array(XPI_ISO_2)
    end
    xf = SVector{n}([
        target_state1;
        target_state2;
        zeros(1 * CONTROL_COUNT);
        zeros(derivative_order * STATE_COUNT * STATE_SIZE_ISO);
    ])

    # control amplitude constraint
    x_max = SVector{n}([
        fill(Inf, STATE_COUNT * STATE_SIZE_ISO);
        fill(MAX_CONTROL_NORM_0, 1); # control
        fill(Inf, derivative_order * STATE_COUNT * STATE_SIZE_ISO)
    ])
    x_min = SVector{n}([
        fill(-Inf, STATE_COUNT * STATE_SIZE_ISO);
        fill(-MAX_CONTROL_NORM_0, 1); # control
        fill(-Inf, derivative_order * STATE_COUNT * STATE_SIZE_ISO)
    ])

    N = Int(floor(evolution_time / dt)) + 1
    U0 = [SVector{m}(
        fill(1e-4, CONTROL_COUNT)
    ) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    if isnothing(qs)
        qs = ones(6)
    end
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * STATE_SIZE_ISO); # state1, state2
        fill(qs[2], 1); # control
        fill(qs[3], eval(:($derivative_order >= 1 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d1state<1,2>
        fill(qs[4], eval(:($derivative_order >= 2 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d2state<1,2>
        fill(qs[5], eval(:($derivative_order >= 3 ? $STATE_COUNT * $STATE_SIZE_ISO : 0))); # d3state<1,2>
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT);
    ]))
    obj = LQRObjective(Q, R, Qf, xf, N)

    # Must satisfy control amplitude bound.
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # Must reach target state. Must have zero net flux.
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX])

    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, target_astate_constraint, N:N);

    # Instantiate problem and solve.
    prob = Problem{IT_RDI[integrator_type]}(model, obj, constraints, x0, xf, Z, N, t0, evolution_time)
    opts = SolverOptions(verbose=VERBOSE)
    solver = AugmentedLagrangianSolver(prob, opts)
    if solver_type == alilqr
        solver = AugmentedLagrangianSolver(prob, opts)
        solver.solver_uncon.opts.square_root = sqrtbp
        solver.solver_uncon.opts.max_cost_value = MAX_COST_VALUE
        solver.opts.constraint_tolerance = al_kickout_tol
        solver.opts.constraint_tolerance_intermediate = al_kickout_tol
        solver.opts.cost_tolerance_intermediate = ILQR_DJ_TOL
        solver.opts.penalty_max = max_penalty
        if smoke_test
            solver.opts.iterations = 1
            solver.solver_uncon.opts.iterations = 1
        end
    elseif solver_type == altro
        solver = ALTROSolver(prob, opts)
        solver.opts.constraint_tolerance = constraint_tol
        solver.solver_al.solver_uncon.opts.square_root = sqrtbp
        solver.solver_al.solver_uncon.opts.max_cost_value = MAX_COST_VALUE
        solver.solver_al.opts.constraint_tolerance = al_kickout_tol
        solver.solver_al.opts.constraint_tolerance_intermediate = al_kickout_tol
        solver.solver_al.opts.cost_tolerance_intermediate = ILQR_DJ_TOL
        solver.solver_al.opts.penalty_max = max_penalty
        solver.solver_pn.opts.constraint_tolerance = constraint_tol
        solver.solver_pn.opts.n_steps = PN_STEPS
        if smoke_test
            solver.solver_al.opts.iterations = 1
            solver.solver_al.solver_uncon.opts.iterations = 1
            solver.solver_pn.opts.n_steps = 1
        end
    end
    Altro.solve!(solver)

    # Post-process.
    acontrols_raw = controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    cmax = TrajectoryOptimization.max_violation(solver)
    cmax_info = TrajectoryOptimization.findmax_violation(get_constraints(solver))
    iterations_ = iterations(solver)

    # Save.
    if SAVE
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            write(save_file, "acontrols", acontrols_arr)
            write(save_file, "controls_idx", cidx_arr)
            write(save_file, "evolution_time", evolution_time)
            write(save_file, "astates", astates_arr)
            write(save_file, "Q", Q_arr)
            write(save_file, "Qf", Qf_arr)
            write(save_file, "R", R_arr)
            write(save_file, "cmax", cmax)
            write(save_file, "cmax_info", cmax_info)
            write(save_file, "dt", dt)
            write(save_file, "derivative_order", derivative_order)
            write(save_file, "solver_type", Integer(solver_type))
            write(save_file, "sqrtbp", Integer(sqrtbp))
            write(save_file, "max_penalty", max_penalty)
            write(save_file, "constraint_tol", constraint_tol)
            write(save_file, "al_kickout_tol", al_kickout_tol)
            write(save_file, "ilqr_dj_tol", ILQR_DJ_TOL)
            write(save_file, "integrator_type", Integer(integrator_type))
            write(save_file, "gate_type", Integer(gate_type))
            write(save_file, "save_type", Integer(jl))
            write(save_file, "iterations", iterations_)
            write(save_file, "max_cost_value", MAX_COST_VALUE)
        end
    end
end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc


text:
## Before we got started
# I added the Images package via
# typing "]" then "add Images"
# alternatively, we could use
# using Pkg; Pkg.add("Images")

## Let's try and make Wikipedia's Julia set figure.
using Images
using Plots

## It says the seed coordinates are:
seeds = (-0.512511498387847167, 0.521295573094847167)

## Let's work with the Julia set.
# What we need to do is to check if, when starting
# from the complex number z=0, and iterating the map
# f(z) = z^2 + c
# for a large number of times (say 256)
# if the value of abs(f(z)) ever goes above 2,
# then we stop when that happens and return the value.
c = 0.25
# if we start at z = 0, then we test if z=0 is in the set.
z = 0.0+0.0im # this is how to make an imaginary number
f(z) = z^2 + c
@show f(z)
@show f(f(z))
@show f(f(f(z)))

##
zi = z
for i=1:10
    global zi
    zi = f(zi)
    @show zi
end
##
"""
`julia_check` is a subroutine useful for working with
the Julia set. We are going take as input two values:
z0 and c. The goal is to test whether or not the value
z0 seems to be in the Julia set for
    f(z) = z^2 + c
by starting from z0 = the input z0 and iterating
z1 = f(z0)
z2 = f(z1)
...
and so on.
Now, if at any point, we find abs(z) >= 2, then we know
that the point is NOT in the Julia set. In which case,
we return the iteration number where we detect this.
We run this for maxiter steps.
"""
function julia_check(z0, c; maxiter::Int=256)
    z = z0
    for i=1:maxiter
        z = z^2 + c
        if abs(z) >= 2
            return i
        end
    end
    return maxiter
end
julia_check(-1, 0.25)

## Let's look at starting positions z0 that range
## from [-1.75 to 1.75] in the real component
## and [-1.75 to 1.75] in the complex component
npts = 10
Z = zeros(npoints,npoints) # this will represent our Images
real_values = range(start=-1.75,stop=1.75,length=npts)
complex_values = range(start=-1.75,stop=1.75,length=npts)
for i=1:npts
    for j=1:npts
        z0 = real_values[i] + complex_values[j]*1.0im
        Z[j,i] = julia_check(z0, 0.25)
    end
end
heatmap(Z)

## The fix!
npts = 10
Z = zeros(npts,npts) # this will represent our Images
real_values = range(-1.75,stop=1.75,length=npts)
complex_values = range(-1.75,stop=1.75,length=npts)
for i=1:npts
    for j=1:npts
        z0 = real_values[i] + complex_values[j]*1.0im
        Z[j,i] = julia_check(z0, 0.25)
    end
end
heatmap(Z)

## Let's turn this into a function!
function julia_set(real_values, complex_values)
    Z = zeros(length(complex_values), length(real_values))
    for i=1:length(real_values)
        for j=1:length(complex_values)
            z0 = real_values[i] + complex_values[j]*1.0im
            Z[j,i] = julia_check(z0, 0.25)
        end
    end
    return Z
end
npts = 500
real_values = range(-1.75,stop=1.75,length=npts)
complex_values = range(-1.75,stop=1.75,length=npts)
heatmap(julia_set(real_values, complex_values))
##
function julia_set(real_values, complex_values; c=0.25)
    Z = zeros(length(complex_values), length(real_values))
    for i=1:length(real_values)
        for j=1:length(complex_values)
            z0 = real_values[i] + complex_values[j]*1.0im
            Z[j,i] = julia_check(z0, c)
        end
    end
    return Z
end
npts = 500
real_values = range(-1.75,stop=1.75,length=npts)
complex_values = range(-1.75,stop=1.75,length=npts)
heatmap(julia_set(real_values, complex_values; c = seeds[1] + seeds[2]*1.0im))
## Let's increase the resolution.
npts = 1000
real_values = range(-1.75,stop=1.75,length=npts)
complex_values = range(-1.75,stop=1.75,length=npts)
Z = julia_set(real_values, complex_values; c = seeds[1] + seeds[2]*1.0im)
## And then try and make our plot.
heatmap(log10.(Z))
## TODO or IDEA - Play around the with the colors to see if you find something you like better
## Show this like an image.
image = Gray.(Z/maximum(Z))
image

## to save this picture, we can use
using FileIO
using ImageMagick
save("juliaset.png", Gray.(Z/maximum(Z)))

ext: jl


text:
mutable struct pose_pose_nh_t  <: LCMType
  utime::Int64         # unix time associated with this factor

  node_1_utime::Int64  # pose 1 (source) utime
  node_1_id::Int64     # pose 1 id

  node_2_utime::Int64  # pose 2 (destination) utime
  node_2_id::Int64     # pose 2 id

  mean_dim::Int32       # default will be 7 (3+4)
  mean::Vector{Float64} # default will be [x y z qw qx qy qz]

  covar_dim::Int32
  covar::Vector{Float64} # usually x y z yaw pitch roll (RotXYZ)
  confidence::Float64    # the confidence value for the "valid factor" hypothesis
end

@lcmtypesetup(pose_pose_nh_t,
  mean => (mean_dim,),
  covar => (covar_dim,)
)


text:
# # Color cycling
# We will try to fake a subtle animation of the water in this public domain image:
using Images
using DitherPunk
using Clustering
using IndirectArrays

img = load("../../assets/waterfall.png")

# When Clustering.jl is loaded, we can import DitherPunk's internal function `get_colorscheme`
# to obtain a color scheme of size 16.
cs = DitherPunk.get_colorscheme(img, 16)
cs = sort(cs; by=c -> -c.l)
# The first three colors are white-ish and correspond to the water.

# Let's look at the dithered result using `color_error_multiplier=1.0`
# for a more visible dithering pattern:
d1 = dither(img, Bayer(;color_error_multiplier=1.0), cs)

# We create two more frames by cycling the first three colors
d2 = IndirectArray(d1.index, cs[[3, 1, 2, 4:end...]])
d3 = IndirectArray(d1.index, cs[[2, 3, 1, 4:end...]]);

# And finally create a gif out of the modified images:
ds = cat(d1, d2, d3; dims=3)
save("waterfall.gif", ds; fps=5)

# The result should look like this:
#
# ![](https://i.imgur.com/nFpH89N.gif)
#
# We really nailed that Web 1.0 look!
#
# To see how good color cycling can look when it is hand-drawn by a professional, take a look at
# [Mark Ferrari's amazing work](http://www.effectgames.com/demos/canvascycle/?sound=1).