text:
using Pkg
#  If NetABM is not available then it should be installed via
#  Uncomment if not installed
] add https://github.com/loopulos/NetABM.jl#develop
#  ] add https://github.com/ollin18/NetABM.jl

function useit(list::Array{Symbol})
    installed = [key for key in keys(Pkg.installed())]
    strpackages = @. string(list)
    uninstalled = setdiff(strpackages,installed)

    map(Pkg.add,uninstalled)
    for package ∈ list
        @eval using $package
    end
end

thep = [:LightGraphs, :StatsBase, :Plots, :GraphPlot, :Statistics, :LaTeXStrings, :Bootstrap, :Distributions, :NetABM]
useit(thep)

################################################################
################################################################
################################################################
################################################################
################################################################
# The same as before but instead of looking at neighbor infections
# we look use a mean cooperation effect by distance

global num_agents = 100
global n_boot = 1000
global cil = 0.95
global j

global adoptstep = (1,[0.1, 0.1])
global maxI = Array{Float64}(undef,0)
global sumI = []

numnei = Int.(ceil.(exp.(range(1,stop=4,length=11))))
numnei = numnei[5]
icol="#E83F6F"
scol="#05668D"
rcol="#90A959"

@time let
global plt = plot(
        xlabel = L"t",
        ylabel = L"N",
        guidefontsize = 16,
        tickfontsize = 10,
        dpi=300
    )

    #  for copera in 0.0:0.1:1
    for copera in [0.5]
        #  for adap in 0.0:0.1:1
        for adap in [0.5]
            global Smean = Array{Float64}(undef,0)
            global Slow = Array{Float64}(undef,0)
            global Shigh = Array{Float64}(undef,0)
            global Imean = Array{Float64}(undef,0)
            global Ilow = Array{Float64}(undef,0)
            global Ihigh = Array{Float64}(undef,0)
            global Rmean = Array{Float64}(undef,0)
            global Rlow = Array{Float64}(undef,0)
            global Rhigh = Array{Float64}(undef,0)
            global Iall = Array{Array}(undef,0)
            global Sall = Array{Array}(undef,0)
            global Rall = Array{Array}(undef,0)

            global RAmean = Array{Float64}(undef,0)
            global RAlow = Array{Float64}(undef,0)
            global RAhigh = Array{Float64}(undef,0)
            global LRAmean = Array{Float64}(undef,0)
            global LRAlow = Array{Float64}(undef,0)
            global LRAhigh = Array{Float64}(undef,0)
            global LRTmean = Array{Float64}(undef,0)
            global LRTlow = Array{Float64}(undef,0)
            global LRThigh = Array{Float64}(undef,0)
            global RTmean = Array{Float64}(undef,0)
            global RTlow = Array{Float64}(undef,0)
            global RThigh = Array{Float64}(undef,0)
            global RAall = Array{Array}(undef,0)
            global LRAall = Array{Array}(undef,0)
            global LRTall = Array{Array}(undef,0)
            global RTall = Array{Array}(undef,0)

            for a in 1:10
                g = barabasi_albert(num_agents,numnei)
                agents = [Agent(i) for i = 1:num_agents];
                init_demographics!(agents;states=["S","I"],initial=[1.0,0.0]);
                pat_0 = rand(1:nv(g))
                agents[pat_0].state = "I"
                set_coop_agents!(agents;p_cop=copera);
                set_adapt_agents!(agents;p_cop=adap);
                set_coop_effet!(agents)
                map(x -> assign_contacts!(g,x), agents);
                global num_S = []
                global num_I = []
                global num_R = []

                global num_ra = []
                global num_lra = []
                global num_lrt = []
                global num_rt = []
                for i in 1:100
                    push!(num_S, ([ag.state for ag in agents if ag.state=="S"] |> length)/num_agents)
                    push!(num_I, ([ag.state for ag in agents if ag.state=="I"] |> length)/num_agents)
                    push!(num_R, ([ag.state for ag in agents if ag.state=="R"] |> length)/num_agents)

                    push!(num_ra,  mean([ag.coop_effect for ag in agents if ag.attitude=="ra" && !ag.adapter]))
                    push!(num_lra,  mean([ag.coop_effect for ag in agents if (ag.attitude=="rt" && ag.adapter)]))
                    push!(num_lrt,  mean([ag.coop_effect for ag in agents if (ag.attitude=="ra" && ag.adapter)]))
                    push!(num_rt,  mean([ag.coop_effect for ag in agents if ag.attitude=="rt" && !ag.adapter]))

                    next_state!(agents;fun=SI_attitude!,inf_prob=0.1, rec_prob=0.03, R=true);
                    update_state!(agents);
                    update_effect_given_distance_coop!(agents,g,adoptstep[1],1,adoptstep[2]);
                end
                if maximum(num_I) > 0.04
                    plot!(plt, 1:100, num_I, alpha=0.15, color = icol, lw = 0.3, label="")
                    plot!(plt, 1:100, num_S, alpha=0.15, color = scol, lw = 0.3, label="")
                    plot!(plt, 1:100, num_R, alpha=0.15, color = rcol, lw = 0.3, label="")
                    push!(Iall,num_I)
                    push!(Sall,num_S)
                    push!(Rall,num_R)

                    push!(RAall,num_ra)
                    push!(LRAall,num_lra)
                    push!(LRTall,num_lrt)
                    push!(RTall,num_rt)
                end
                GC.gc()
            end
            for step in eachindex(Iall[1])
                bs1 = bootstrap(mean,getindex.(Iall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(Iall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(Imean,bci1[1][1])
                push!(Ilow,bci2[1][2])
                push!(Ihigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(Sall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(Sall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(Smean,bci1[1][1])
                push!(Slow,bci2[1][2])
                push!(Shigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(Rall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(Rall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(Rmean,bci1[1][1])
                push!(Rlow,bci2[1][2])
                push!(Rhigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(RAall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(RAall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(RAmean,bci1[1][1])
                push!(RAlow,bci2[1][2])
                push!(RAhigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(LRAall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(LRAall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(LRAmean,bci1[1][1])
                push!(LRAlow,bci2[1][2])
                push!(LRAhigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(LRTall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(LRTall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(LRTmean,bci1[1][1])
                push!(LRTlow,bci2[1][2])
                push!(LRThigh,bci2[1][3])

                bs1 = bootstrap(mean,getindex.(RTall,step),BasicSampling(n_boot))
                bs2 = bootstrap(std,getindex.(RTall,step),BasicSampling(n_boot))
                bci1 = confint(bs1, BasicConfInt(cil))
                bci2 = confint(bs2, BasicConfInt(cil))
                push!(RTmean,bci1[1][1])
                push!(RTlow,bci2[1][2])
                push!(RThigh,bci2[1][3])
            end
        end
        GC.gc()
    end
end

plot!(plt, 1:100, Imean, color = icol,   lab = "I", lw = 1.25)
plot!(plt, 1:100, Smean, color = scol,   lab = "S", lw = 1.25)
plot!(plt, 1:100, Rmean, color = rcol,   lab = "R", lw = 1.25)
plt

png(plt,"figs/cop05_adap05_k10.png")
savefig(plt,"figs/cop05_adap05_k10.pdf")

plt2 = plot(
        xlabel = L"t",
        ylabel = L"N",
        guidefontsize = 16,
        tickfontsize = 10,
        dpi=300
    )

plot!(plt2, 1:100, Imean, ribbon = (Ilow, Ihigh), fillalpha=0.35, color = "#E83F6F",   lab = "I", lw = 1.25)
plot!(plt2, 1:100, Smean, ribbon = (Slow, Shigh), fillalpha=0.35, color = "#05668D",   lab = "S", lw = 1.25)
plot!(plt2, 1:100, Rmean, ribbon = (Rlow, Rhigh), fillalpha=0.35, color = "#90A959",   lab = "R", lw = 1.25)

png(plt2,"figs/cop05_adap05_k10_95.png")
savefig(plt2,"figs/cop05_adap05_k10_95.pdf")
# plot!(plt, tau, num_R, color = :green, lab = "R", lw = 1.25)

plt3 = plot(
        xlabel = L"t",
        ylabel = L"N",
        guidefontsize = 16,
        tickfontsize = 10,
        dpi=300,
        ylim=[0,1]
    )

plot!(plt3, 1:100, RAmean, ribbon = (RAlow, RAhigh), fillalpha=0.35,    lab = "RA", lw = 1.25)
plot!(plt3, 1:100, LRAmean, ribbon = (LRAlow, LRAhigh), fillalpha=0.35,   lab = "LRA", lw = 1.25)
plot!(plt3, 1:100, LRTmean, ribbon = (LRTlow, LRThigh), fillalpha=0.35,   lab = "LRT", lw = 1.25)
plot!(plt3, 1:100, RTmean, ribbon = (RTlow, RThigh), fillalpha=0.35,   lab = "RT", lw = 1.25)

for i in 1:11
    j = numnei[i]
    z1 = z[:,:,i]
    #  plt = heatmap!(xs, ys, z1, title=L"\langle k\rangle =$i", aspect_ratio = 1, clim=(0,1), c = :algae,dpi=300)
    plt = heatmap!(xs, ys, z1, title= h, aspect_ratio = 1, clim=(0,1), c = :algae,dpi=300)
    png(plt,"figs/heatmap_scaled_maximum_coop_neigh_$i.png")
    savefig(plt,"figs/heatmap_scaled_maximum_coop_neigh_$i.pdf")
end


text:
using NPFinancial
using Test

@testset "Tests" begin

    @test NPFinancial._convert_when(:begin) == 1
    @test NPFinancial._convert_when(:end) == 0

    @info "Custom tests"
    @test fv(0.01, 1, 0, -100, :end) == 101.0
    @test map(x -> fv(x, 10*12, -100, -100), [0.05, 0.06, 0.07]/12) ≈ [15692.92889434,  16569.87435405,  17509.44688102]
    @test pmt(0.075/12, 12*15, 200000) ≈ -1854.0247200054619
    @test nper(0.07/12, -150, 8000) ≈ 64.07334877066185
    @test ipmt(0.01, 4, 10, -100.0) ≈ 0.7103767498422946
    @test rate(1, 0, -100, 101) ≈ 0.01
    @test isnan(rate(10, -10, 100, 102))

    # unit tests borrowed from numpy

    @info "Testing irr"
    v = [-150000, 15000, 25000, 35000, 45000, 60000]
    @test round(irr(v), digits = 4) == 0.0524

    v = [-100, 0, 0, 74]
    @test round(irr(v), digits = 4) == -0.0955

    v = [-100, 39, 59, 55, 20]
    @test round(irr(v), digits = 5) == 0.28095

    v = [-100, 100, 0, -7]
    @test round(irr(v), digits = 4) == -0.0833

    v = [-100, 100, 0, 7]
    @test round(irr(v), digits = 5) == 0.06206

    v = [-5, 10.5, 1, -8, 1]
    @test round(irr(v), digits = 4) == 0.0886

    v = [-1, -2, -3]
    @test isnan(irr(v))    # in Julia, we cannot compare NaN with NaN

    @info "Testing pv"
    @test round(pv(0.07, 20, 12000, 0), digits = 2) == -127128.17

    @info "Testing pmt"
    res = pmt(0.08 / 12, 5 * 12, 15000)
    tgt = -304.145914
    @test round(res, digits = 6) == round(tgt, digits = 6)

    # Test the edge case where rate == 0.0
    res = pmt(0.0, 5 * 12, 15000)
    tgt = -250.0
    @test round(res, digits = 6) == round(tgt, digits = 6)

    # julia broadcast
    res = pmt.([0.3, 0.8], [12, 3], [2000, 20000])
    tgt = [-626.90814, -19311.258]
    # @info "res=$res"
    # @info "tgt=$tgt"
    # @info "result=$(round.(res,3) .== round.(tgt,3))"
    @test all(round.(res, digits = 3) .== round.(tgt, digits = 3)) == true

    @info "Testing ppmt"
    @test round(ppmt(0.1 / 12, 1, 60, 55000), digits = 2) == -710.25
    @test round(ppmt(0.23 / 12, 1, 60, 10000000000), digits = 8) == -90238044.232277036

    @info "Testing ipmt"
    @test round(ipmt(0.1 / 12, 1, 24, 2000), digits = 2) == -16.67

    @info "Testing nper"
    @test round(nper(0.075, -2000, 0, 100000.), digits = 2) == 21.54
    @test round(nper(0.0, -2000, 0, 100000.), digits = 1) == 50.0

    @info "Testing npv"
    @test round(npv(0.05, [-15000, 1500, 2500, 3500, 4500, 6000]), digits = 2) == 122.89

    @info "Testing mirr"
    val = [-4500, -800, 800, 800, 600, 600, 800, 800, 700, 3000]
    @test round(mirr(val, 0.08, 0.055), digits = 4) == 0.0666

    val = [-120000, 39000, 30000, 21000, 37000, 46000]
    @test round(mirr(val, 0.10, 0.12), digits = 6) == 0.126094

    val = [100, 200, -50, 300, -200]
    @test round(mirr(val, 0.05, 0.06), digits = 4) == 0.3428

    val = [39000, 30000, 21000, 37000, 46000]
    @test isnan(mirr(val, 0.10, 0.12))

    @info "Testing eir"
    @test round(eir(0.08, 2), digits = 3) == 0.082

    @info "All tests completed"
end

ext: jl


text:
module Day23

export day23

mutable struct GameState
    current::Int
    next::Dict{Int, Int}
end

function destination(state)
    n = length(state.next)
    candidate = state.current > 1 ? state.current - 1 : n
    slice = Vector{Int}()
    c = state.current
    for i in 1:3
        c = state.next[c]
        push!(slice,c)
    end
    while in(candidate, slice)
        candidate = candidate > 1 ? candidate - 1 : n
    end
    candidate, slice
end

function step!(state)
    (dest, slice) = destination(state)
    state.next[state.current] = state.next[slice[3]]
    tmp = state.next[dest]
    state.next[dest] = slice[1]
    state.next[slice[3]] = tmp
    state.current = state.next[state.current]
end

function day23(::Val{:a})
    input = [7,8,4,2,3,5,9,1,6]
    #    input = [3,8,9,1,2,5,4,6,7]
    next = Dict{Int, Int}()
    for i in 1:length(input)-1
        next[input[i]] = input[i+1]
    end
    next[input[end]] = input[1]
    state = GameState(input[1],next)
    for i in 1:100
        step!(state)
    end
    answer = Vector{String}()
    c = 1
    for i in 1:length(input)-1
        push!(answer, string(state.next[c]))
        c = state.next[c]
    end
    join(answer)
end

function day23(::Val{:b})
    input = [7,8,4,2,3,5,9,1,6]
    # input = [3,8,9,1,2,5,4,6,7]
    next = Dict{Int, Int}()
    for i in 1:length(input)-1
        next[input[i]] = input[i+1]
    end
    next[input[end]] = length(input) + 1
    for k in length(input)+1:10^6-1
        next[k] = k+1
    end
    next[10^6] = input[1]
    state = GameState(input[1],next)
    for i in 1:10^7
        step!(state)
    end
    a = state.next[1]
    b = state.next[a]
    a * b
end

end


text:
include("DGModel_kernels.jl")

struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD, GF, TF} <:
       SpaceDiscretization
    balance_law::BL
    grid::G
    numerical_flux_first_order::NFND
    numerical_flux_second_order::NFD
    numerical_flux_gradient::GNF
    state_auxiliary::AS
    state_gradient_flux::DS
    states_higher_order::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
    gradient_filter::GF
    tendency_filter::TF
    check_for_crashes::Bool
end

function DGModel(
    balance_law,
    grid,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    fill_nan = false,
    check_for_crashes = false,
    state_auxiliary = create_state(
        balance_law,
        grid,
        Auxiliary(),
        fill_nan = fill_nan,
    ),
    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
    states_higher_order = (
        create_state(balance_law, grid, GradientLaplacian()),
        create_state(balance_law, grid, Hyperdiffusive()),
    ),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
    gradient_filter = nothing,
    tendency_filter = nothing,
)
    state_auxiliary =
        init_state(state_auxiliary, balance_law, grid, direction, Auxiliary())
    DGModel(
        balance_law,
        grid,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
        gradient_filter,
        tendency_filter,
        check_for_crashes,
    )
end

# Include the remainder model for composing DG models and balance laws
include("remainder.jl")

"""
    (dg::DGModel)(tendency, state_prognostic, _, t, α, β)

Uses spectral element discontinuous Galerkin in all direction to compute the
tendency.

Computes the tendency terms compatible with `IncrementODEProblem`

    tendency .= α .* dQdt(state_prognostic, p, t) .+ β .* tendency

The 4-argument form will just compute

    tendency .= dQdt(state_prognostic, p, t)
"""
function (dg::DGModel)(tendency, state_prognostic, _, t, α, β)

    device = array_device(state_prognostic)
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order

    FT = eltype(state_prognostic)
    num_state_prognostic = number_states(dg.balance_law, Prognostic())
    num_state_gradient_flux = number_states(dg.balance_law, GradientFlux())
    nhyperviscstate = number_states(dg.balance_law, Hyperdiffusive())
    num_state_tendency = size(tendency, 2)

    @assert num_state_prognostic ≤ num_state_tendency

    if num_state_prognostic < num_state_tendency && β != 1
        # If we don't operate on the full state, then we need to scale here instead of volume_tendency!
        tendency .*= β
        β = β != 0 # if β==0 then we can avoid the memory load in volume_tendency!
    end

    communicate =
        !(
            isstacked(dg.grid.topology) &&
            typeof(dg.direction) <: VerticalDirection
        )

    update_auxiliary_state!(
        dg,
        dg.balance_law,
        state_prognostic,
        t,
        dg.grid.topology.realelems,
    )

    exchange_state_prognostic = NoneEvent()
    exchange_state_gradient_flux = NoneEvent()
    exchange_Qhypervisc_grad = NoneEvent()
    exchange_Qhypervisc_div = NoneEvent()

    comp_stream = Event(device)

    if communicate
        exchange_state_prognostic = MPIStateArrays.begin_ghost_exchange!(
            state_prognostic;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0 || nhyperviscstate > 0
        ########################
        # Gradient Computation #
        ########################

        comp_stream = launch_volume_gradients!(
            dg,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dg.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dg.check_for_crashes,
            )
            update_auxiliary_state!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end

        comp_stream = launch_interface_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_state_prognostic),
        )

        if dg.gradient_filter !== nothing
            comp_stream = Filters.apply_async!(
                dg.state_gradient_flux,
                1:num_state_gradient_flux,
                dg.grid,
                dg.gradient_filter;
                dependencies = comp_stream,
            )
        end

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        dg.state_gradient_flux,
                        dependencies = comp_stream,
                    )
            end
            if nhyperviscstate > 0
                exchange_Qhypervisc_grad = MPIStateArrays.begin_ghost_exchange!(
                    Qhypervisc_grad,
                    dependencies = comp_stream,
                )
            end
        end

        if num_state_gradient_flux > 0
            # Update_aux_diffusive may start asynchronous work on the compute device
            # and we synchronize those here through a device event.
            checked_wait(device, comp_stream, nothing, dg.check_for_crashes)
            update_auxiliary_state_gradient!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

    if nhyperviscstate > 0
        #########################
        # Laplacian Computation #
        #########################

        comp_stream = launch_volume_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_Qhypervisc_grad = MPIStateArrays.end_ghost_exchange!(
                Qhypervisc_grad,
                dependencies = exchange_Qhypervisc_grad,
                check_for_crashes = dg.check_for_crashes,
            )
        end

        comp_stream = launch_interface_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_Qhypervisc_grad),
        )

        if communicate
            exchange_Qhypervisc_div = MPIStateArrays.begin_ghost_exchange!(
                Qhypervisc_div,
                dependencies = comp_stream,
            )
        end

        ####################################
        # Hyperdiffusive terms computation #
        ####################################

        comp_stream = launch_volume_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t,
            dependencies = (comp_stream,),
        )

        comp_stream = launch_interface_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = (comp_stream,),
        )

        if communicate
            exchange_Qhypervisc_div = MPIStateArrays.end_ghost_exchange!(
                Qhypervisc_div,
                dependencies = exchange_Qhypervisc_div,
                check_for_crashes = dg.check_for_crashes,
            )
        end

        comp_stream = launch_interface_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_Qhypervisc_div),
        )

        if communicate
            exchange_Qhypervisc_grad = MPIStateArrays.begin_ghost_exchange!(
                Qhypervisc_grad,
                dependencies = comp_stream,
            )
        end
    end

    ###################
    # RHS Computation #
    ###################
    comp_stream = launch_volume_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        dependencies = (comp_stream,),
    )

    comp_stream = launch_interface_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :interior,
        dependencies = (comp_stream,),
    )

    if communicate
        if num_state_gradient_flux > 0 || nhyperviscstate > 0
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.end_ghost_exchange!(
                        dg.state_gradient_flux;
                        dependencies = exchange_state_gradient_flux,
                        check_for_crashes = dg.check_for_crashes,
                    )

                # Update_aux_diffusive may start asynchronous work on the
                # compute device and we synchronize those here through a device
                # event.
                checked_wait(
                    device,
                    exchange_state_gradient_flux,
                    nothing,
                    dg.check_for_crashes,
                )
                update_auxiliary_state_gradient!(
                    dg,
                    dg.balance_law,
                    state_prognostic,
                    t,
                    dg.grid.topology.ghostelems,
                )
                exchange_state_gradient_flux = Event(device)
            end
            if nhyperviscstate > 0
                exchange_Qhypervisc_grad = MPIStateArrays.end_ghost_exchange!(
                    Qhypervisc_grad;
                    dependencies = exchange_Qhypervisc_grad,
                    check_for_crashes = dg.check_for_crashes,
                )
            end
        else
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dg.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dg.check_for_crashes,
            )
            update_auxiliary_state!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end
    end

    comp_stream = launch_interface_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :exterior,
        dependencies = (
            comp_stream,
            exchange_state_prognostic,
            exchange_state_gradient_flux,
            exchange_Qhypervisc_grad,
        ),
    )

    # The synchronization here through a device event prevents CuArray based and
    # other default stream kernels from launching before the work scheduled in
    # this function is finished.
    if dg.tendency_filter !== nothing
        comp_stream = Filters.apply_async!(
            tendency,
            1:num_state_tendency,
            dg.grid,
            dg.tendency_filter;
            dependencies = comp_stream,
        )
    end
    checked_wait(device, comp_stream, nothing, dg.check_for_crashes)
end

function restart_ode_state(dg::DGModel, state_data; init_on_cpu = false)
    bl = dg.balance_law
    grid = dg.grid

    state = create_state(bl, grid, Prognostic())
    state .= state_data

    device = arraytype(dg.grid) <: Array ? CPU() : CUDADevice()
    event = Event(device)
    event = MPIStateArrays.begin_ghost_exchange!(state; dependencies = event)
    event = MPIStateArrays.end_ghost_exchange!(state; dependencies = event)
    wait(device, event)

    return state
end

function indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    state_prognostic::MPIStateArray,
    state_auxiliary::MPIStateArray,
    t::Real,
    elems::UnitRange = dg.grid.topology.elems,
)

    device = array_device(state_prognostic)

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    FT = eltype(state_prognostic)

    # Compute integrals
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_indefinite_stack_integral!(device, (Nq[1], Nqj))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_prognostic.data,
        state_auxiliary.data,
        grid.vgeo,
        # Only need the vertical Imat since this kernel is vertically oriented
        grid.Imat[dim],
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    checked_wait(device, event, nothing, dg.check_for_crashes)
end

function reverse_indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    state_prognostic::MPIStateArray,
    state_auxiliary::MPIStateArray,
    t::Real,
    elems::UnitRange = dg.grid.topology.elems,
)

    device = array_device(state_auxiliary)

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    FT = eltype(state_auxiliary)

    # Compute integrals
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_reverse_indefinite_stack_integral!(device, (Nq[1], Nqj))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_prognostic.data,
        state_auxiliary.data,
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    checked_wait(device, event, nothing, dg.check_for_crashes)
end


text:
mutable struct TrackedState <: State
    queues::Vector{Vector{Job}} #A vector of queues holding the jobs waiting in buffer
    in_transit::BinaryMinHeap{Job} #Jobs in transit between queues
    left_system::Vector{Job} #Jobs that have left the system
    params::NetworkParameters #The parameters of the queueing system
    arrival_distr::Gamma{Float64}
    travel_distr::Gamma{Float64}
    service_distrs::Vector{Gamma{Float64}}
end

"""
Constructs an initilized tracked network state from the given parameters
"""
TrackedState(params::NetworkParameters) = TrackedState(
    [Vector{Job}[] for _ in 1:params.L], # Initial queues
    BinaryMinHeap{Job}(),                # Intial transit
    Vector{Job}[],                       # Initial left system
    params,
    rate_scv_gamma(params.λ, params.gamma_scv),
    rate_scv_gamma(params.η, params.gamma_scv),
    map((μ) -> rate_scv_gamma(μ, params.gamma_scv), params.μ_vector)
    )

queued_count(state::TrackedState)::Int = sum([length(state.queues[i]) for i in 1:length(state.queues)])

transit_count(state::TrackedState)::Int = length(state.in_transit)

num_in_queue(q::Int, state::TrackedState)::Int = length(state.queues[q])

function clear_left(state::TrackedState)
    empty!(state.left_system)
end

""" Prints the current state of the system"""
function show(io::IO, state::TrackedState)
    print(io, "State:\n")
    for (i, queue) in enumerate(state.queues)
        print(io, "Q$i: $(length(queue))\n")
    end
    print(io, "T: $(length(state.in_transit))\n")
end

function new_job(q::Int, time::Float64, state::TrackedState)
    push!(state.queues[q], Job(time))
end

function new_transit(time::Float64, transit_time::Float64, state::TrackedState)
    push!(state.in_transit, Job(time, transit_time))
end

function pop_queue(q::Int, time::Float64, state::TrackedState)
    job = popfirst!(state.queues[q])
    job.exit_time = time
    push!(state.left_system, job)
end

function queue_to_transit(q::Int, time::Float64, transit_time::Float64, state::TrackedState)
    job = popfirst!(state.queues[q])
    job.exit_time = transit_time
    push!(state.in_transit, job)
end

function pop_transit(time::Float64, state::TrackedState)
    job = pop!(state.in_transit)
    job.exit_time = time
    push!(state.left_system, job)
end

function transit_to_queue(q::Int, time::Float64, state::TrackedState)
    job = pop!(state.in_transit)
    job.exit_time = time
    push!(state.queues[q], job)
end

function update_transit(time::Float64, transit_time::Float64, state::TrackedState)
    job = pop!(state.in_transit)
    job.exit_time = transit_time
    push!(state.in_transit, job)
end

function failed_arrival(time::Float64, state::TrackedState)
    push!(state.left_system, Job(time))
end


text:
module SnoopCompileWithMocking

using Mocking

export make_fancy

function make_fancy(normal)::Nothing
    @mock fancynize(normal)
end

function fancynize(normal)::Nothing
    println("This is now awesome: $normal")
end

include("precompile_includer.jl")
end # module SnoopCompileWithMocking


text:
### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 9739e5dc-369e-11eb-2c56-f36ccd348299
begin

	using Plots
	using Roots
	using PlutoUI
	using Latexify
	using StaticArrays
	using LinearAlgebra
	using ModelingToolkit
	using UnitfulAstrodynamics
	using DifferentialEquations

	μ = nondimensionalize(Earth.μ, Moon.μ)

end;

# ╔═╡ 860ea13a-369e-11eb-34b2-3fc6557a68d6
md"""

# 🎢 Halo Orbits & Invariant Manifolds

_Invariant manifolds about Halo orbits and their applications in the Circular Restricted Three-body Problem._

- Joe Carpinelli
- December, $2020$
- ENAE$601$ Final Project

###### Presentation Mode: $(html"<button onclick=present()>Toggle Presentation</button>")

"""

# ╔═╡ 2f3bcf58-3758-11eb-352c-a1181281eba0
md"""

## 📋 Project Overview

In the context of astrodynamics,  __manifolds are groups of trajectories that move toward or away from Lagrange points__. To use invariant manifolds for interplanetary travel, several concepts must be developed and built on: lagrange points, periodic and quasi-periodic orbits within the Circular Restricted Three-body Problem, and manifolds about periodic orbits $\verb|[1]|$.

### Outline

- Brief review of Lagrange points
- Periodic orbits (specifically, the subsection known as Halo orbits)
- Finding Halo orbits (both analytically, and numerically)
- Invariant manifolds about Halo orbits

### Primary Reference

- Megan Rund's Masters Thesis at California Polytechnic State University $\verb|[1]|$

"""

# ╔═╡ fc4106f6-399e-11eb-1f9f-e7c4d5539f20
md"""

## Stability at Lagrange Points

* Like equilibrium points for all nonlinear systems, Lagrange points can be __stable or unstable__
* We can find the stability of an equilibrium point by analysing __eigenvalues__ of the Jacobian of the state vector

In the $2$D case:

$\begin{bmatrix}\dot{\zeta} \\ \dot{\beta} \\ \ddot{\zeta} \\ \ddot{\eta} \end{bmatrix} = \begin{bmatrix} 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 1 \\ U_{xx} & U_{xy} & 0 & 2 \\ U_{yx} & U_{yy} & -2 & 0\end{bmatrix} \begin{bmatrix} \zeta \\ \eta \\ \dot{\zeta} \\ \dot{\eta} \end{bmatrix}$

$\lambda^4 + (4-U_{xx}-U_{yy})\lambda^2 + U_{xx} U_{yy} - U_{xy}^2 = 0$

"""

# ╔═╡ 8c0a0a9c-39a1-11eb-1fe6-c7177fa4b5a8
md"""

## Periodic Orbits

#### Libration Orbit Families
* Orbits about Lagrange points are known as __Libration orbits__
* __Lyapunov__ orbits are two-dimensional Libration orbits
* __Lissajous__ orbits are three-dimensional, semi-periodic Libration orbits
* __Halo__ orbits are three-dimensional, theoretically periodic Libration orbits

#### Finding Halo Orbits
* The calculations required to solve for Halo orbits are numerically sensitive
* We covered the procedure for iteratively solving for Halo orbits __numerically__ in __Lecture 16__; this requires an initial guess

* __How do we find an initial guess for a Halo orbit?__

"""

# ╔═╡ 95a851d0-39a1-11eb-05b1-81fbb6c0bd54
md"""

## 😇 Analytical Halo Solution

* The Circular Restricted Three-body Dynamics can be derived using "Legendre poly- nomials", and there exists a third-order approximation (shown below) $\verb|[5]|$
* We can choose parameters to remove __unstable__ (secular) terms from the expansion

#### Third Order CR3BP Expansion

   $(LocalResource("./media/poly.png", :width=>400))

#### Algorithm Inputs

* Nondimensional mass parameter $\mu$
* Nondimensional $Z-axis$ amplitude for the desired Halo orbit
* Lagrange point to orbit (L1 or L2)

#### Algorithm Outputs

* Position vector $\overrightarrow{r}_0$
* Velocity vector $\overrightarrow{v}_0$
* Estimated orbital period $T$

"""

# ╔═╡ bf3ee972-39ac-11eb-2bdb-fd28bebd418e
md"""

## 🧮 Numerical Halo Solution

* As discussed in __Lecture 16__, we can append the __state transition matrix__ $\Phi(t_0 + t_i, t_0)$ to our state vector, and iteratively change initial conditions to __numerically find a periodic orbit__ $\verb|[1]|$

1.  $\Phi = I,\ \overrightarrow{x}_0 = \begin{bmatrix} x_0 & 0 & z_0 & 0 & \dot{y}_0 & 0 & \Phi_1 & \Phi_2 & \Phi_3 & \Phi_4 & \Phi_5 & \Phi_6 \end{bmatrix}$
2. Propagate until $y=0$ again; $\dot{\Phi} = F \Phi$ where $F = \begin{bmatrix} 0 & I_3 \\ U_{XX} & 2\Omega \end{bmatrix}$, $U_{XX}$ is the matrix of second partial derivatives of potential $U$
3. Calculate $\begin{bmatrix} \delta x_0 \\ \delta \dot{y}_0 \end{bmatrix} = \left( \begin{bmatrix} \Phi_{4 1} & \Phi_{4 5} \\ \Phi_{6 1} & \Phi_{6 5} \end{bmatrix}    \frac{1}{\dot{y}} \begin{bmatrix} \ddot{x} \\ \ddot{z} \end{bmatrix} \right)^{-1} \begin{bmatrix}  -\dot{x} \\ -\dot{z} \end{bmatrix}$
4. Set $x_0 \leftarrow x_0 + \delta x_0$, $\dot{y}_0 \leftarrow \dot{y}_0 + \delta \dot{y}$ and jump back to __Step 1__ (until $\dot{x}$, $\dot{z}$ are both within some tolerance of zero)

"""

# ╔═╡ b123e506-39af-11eb-06b3-cfb02d9d50fc
md"""

## 🎢 Dynamics along Halo Orbits

### Manifolds Exist
* Each point along a Halo orbit is connected with an __unstable__ manifold, and a __stable__ manifold
* The unstable manifold __departs__ the Halo orbit, and the stable manifold __arrives__ at the Halo orbit

* __How can we calculate the perturbation required to shift the spacecraft onto the manifold?__

### Finding Manifolds
* We can use __Eigenvectors__ of the Jacobian to calculate a state perturbation which will place the spacecraft onto a manifold at each point along the Halo orbit

1. Propagate the Halo orbit for one period $T$, including the state transition matrix $\Phi(t_0 + t_i, t_0)$
2. Let the final state transition matrix be $M$; $M = \Phi(t_0 + T, t_0)$
3. Calculate eigenvectors $V^{S} \leftrightarrow \text{minreal}(\text{eig}(M))$, and $V^{U} \leftrightarrow \text{maxreal}(\text{eig}(M))$
4. For each point $i$ along the Halo orbit...
   *  $V_i^S = \Phi(t_0 + t_i, t_0) V^S$, and $V_i^U = \Phi(t_0 + t_i, t_0) V^U$
   *  $X_i^S = X_i \pm \epsilon \frac{V_i^S}{|V_i^S|}$, and $X_i^U = X_i \pm \epsilon \frac{V_i^U}{|V_i^U|}$

"""

# ╔═╡ 542fc7e8-3a27-11eb-0284-07c4a6f626f5
md"""

## Invariant Manifold Example

* Pulled from Rund's Thesis $\verb|[1]|$
* Note that stable manifolds need to be propagated _backward_ in time from the perturbation along the Halo orbit, becuase they converge on the Halo orbit $\verb|[1]|$

$(LocalResource("media/fig3.6.png"))

"""

# ╔═╡ a68f4086-3a27-11eb-24d4-692815414617
md"""

## 🪐 Manifold-based Transfer Designs

* __Design Overview__ summarized from $\verb|[1]|$

### Design Overview

1. Find a desired Sun-Earth Halo orbit
2. Place the spacecraft within the __stable manifold__ of this Halo orbit
3. Perturb the spacecraft from the Halo onto the __unstable manifold__ to return towards Earth, and apply a maneuver to place the spacecraft on a Hyperbolic escape trajectory toward your destination planet
4. At destination planet, apply a maneuver to place spacecraft onto __stable manifold__ of destination Halo orbit

"""

# ╔═╡ 4b4fca98-3a2b-11eb-09d1-b358440b75fc
md"""

## 🚀 Conclusions

### Manifold Transfers
* Lagrange points, and periodic orbits about Lagrange points are surrounded by collections of trajectories called __manifolds__
* Manifolds can depart, or arrive at the Lagrange point / Libration orbit
* Stable manifolds can bring us to Halo orbits for free$^*$ (plus the cost to place the spacecraft onto the manifold)

### Lessons Learned

* Linear (Eigenvector) analysis works in Astrodynamics too
* Calculations for iterating on Halo orbits are __extremely__ numerically sensitive
* Even with more complicated models (CR3BP), we need a patch-conic-like approach for interplanetary mission design

"""

# ╔═╡ ddf8378a-3757-11eb-039b-254488a45127
md"""

## 📚 References

 $\verb|[1]|$ Rund, M. S., “Interplanetary Transfer Trajectories Using the Invariant Manifolds of Halo Orbits,” , 2018.

 $\verb|[2]|$ Howell, K. C., “Three-dimensional, periodic,‘halo’orbits,” Celestial mechanics, Vol. 32, No. 1, 1984, pp. 53–71.

 $\verb|[3]|$ Vallado, D. A., Fundamentals of astrodynamics and applications, Vol. 12, Springer Science & Business Media, 2001.

 $\verb|[4]|$ Richardson, D., “Analytical construction of periodic orbits about the collinear points of the Sun-Earth system.” asdy, 1980, p. 127.

 $\verb|[5]|$ Lara, M., Russell, R., and Villac, B., “Classification of the distant stability regions at Europa,” Journal of Guidance, Control, and Dynamics, Vol. 30, No. 2, 2007, pp. 409–418.

 $\verb|[6]|$ Koon, W. S., Lo, M. W., Marsden, J. E., and Ross, S. D., “Dynamical systems, the three-body problem and space mission design,” Free online Copy: Marsden Books, 2008.

 $\verb|[7]|$ Williams, J., Lee, D. E., Whitley, R. J., Bokelmann, K. A., Davis, D. C., and Berry, C. F., “Targeting cislunar near rectilinear halo orbits for human space exploration,” 2017.

 $\verb|[8]|$ Zimovan-Spreen, E. M., Howell, K. C., and Davis, D. C., “Near rectilinear halo orbits and nearby higher-period dynamical structures: orbital stability and resonance properties,” Celestial Mechanics and Dynamical Astronomy, Vol. 132, No. 5, 2020, pp. 1–25.

 $\verb|[9]|$ NASA, NASA’s Lunar Exploration Program Overview, 2020.

 $\verb|[10]|$ Carpinelli, J., “UnitfulAstrodynamics.jl,” https://juliahub.com/ui/Packages/UnitfulAstrodynamics/uJGLZ/, 2020.

"""

# ╔═╡ 1f0ebb78-3707-11eb-00da-2f08e79fd676
md"# ⌨ Source Code"

# ╔═╡ d74f9638-369e-11eb-067c-05720a8f4d94
md"""

## Package Dependencies

* The following packages were used -- all are available in Julia's [General](https://juliahub.com) package registry

"""

# ╔═╡ 38842964-36a1-11eb-3b0c-bf5bb8000341
md"""

## Finding Lagrange Points

"""

# ╔═╡ f4c7e16e-36a7-11eb-17f3-2f62591cebe8
f₁ = let

	f₁ = lagrangeplot(nondimensionalize(Earth.μ, Moon.μ); labels=["Earth" "Moon"])
	savefig(f₁, "media/fig1.png")

	f₁

end;

# ╔═╡ deb34994-399e-11eb-28a7-7d63db867a11
md"""

## Lagrange Points

* Lagrange points are __equilibrium points__ within the Circular Restricted Three-body Problem

$(f₁)

"""

# ╔═╡ d54ed67a-39a0-11eb-2d9b-a7800fd0163a
md"## Unstable Lagrange Point (L2)"

# ╔═╡ 9fc039b8-38c4-11eb-305d-7bff22a6d965
f₂ = let

	μ = nondimensionalize(Earth.μ, Moon.μ)
	r = lagrange(μ, 2)[:]
	v = [0.0, 0.0, 0.0]

	sys = NondimensionalThreeBodyState(r, v, μ, NaN * u"km", NaN * u"s")
	sols = propagate(sys, 30; save_everystep=true, reltol=1e-16, abstol=1e-16)

	x = [sols.step[i].rₛ[1] for i ∈ 1:length(sols.step)]
	y = [sols.step[i].rₛ[2] for i ∈ 1:length(sols.step)]
	z = [sols.step[i].rₛ[3] for i ∈ 1:length(sols.step)]

	f₂ = plot(x,y; label="Spacecraft Position")
	scatter!(f₂, [[i] for i ∈ lagrange(μ, 2)[1:2]]...; label="L2", markershape=:x)
	scatter!(f₂, [[i] for i ∈ sys.r₁][1:2]...; label="Earth", markersize=7)
	scatter!(f₂, [[i] for i ∈ sys.r₂][1:2]...; label="Moon")
	plot!(f₂; title="Spacecraft Propagated from Earth-Moon L2",
			  xlabel="X (AU)",
			  ylabel="Y (AU)")

	savefig(f₂, "media/fig2.png")

	f₂

end;

# ╔═╡ e3cf8532-39a0-11eb-0244-e509c5e55b75
md"## Stable Lagrange Point (L4)"

# ╔═╡ 37a8538c-3997-11eb-2649-2daacb98b8c8
f₇ = let

	μ = nondimensionalize(Earth.μ, Moon.μ)
	r = lagrange(μ, 4)[:] * (1 + 1e-3)
	v = [0.0, 0.0, 0.0]

	sys = NondimensionalThreeBodyState(r, v, μ, NaN * u"km", NaN * u"s")
	sols = propagate(sys, 1000; save_everystep=true, reltol=1e-16, abstol=1e-16)

	x = [sols.step[i].rₛ[1] for i ∈ 1:length(sols.step)]
	y = [sols.step[i].rₛ[2] for i ∈ 1:length(sols.step)]
	z = [sols.step[i].rₛ[3] for i ∈ 1:length(sols.step)]

	f₇ = plot(x,y; label="Spacecraft Position")
	scatter!(f₇, [[i] for i ∈ lagrange(μ, 4)[1:2]]...; label="L4", markershape=:x)
	scatter!(f₇, [[i] for i ∈ sys.r₁][1:2]...; label="Earth", markersize=7)
	scatter!(f₇, [[i] for i ∈ sys.r₂][1:2]...; label="Moon")
	plot!(f₇; title="Spacecraft Perturbed from Earth-Moon L4",
			  xlabel="X (AU)",
			  ylabel="Y (AU)")

	savefig(f₂, "media/fig2.png")

	f₇

end;

# ╔═╡ 7d9da4e8-399f-11eb-35b2-11b9f4c1371f
md"""

## Stability at Lagrange Points: Examples

* Earth-Moon L4 is __stable__, and Earth-Moon L2 is __unstable__

$(plot(plot(f₂; title="Earth-Moon L2", labels=["" "" "" ""]), plot(f₇; title="Earth-Moon L4", legend=:right, labels=["Spacecraft Position" "L2 & L4" "Earth" "Moon"]))))

"""

# ╔═╡ e9e0717a-39a0-11eb-28bc-1bb42537ff43
md"## Plot Analytical Halo Orbit"

# ╔═╡ d251ce38-38dc-11eb-0fa6-653456709b32
function haloplot(μ, L, Z, H, str1="Mass 1", str2="Mass 2"; kwargs...)

	if L==:L1
		l = 1
	else
		l = 2
	end

	defaults = (; title="Analytical Halo Solutions",
				  xlabel="X (DU)", ylabel="Y (DU)", zlabel="Z (DU)")
	options  = merge(defaults, kwargs)

	fig = plot(; options...)
	for z ∈ Z
		r,v,Τ = halo_analytic(μ; L=L, Zₐ=z, hemisphere=H, steps=1000)
		x = r[:,1]
		y = r[:,2]
		z = r[:,3]
		plot!(fig, x, y, z; label=:none)

	end

	scatter!(fig, [[v] for v ∈ lagrange(μ, l)]...;
			 label=string("L",L), markershape=:x)
	scatter!(fig, [-μ], [0], [0]; label=str1)
	scatter!(fig, [1-μ], [0], [0]; label=str2)

	return fig

end;

# ╔═╡ f24c96cc-39a0-11eb-12ba-d33ac3e6a77d
md"## Northern Sun-Jupiter Halos"

# ╔═╡ cd021b9a-38dc-11eb-1f5b-1b3a2294ef56
f₃ = let

	μ = nondimensionalize(Sun.μ, Jupiter.μ)
	Z = [x / 10 for x ∈ 1:10]
	f₃ = haloplot(μ, 1, Z, :northern, "Sun", "Jupiter";
				  title="Northern Analytical Halo Solutions")

	savefig(f₃, "media/fig3.png")

	f₃

end;

# ╔═╡ fc4671ac-39a0-11eb-2867-f99155906347
md"## Southern Sun-Jupiter Halos"

# ╔═╡ c544f248-38ea-11eb-1a76-f763700e5e66
f₄ = let

	μ = nondimensionalize(Sun.μ, Jupiter.μ)
	Z = [x / 10 for x ∈ 1:10]
	f₄ = haloplot(μ, 1, Z, :southern, "Sun", "Jupiter";
				  title="Southern Analytical Halo Solutions")

	savefig(f₄, "media/fig4.png")

	f₄

end;

# ╔═╡ 41c340f0-39a8-11eb-375a-c311f97cb062
md"""

## Analytical Halo Examples: Sun-Jupiter L1

__NOT numerically propagated!__

$(plot(f₄; title=" Sun-Jupiter Halo Examples"))

"""

# ╔═╡ 032912e0-39a1-11eb-23d7-a1be4cc91e98
md"## Northern Earth-Moon Halos"

# ╔═╡ 2b968ad4-38eb-11eb-106a-7fc6c7b60bc0
f₅ = let

	μ = nondimensionalize(Earth.μ, Moon.μ)
	Z = [x / 10 for x ∈ 1:10]
	f₅ = haloplot(μ, 2, Z, :northern, "Earth", "Moon";
				  title="Northern Analytical Halo Solutions")

	savefig(f₅, "media/fig5.png")

	f₅

end;

# ╔═╡ 868a7c00-39a6-11eb-23d8-43efa8cea0a2
md"""

## Analytical Halo Examples: Earth-Moon L2

__NOT numerically propagated!__

$(plot(f₅; title="Earth-Moon Halo Examples"))

"""

# ╔═╡ 0c26f664-39a1-11eb-2d2c-75cc99d2a193
md"## Southern Earth-Moon Halos"

# ╔═╡ 4e5b126a-38eb-11eb-30d4-03206b35fd53
f₆ = let

	μ = nondimensionalize(Earth.μ, Moon.μ)
	Z = [x / 10 for x ∈ 1:10]
	f₆ = haloplot(μ, 2, Z, :southern, "Earth", "Moon";
				  title="Southern Analytical Halo Solutions")

	savefig(f₆, "media/fig6.png")

	f₆

end;

# ╔═╡ bbc1be44-3a44-11eb-3da2-19023d2871d1
md"## Numerically Produced Halo"

# ╔═╡ 5c42be9c-39be-11eb-1fb0-05e460b4bd2f
f₈ = let

	r,v,T = halo(μ, L=2, Zₐ=0.05, ϕ=0.05, max_iter=20)
	sys = NondimensionalThreeBodyState(r,v,μ, NaN*u"km", NaN*u"s")
	sols = propagate(sys, T)
	x = [x.rₛ[1] for x ∈ sols.step]
	y = [x.rₛ[2] for x ∈ sols.step]
	z = [x.rₛ[3] for x ∈ sols.step]

	f₈ = plot(x,y,z; label="Spacecraft Position")
	plot!(f₈; title="Numerically Produced Earth-Moon L2 Halo",
			xlabel="X (DU)", ylabel="Y (DU)", zlabel="Z (DU)")

	savefig(f₈, "media/fig8.png")

	f₈

end;

# ╔═╡ Cell order:
# ╟─860ea13a-369e-11eb-34b2-3fc6557a68d6
# ╟─2f3bcf58-3758-11eb-352c-a1181281eba0
# ╟─deb34994-399e-11eb-28a7-7d63db867a11
# ╟─fc4106f6-399e-11eb-1f9f-e7c4d5539f20
# ╟─7d9da4e8-399f-11eb-35b2-11b9f4c1371f
# ╟─8c0a0a9c-39a1-11eb-1fe6-c7177fa4b5a8
# ╟─95a851d0-39a1-11eb-05b1-81fbb6c0bd54
# ╟─868a7c00-39a6-11eb-23d8-43efa8cea0a2
# ╟─41c340f0-39a8-11eb-375a-c311f97cb062
# ╟─bf3ee972-39ac-11eb-2bdb-fd28bebd418e
# ╟─b123e506-39af-11eb-06b3-cfb02d9d50fc
# ╟─542fc7e8-3a27-11eb-0284-07c4a6f626f5
# ╟─a68f4086-3a27-11eb-24d4-692815414617
# ╟─4b4fca98-3a2b-11eb-09d1-b358440b75fc
# ╟─ddf8378a-3757-11eb-039b-254488a45127
# ╟─1f0ebb78-3707-11eb-00da-2f08e79fd676
# ╟─d74f9638-369e-11eb-067c-05720a8f4d94
# ╠═9739e5dc-369e-11eb-2c56-f36ccd348299
# ╟─38842964-36a1-11eb-3b0c-bf5bb8000341
# ╠═f4c7e16e-36a7-11eb-17f3-2f62591cebe8
# ╟─d54ed67a-39a0-11eb-2d9b-a7800fd0163a
# ╠═9fc039b8-38c4-11eb-305d-7bff22a6d965
# ╟─e3cf8532-39a0-11eb-0244-e509c5e55b75
# ╠═37a8538c-3997-11eb-2649-2daacb98b8c8
# ╟─e9e0717a-39a0-11eb-28bc-1bb42537ff43
# ╠═d251ce38-38dc-11eb-0fa6-653456709b32
# ╟─f24c96cc-39a0-11eb-12ba-d33ac3e6a77d
# ╠═cd021b9a-38dc-11eb-1f5b-1b3a2294ef56
# ╟─fc4671ac-39a0-11eb-2867-f99155906347
# ╠═c544f248-38ea-11eb-1a76-f763700e5e66
# ╟─032912e0-39a1-11eb-23d7-a1be4cc91e98
# ╠═2b968ad4-38eb-11eb-106a-7fc6c7b60bc0
# ╟─0c26f664-39a1-11eb-2d2c-75cc99d2a193
# ╠═4e5b126a-38eb-11eb-30d4-03206b35fd53
# ╟─bbc1be44-3a44-11eb-3da2-19023d2871d1
# ╠═5c42be9c-39be-11eb-1fb0-05e460b4bd2f

max_stars_repo_name: cadojo/Halo-Orbit-Explorations

max_issues_repo_name: cadojo/Halo-Orbit-Explorations

max_forks_repo_name: cadojo/Halo-Orbit-Explorations


text:
module Optini

using LinearAlgebra
using Zygote

import Base.MathConstants: φ

# Generic types
include("types.jl")

# Univariate algorithms
include("univariate/algorithms/fibonacci.jl")
include("univariate/algorithms/golden_section.jl")
include("univariate/algorithms/quadratic_fit.jl")
include("univariate/algorithms/bisection.jl")
include("univariate/bracket.jl")
include("univariate/optimize.jl")

# Multivariate algorithms
include("multivariate/types.jl")
include("multivariate/optimize.jl")
include("multivariate/misc/steepest.jl")
include("multivariate/misc/newton.jl")

# Line search methods
include("multivariate/line_search/types.jl")
include("multivariate/line_search/interface.jl")
include("multivariate/line_search/initial_guess.jl")
include("multivariate/line_search/algorithms/static.jl")
include("multivariate/line_search/algorithms/exact.jl")
include("multivariate/line_search/algorithms/backtracking.jl")
include("multivariate/line_search/algorithms/interpolation.jl")
include("multivariate/line_search/algorithms/strong_wolfe.jl")

# Trust region methods
include("multivariate/trust_region/types.jl")
include("multivariate/trust_region/interface.jl")
include("multivariate/trust_region/algorithms/cauchy.jl")
include("multivariate/trust_region/algorithms/dogleg.jl")
include("multivariate/trust_region/algorithms/2d_subspace.jl")

export bracket
export Bisection, Fibonacci, GoldenSection, QuadraticFit

export Steepest
export Newton

export StaticInitial, PreviousDecreaseInitial, QuadraticInitial
export StaticLineSearch, ExactLineSearch, BacktrackingLineSearch, InterpolationLineSearch,
    StrongWolfeLineSearch
export LineSearch, GradientDescent, NewtonDescent

export CauchyPoint, Dogleg, TwoDimSubspace
export TrustRegion

export optimize

end


text:
# This module contains basic utility functions and utility functions to
# simulate modelica models using OMJulia

"""
    MoSTError

Error class for OMJulia-related errors that contains the OMC error message.
"""
struct MoSTError <: Exception
    msg:: String
    omc:: String
end

Base.showerror(io::IO, e::MoSTError) = print(io, e.msg, "\n---\nOMC error string:\n", e.omc)

"""
    MoSTError(omc:: OMCSession, msg:: String)

Creates MoSTError with message `msg` and current result of `getErrorString()`
as OMC error message.
"""
MoSTError(omc:: OMCSession, msg:: String) = MoSTError(msg, getErrorString(omc))

"""
    loadModel(omc:: OMCSession, name:: String; check=true, instantiate=true)

Loads the model with fully qualified name `name` from a source file available
from the model directory.
Note that this refers to the model *name*, not the model *file*.

Example:

    loadModel(omc, "Modelica.SIunits.Voltage")

This function will actually call several OM scripting functions to
ensure that as many errors in the model are caught and thrown as
[`MoSTError`](@ref)s as possible:

* First, `loadModel(name)` is called to load the model if it exists. This
    call does only fail if the toplevel model does not exist. E.g.,
    `loadModel(Modelica.FooBar)` would still return true, because `Modelica`
    could be loaded, although `FooBar` does not exist.
* We then check with `getClassRestriction(name)` if the model actually exists
    (which is the case when the return value is nonempty).
* With `checkModel(name)` we find errors such as missing or mistyped variables.
* Finally, we use `instantiateModel(name)` which can sometimes find additional
    errors in the model structure (e.g. since Modelica 1.16, unit consistency
    checks are performed here).

If `check`, or `instantiate` are false, the loading process is stopped at the
respective steps.
""" # TODO: which errors are found by instantiateModel that checkModel does not find?
function loadModel(omc:: OMCSession, name:: String; ismodel=true, check=true, instantiate=true)
    # only load model if it was not created by sending a class definition
    # string directly to the OMC
    if filename(omc, name) != "<interactive>"
        success = sendExpression(omc, "loadModel($name)")
        es = getErrorString(omc)
        if isnothing(success)
            # i have seen this happen, but do not know why it does occur
            throw(MoSTError("Unexpected error: loadModel($name) returned nothing", es))
        end
        if !success || length(es) > 0
            throw(MoSTError("Could not load $name", es))
        end
    end
    # loadModel will only fail if the *toplevel* class does not exist
    # => check that the full class name could actually be loaded
    if !isloaded(omc, name)
        throw(MoSTError("Model $name not found in MODELICAPATH", ""))
    end
    if !ismodel
        @warn(string(
            "The keyword parameter ismodel is deprecated since the existance",
            " of a model/class is now checked with new isloaded() function",
            " which does not fail like isModel() did.",
            " You can simply replace `ismodel=false` with `check=false`,",
            " which will now have the same effect."
        ))
        return
    end
    if !check
        return
    end
    check = sendExpression(omc, "checkModel($name)")
    es = getErrorString(omc)
    if !startswith(check, "Check of $name completed successfully")
        throw(MoSTError("Model check of $name failed", join([check, es], "\n")))
    end
    if !instantiate
        return
    end
    inst = sendExpression(omc, "instantiateModel($name)")
    es = getErrorString(omc)
    if length(es) > 0
        throw(MoSTError("Model $name could not be instantiated", es))
    end
end

"""
    isloaded(omc:: OMCSession, name:: String)

Checks that the model/class/package/... `name` was correctly loaded and can
be queried by other functions.
"""
function isloaded(omc:: OMCSession, name:: String)
    # One of the few functions that we can use here is getClassRestriction,
    # because it gives an empty string for nonexistent models and "model",
    # "class", "connector", ... for other class types
    classrest = sendExpression(omc, "getClassRestriction($name)")
    return !isempty(classrest)
end

"""
    filename(omc:: OMCSession, name:: String)

Returns the file name where the model/class/package/... `name` is stored.
If `name` was defined by directly sending a class definition to the OMC the
return value will be `"<interactive>"`. If the model could not be found,
the return value will be an empty string.
"""
function filename(omc:: OMCSession, name:: String)
    res = sendExpression(omc, "getClassInformation($name)")
    return res[6]
end

"""
    moescape(s:: String)

Escapes string according to Modelica specification for string literals.

Escaped characters are: `['\\\\', '"', '?', '\\a', '\\b', '\\f', '\\n', '\\r', '\\t', '\\v',]`
"""
function moescape(s:: String)
    rep = Dict(
        '\\' => "\\\\",
        '"' => "\\\"",
        '?' => "\\?",
        '\a' => "\\a",
        '\b' => "\\b",
        '\f' => "\\f",
        '\n' => "\\n",
        '\r' => "\\r",
        '\t' => "\\t",
        '\v' => "\\v",
    )
    return join([(x in keys(rep) ? rep[x] : x) for x in s])
end

"""
    moescape(s:: String)
    moescape(io:: IO, s:: String)

Unescapes string that was escaped by [`moescape(s:: String)`](@ref) or that
was returned from the OMC compiler. If `io` is given the string is printed to
the `IO` object, otherwise it is returned directly.
"""
function mounescape(io:: IO, s:: String)
    rev = Dict(
        "\\\\" => '\\',
        "\\\"" => '"',
        "\\?" => '?',
        "\\a" => '\a',
        "\\b" => '\b',
        "\\f" => '\f',
        "\\n" => '\n',
        "\\r" => '\r',
        "\\t" => '\t',
        "\\v" => '\v'
    )
    i = Iterators.Stateful(s)
    while !isempty(i)
        c = popfirst!(i)
        if c != '\\' || isempty(i)
            print(io, c)
        else
            nxt = popfirst!(i)
            print(io, rev[join([c, nxt])])
        end
    end
end
mounescape(s::String) = sprint(mounescape, s; sizehint=lastindex(s))

function getErrorString(omc:: OMCSession)
    es = sendExpressionRaw(omc, "getErrorString()")
    parsed = strip(strip(mounescape(es)),'"')
    # FIXME this should be removed if there is a way to fix the model or if OpenModelica 1.16 is updated
    # we ignore a specific cryptic error message from OpenModelica 1.16.0
    ignore = "Warning: function Unit.unitString failed for \"MASTER()\".\n"
    return parsed == ignore ? "" : parsed
end

function sendExpressionRaw(omc:: OMCSession, expr)
    # FIXME this function should be replaced by sendExpression(omc, parsed=false)
    send(omc.socket, expr)
    message=recv(omc.socket)
    return unsafe_string(message)
end

"""
    getSimulationSettings(omc:: OMCSession, name:: String; override=Dict())

Reads simulation settings from the model `name`.
Any content in `override` will override the setting with the respective key.

Returns a Dict with the keys `"startTime"`, `"stopTime"`, `"tolerance"`,
`"numberOfIntervals"`, `"outputFormat"` and `"variableFilter"`.
If any of these settings are not defined in the model file, they will be
filled with default values.

In `override`, an additional key `"interval"` is allowed to recalculate the
`"numberOfIntervals"` based on the step size given as value to this key.

Throws a [`MoSTError`](@ref) if the model `name` was not loaded beforehand using
[`loadModel(omc:: OMCSession, name:: String)`](@ref).
"""
function getSimulationSettings(omc:: OMCSession, name:: String; override=Dict())
    values = sendExpression(omc, "getSimulationOptions($name)")
    settings = Dict(
        "startTime"=>values[1], "stopTime"=>values[2],
        "tolerance"=>values[3], "numberOfIntervals"=>values[4],
        "outputFormat"=>"csv", "variableFilter"=>".*"
    )
    interval = values[5]
    settings["variableFilter"] = getVariableFilter(omc, name)
    for x in keys(settings)
        if x in keys(override)
            settings[x] = override[x]
        end
    end
    # the overriding of simulation time or interval size may require additional
    # changes to the numberOfIntervals setting
    hasinterval = haskey(override, "interval")
    onlytime = (haskey(override, "startTime") || haskey(override, "stopTime")
        && !haskey(override, "interval")
        && !haskey(override, "numberOfIntervals")
    )
    if hasinterval || onlytime
        timespan = settings["stopTime"] - settings["startTime"]
        interval = get(override, "interval", interval)
        settings["numberOfIntervals"]  = trunc(Int, timespan / interval)
    end
    return settings
end

"""
    getVariableFilter(omc:: OMCSession, name:: String)

Reads the value for the `variableFilter` simulation setting from the model
file if it has been defined.
MoST assumes that this value will be given in a vendor-specific annotation
of the form `__MoST_experiment(variableFilter=".*")`.
If such an annotation is not found, the default filter `".*"` is returned.

Throws a [`MoSTError`](@ref) if the model `name` does not exist.
"""
function getVariableFilter(omc:: OMCSession, name:: String)
    mostann = sendExpression(omc, "getAnnotationNamedModifiers($name, \"__MoST_experiment\")")
    if isnothing(mostann)
        throw(MoSTError("Model $name not found", ""))
    end
    varfilter = ".*"
    if "variableFilter" in mostann
        varfilter = sendExpression(omc, "getAnnotationModifierValue($name, \"__MoST_experiment\", \"variableFilter\")")
    end
    return varfilter
end

"""
    getVersion(omc:: OMCSession)

Returns the version of the OMCompiler as a triple (major, minor, patch).
"""
function getVersion(omc:: OMCSession)
    versionstring = sendExpression(omc, "getVersion()")
    # example: OMCompiler v1.17.0-dev.94+g4da66238ab
    # example: OpenModelica 1.14.2
    vmatch = match(r"^(?:OMCompiler v|OpenModelica )(\d+)\.(\d+).(\d+)", versionstring)
    if isnothing(vmatch)
        throw(MoSTError(omc, "Got unexpected version string: $versionstring"))
    end
    cap = map(x -> parse(Int, x), vmatch.captures)
    major, minor, patch = cap
    return Tuple([major, minor, patch])
end

"""
    simulate(omc:: OMCSession, name::String)
    simulate(omc:: OMCSession, name::String, settings:: Dict{String, Any})

Simulates the model `name` which must have been loaded before with
[`loadModel(omc:: OMCSession, name:: String)`](@ref).
The keyword-parameters in `settings` are directly passed to the OpenModelica
scripting function `simulate()`.
If the parameter is not given, it is obtained using
[`getSimulationSettings(omc:: OMCSession, name:: String; override=Dict())`](@ref).

The simulation output will be written to the current working directory of the
OMC that has been set by
[`setupOMCSession(outdir, modeldir; quiet=false, checkunits=true)`](@ref).

The simulation result is checked for errors with the following methods:

* The messages returned by the OM scripting call are checked for
    the string `Simulation execution failed`. This will, e.g., be the case
    if there is an arithmetic error during simulation.
* The abovementioned messages are checked for the string `| warning |` which
    hints at missing initial values and other non-critical errors.
* The error string returned by the OM scripting function `getErrorString()`
    should be empty if the simulation was successful.

If any of the abovementioned methods reveals errors, a [`MoSTError`](@ref)
is thrown.
""" # TODO which class of errors can be found using the error string?
function simulate(omc:: OMCSession, name::String, settings:: Dict{String, Any})
    prepare(s:: String) = "\"$(moescape(s))\""
    prepare(x:: Number) = x
    setstring = join(("$k=$(prepare(v))" for (k,v) in settings), ", ")
    r = sendExpression(omc, "simulate($name, $setstring)")
    if startswith(r["messages"], "Simulation execution failed")
        throw(MoSTError("Simulation of $name failed", r["messages"]))
    end
    if occursin("| warning |", r["messages"])
        throw(MoSTError("Simulation of $name produced warning", r["messages"]))
    end
    es = getErrorString(omc)
    if length(es) > 0
        throw(MoSTError("Simulation of $name failed", es))
    end
end
simulate(omc:: OMCSession, name::String) = simulate(omc, name, getSimulationSettings(omc, name))

"""
    avoidStartupFreeze(omc:: OMCSession)

Helper function to avoid freezes that can occur when the first message is sent
to a newly created OMCSession.

The current strategy for this is to detect the freeze, discard the frozen
session and create a new session, repeating this process until a non-frozen
connection is obtained.
"""
function avoidStartupFreeze(omc:: OMCSession)
    # TODO if this does not work, we can try this instead:
    #      https://github.com/JuliaInterop/ZMQ.jl/issues/198#issuecomment-576689600
    # sleep(0.5)
    function reconnect(omc:: OMCSession)
        @warn(string(
            "Discarding frozen connection to OMC with file descriptor $(omc.socket.fd)",
            " and starting new OMC instance. This may leave the old OMC",
            " instance still running on your machine."
        ))
        try
            closeOMCSession(omc)
        catch e
            @warn(string(
                "Closing old OMC instance failed with the following error:\n",
                e
            ))
        end
        return safeOMCSession()
    end
    status = :started
    timeout = 0.1
    while status != :received
        # send a simple command to OMC
        send(omc.socket, "getVersion()")
        # use julia task to allow recv to run into a timeout
        # idea from https://github.com/JuliaInterop/ZMQ.jl/issues/87#issuecomment-131153884
        c = Channel()
        @async put!(c, (recv(omc.socket), :received));
        @async (sleep(timeout); put!(c, (nothing, :timedout));)
        data, status = take!(c)
        if status == :timedout
            omc = reconnect(omc)
        end
    end
    return omc
end

"""
    safeOMCSession()

Helper function to avoid ZMQ.StateError that can occur when calling
`OMCSession()`.

The current strategy is to simply retry the connector call up to 10 times.
"""
function safeOMCSession()
    created = false
    tries = 0
    omc = nothing
    while !created && tries <= 10
        tries += 1
        try
            omc = OMCSession()
            created = true
        catch e
            if !isa(e, ZMQ.StateError)
                rethrow(e)
            end
            @warn(string(
                "OMCSession() constructor errored, attempting retry no $tries/10:\n",
                e
            ))
        end
    end
    if !created
        throw(MoSTError("OMCSession could not be created after $tries retries", ""))
    end
    return omc
end

"""
    setupOMCSession(outdir, modeldir; quiet=false, checkunits=true)

Creates an `OMCSession` and prepares it by preforming the following steps:

* create the directory `outdir` if it does not already exist
* change the working directory of the OMC to `outdir`
* add `modeldir` to the MODELICAPATH
* enable unit checking with the OMC command line option
    `--unitChecking` (unless `checkunits` is false)

If `quiet` is false, the resulting MODELICAPATH is printed to stdout.

Returns the newly created OMCSession.
"""
function setupOMCSession(outdir, modeldir; quiet=false, checkunits=true)
    # create output directory
    if !isdir(outdir)
        mkpath(outdir)
    end
    # create sessions
    omc = safeOMCSession()
    # sleep for a short while, because otherwise first ZMQ call may freeze
    omc = avoidStartupFreeze(omc)
    # move to output directory
    sendExpression(omc, "cd(\"$(moescape(outdir))\")")
    # set modelica path
    mopath = sendExpression(omc, "getModelicaPath()")
    mopath = "$mopath:$(moescape(abspath(modeldir)))"
    if !quiet
        println("Setting MODELICAPATH to ", mopath)
    end
    sendExpression(omc, "setModelicaPath(\"$mopath\")")
    # enable unit checking
    if checkunits
        flag = if getVersion(omc) >= Tuple([1, 16, 0])
            "--unitChecking"
        else
            "--preOptModules+=unitChecking"
        end
        sendExpression(omc, "setCommandLineOptions(\"$flag\")")
    end
    if !quiet
        opts = sendExpression(omc, "getCommandLineOptions()")
        println("Using command line options: $opts")
    end
    return omc
end

"""
    installAndLoad(omc:: OMCSession, lib:: AbstractString; version="latest")

Loads the Modelica library `lib` in version `version` and also installs it
if necessary.
"""
function installAndLoad(omc:: OMCSession, lib:: AbstractString; version="latest")
    lmver = version
    instver = version
    if getVersion(omc) < Tuple([1,16,0]) && version != "latest"
        # OpenModelica 1.14.2 does not have installPackage
        # => try to load default version and fail on error
        @warn(string(
            "Cannot install specific version $version of library $lib on OpenModelica < 1.16.0",
            " Attempting to load default version if it is installed."
        ))
        version = "latest"
    end
    if version == "latest"
        lmver = "default"
        instver = ""
    end
    sendExpression(omc, "loadModel($lib, {\"$lmver\"})")
    es = getErrorString(omc)
    if length(es) > 0 # can happen on OpenModelica 1.16 if MSL is not installed by default
        if getVersion(omc) < Tuple([1,16,0])
            throw(MoSTError("Cannot load library $lib with version $version on OpenModelica < 1.16.0", es))
        end
        sendExpression(omc, "installPackage($lib, \"$instver\")")
        # expected content: "Notification: Package installed successfully" for Modelica, ModelicaServices and Complex
        es = getErrorString(omc)
        if !occursin("Package installed successfully", es)
            throw(MoSTError("Failed to install library $lib with version $version", es))
        end
        sendExpression(omc, "loadModel($lib, {\"$lmver\"})")
        es = getErrorString(omc)
        if length(es) > 0
            throw(MoSTError("Failed to load library $lib with version $version after successful installation", es))
        end
    end
end

"""
    closeOMCSession(omc:: OMCSession; quiet=false)

Closes the OMCSession given by `omc`, shutting down the OMC instance.

Due to a [bug in the current release version of OMJulia](https://github.com/OpenModelica/jl/issues/32)
the function may occasionally freeze.
If this happens, you have to stop the execution with CTRL-C.
You can tell that this is the case if the output `Closing OMC session` is
printed to stdout, but it is not followed by `Done`.
If desired, these outputs can be disabled by setting `quiet=true`.

If you want to use a MoST.jl script for continuous integration, you can use
the following shell command to add a timeout to your script and treat the
timeout as a successful test run (which is, of course, unsafe).

```bash
(timeout 2m julia myTestScript.jl; rc=\$?; if [ \${rc} -eq 124 ]; then exit 0; else exit \${rc}; fi;)
```
"""
function closeOMCSession(omc:: OMCSession; quiet=false)
    if !quiet
        println("Closing OMC session")
    end
    # only send, do not wait for response since this may lead to freeze
    send(omc.socket, "quit()")
    # curently disabled due to error
    #close(omc.socket) # also close ZMQ socket
    if !quiet
        println("Done")
    end
end

"""
    withOMC(f:: Function, outdir, modeldir; quiet=false, checkunits=true)

Allows to use OMCSession with do-block syntax, automatically closing the
session after the block has been executed.
For the parameter definition see [`setupOMCSession(outdir, modeldir; quiet=false, checkunits=true)`](@ref).

Example:

```julia
withOMC("test/out", "test/res") do omc
    loadModel(omc, "Example")
end
```
"""
function withOMC(f:: Function, outdir, modeldir; quiet=false, checkunits=true)
    omc = setupOMCSession(outdir, modeldir; quiet=quiet, checkunits=checkunits)
    try
        f(omc)
    finally
        closeOMCSession(omc)
    end
end

max_stars_repo_name: THM-MoTE/MoST

max_issues_repo_name: THM-MoTE/MoST

max_forks_repo_name: THM-MoTE/MoST


text:
using BenchmarkTools

mutable struct Foo
    data::Int
end

function foo()
    f = Foo(2)
    r = 0
    for i = 1:1000_000_000
        r += i + f.data
    end
    return r
end

function foo2()
    r = 0
    for i = 1:1000_000_000
        r += i + 2
    end
    return r
end

baseline = @benchmark foo2()
bench = @benchmark foo()

display(baseline)
println()
display(bench)


text:
function trunc(x)
if x < 0
return 0
else
return x
end
end

max_stars_repo_path: Module 2/Chapter03/trunc.jl

max_issues_repo_path: Module 2/Chapter03/trunc.jl

max_forks_repo_path: Module 2/Chapter03/trunc.jl


text:
# Computes a polynomial Lyapunov for the switched system from Example 2.8 of
# [PJ08].
#   x_{k+1} = A_σ x_k, σ ∈ {1, 2}
# Where
#   A1 = [1 0   A2 = [0  1
#         1 0]        0 -1]
# There is a quartic Lyapunov function for the system but only a quadratic for
# the system with A1' = A1/√2, A2' = A2/√2.
# See https://github.com/blegat/SwitchOnSafety.jl for computing Lyapunov
# functions for switched systems.
#
# [PJ08] P. Parrilo and A. Jadbabaie
# Approximation of the joint spectral radius using sum of squares
# Linear Algebra and its Applications, Elsevier, 2008, 428, 2385-2402
# Inspired from the construction of:
# Ando, T. and Shih, M.-h.
# Simultaneous Contractibility.
# SIAM Journal on Matrix Analysis & Applications, 1998, 19, 487

using LinearAlgebra, Test
using SumOfSquares
using DynamicPolynomials

# Search for polynomial Lyapunov functions of degree `2d` for `A1/γ` and `A2/γ`
function lyapunov_switched_system_test(
    optimizer, config::MOIT.Config, degree::Int, γ, feasible::Bool, basis)

    atol = config.atol
    rtol = config.rtol

    A1 = [1 0
          1 0]
    A2 = [0  1
          0 -1]

    model = _model(optimizer)

    @polyvar x[1:2]

    # p is should be strictly positive so we create a nonnegative `p`
    # and sum it with a strictly positive `q`.
    # Since the problem is homogeneous (i.e. given any `λ > 0`, `p` is
    # feasible iff `λp` is feasible), this is wlog.
    p0 = @variable(model, variable_type=SOSPoly(basis(monomials(x, degree))))
    q = GramMatrix(SOSDecomposition(x.^degree))

    # Keep `p` in a `GramMatrix` form while `q + p0` would transform it to
    # a polynomial. It is not mandatory to keep it in its `GramMatrix` form
    # but it will allow us to test `JuMP.value(::GramMatrix{JuMP.AffExpr})`.
    p = gram_operate(+, q, p0)

    c1 = @constraint(model, p - p(x => (A1 / γ) * vec(x)) in SOSCone(), basis = basis)
    c2 = @constraint(model, p - p(x => (A2 / γ) * vec(x)) in SOSCone(), basis = basis)

    JuMP.optimize!(model)

    if feasible
        @test JuMP.termination_status(model) == MOI.OPTIMAL
        @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
        @test all(eigvals(Matrix(getmat(JuMP.value(p0)))) .≥ -atol)
        @test JuMP.value(p0).basis isa basis
        @test getmat(JuMP.value(p)) ≈ getmat(gram_operate(+, q, JuMP.value(p0))) atol=atol rtol=rtol
        @test gram_matrix(c1).basis isa basis
        @test gram_matrix(c2).basis isa basis
    else
        @test JuMP.termination_status(model) == MOI.INFEASIBLE
        @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE
        for (μ1, μ2) in [(JuMP.dual(c1), JuMP.dual(c2)),
                         (moments(c1), moments(c2))]
            # The dual constraint should work on any polynomial.
            # Let's test it with q
            lhs = dot(μ1, q(x => A1 * vec(x))) + dot(μ2, q(x => A2 * vec(x)))
            rhs = dot(μ1, q) + dot(μ2, q)
            @test atol + rtol * max(abs(lhs), abs(rhs)) + lhs >= rhs
        end
        @test moment_matrix(c1).basis isa basis
        @test moment_matrix(c2).basis isa basis
    end
end

quadratic_infeasible_lyapunov_switched_system_test(optimizer, config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 - ε, false, MonomialBasis)
sd_tests["quadratic_infeasible_lyapunov_switched_system"] = quadratic_infeasible_lyapunov_switched_system_test
quadratic_infeasible_scaled_lyapunov_switched_system_test(optimizer, config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 - ε, false, ScaledMonomialBasis)
sd_tests["quadratic_infeasible_scaled_lyapunov_switched_system"] = quadratic_infeasible_scaled_lyapunov_switched_system_test
quadratic_feasible_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 + ε, true, MonomialBasis)
sd_tests["quadratic_feasible_lyapunov_switched_system"]   = quadratic_feasible_lyapunov_switched_system_test
quadratic_feasible_scaled_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 1, √2 + ε, true, ScaledMonomialBasis)
sd_tests["quadratic_feasible_scaled_lyapunov_switched_system"]   = quadratic_feasible_scaled_lyapunov_switched_system_test
quartic_infeasible_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 - ε, false, MonomialBasis)
sd_tests["quartic_infeasible_lyapunov_switched_system"]   = quartic_infeasible_lyapunov_switched_system_test
quartic_infeasible_scaled_lyapunov_switched_system_test(optimizer,   config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 - ε, false, ScaledMonomialBasis)
sd_tests["quartic_infeasible_scaled_lyapunov_switched_system"]   = quartic_infeasible_scaled_lyapunov_switched_system_test
quartic_feasible_lyapunov_switched_system_test(optimizer,     config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 + ε, true, MonomialBasis)
sd_tests["quartic_feasible_lyapunov_switched_system"]     = quartic_feasible_lyapunov_switched_system_test
quartic_feasible_scaled_lyapunov_switched_system_test(optimizer,     config, ε = 2e-1) = lyapunov_switched_system_test(optimizer, config, 2,  1 + ε, true, ScaledMonomialBasis)
sd_tests["quartic_feasible_scaled_lyapunov_switched_system"]     = quartic_feasible_scaled_lyapunov_switched_system_test

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares


text:
struct CoNLL2000{S}
    filepaths::Vector{S}
    year::Int
    trainpath::String
    testpath::String
end

function CoNLL2000(path1, path2, year=2000)
    @assert(isdir(path1), path1)
    @assert(isdir(path2), path2)

    files = Dict("train" => joinpath(path1, "train.txt"),
             "test" => joinpath(path2, "test.txt"))

    return CoNLL2000(collect(values(files)), year, files["train"],
                  files["test"])
end

CoNLL2000() = CoNLL2000(datadep"CoNLL 2000 train", datadep"CoNLL 2000 test")

MultiResolutionIterators.levelname_map(::Type{CoNLL2000}) = [
    :sent=>1, :sentence=>1,
    :word=>2, :token=>2,
    :char=>3, :character=>3
    ]

function parse_conll2000_tagged_word(line::AbstractString)
    tokens_tags = split(line)
    length(tokens_tags) != 3 && throw("Error parsing line: \"$line\". Invalid Format.")
    return POSTaggedWord(tokens_tags[2], tokens_tags[3], tokens_tags[1])
end

function parse_conll2000file(filename)
    local sent
    sents = @NestedVector(POSTaggedWord, 2)()
    context = Document(intern(basename(filename)), sents)

    function new_sentence()
        sent = @NestedVector(POSTaggedWord,1)()
        push!(sents, sent)
    end

    # words
    get_tagged(line) = push!(sent, parse_conll2000_tagged_word(line))

    new_sentence()
    # parse
    for line in eachline(filename)
        if length(line) == 0 && length(sent) > 0
            new_sentence()
        else
            get_tagged(line)
        end
    end
    isempty(sents[end]) && deleteat!(sents, lastindex(sents))

    return context
end

function load(corpus::CoNLL2000, file="train")
    file == "train" && return parse_conll2000file(corpus.trainpath)
    file == "test" && return parse_conll2000file(corpus.testpath)
    file == "dev" && throw("Dev set not available in $year.")
    throw("Invalid filename! Available datasets are `train` & `test`")
end


text:
nextL(L, l::Integer) = L*l
nextL(L, r::AbstractUnitRange) = L*unsafe_length(r)
nextL(L, r::Slice) = L*unsafe_length(r.indices)
offsetin(i, l::Integer) = i-1
offsetin(i, r::AbstractUnitRange) = i-first(r)

# Generic cases
_sub2ind(dims::DimsInteger, I::Integer...) = (@_inline_meta; _sub2ind_recurse(dims, 1, 1, I...))
_sub2ind(inds::Indices, I::Integer...) = (@_inline_meta; _sub2ind_recurse(inds, 1, 1, I...))

# In 1d, there's a question of whether we're doing cartesian indexing
# or linear indexing. Support only the former.
_sub2ind(inds::Indices{1}, I::Integer...) =
    throw(ArgumentError("Linear indexing is not defined for one-dimensional arrays"))
_sub2ind(inds::Tuple{OneTo}, I::Integer...) = (@_inline_meta; _sub2ind_recurse(inds, 1, 1, I...)) # only OneTo is safe
_sub2ind(inds::Tuple{OneTo}, i::Integer)    = i

@propagate_inbounds function _to_linear(axs, L, idx, inds)
end
_sub2ind_recurse(axs::AbstractIndex, L, ind) = ind
function _sub2ind_recurse(::Tuple{}, L, ind, i::Integer, I::Integer...)
    @_inline_meta
    _sub2ind_recurse((), L, ind+(i-1)*L, I...)
end
function _sub2ind_recurse(axs, L, idx, inds::Tuple)
    @_inline_meta
    r1 = first(axs)
    i = to_index(first(axs), )
    _sub2ind_recurse(tail(axs), nextL(L, r1),
                     ind + offsetin(i, r1) * L,
                     first(I), tail(inds)
                    )
end

# Vectorized forms
function _sub2ind(inds::Indices{1}, I1::AbstractVector{T}, I::AbstractVector{T}...) where T<:Integer
    throw(ArgumentError("Linear indexing is not defined for one-dimensional arrays"))
end
_sub2ind(inds::Tuple{OneTo}, I1::AbstractVector{T}, I::AbstractVector{T}...) where {T<:Integer} =
    _sub2ind_vecs(inds, I1, I...)
_sub2ind(inds::Union{DimsInteger,Indices}, I1::AbstractVector{T}, I::AbstractVector{T}...) where {T<:Integer} =
    _sub2ind_vecs(inds, I1, I...)
function _sub2ind_vecs(inds, I::AbstractVector...)
    I1 = I[1]
    Iinds = axes1(I1)
    for j = 2:length(I)
        axes1(I[j]) == Iinds || throw(DimensionMismatch("indices of I[1] ($(Iinds)) does not match indices of I[$j] ($(axes1(I[j])))"))
    end
    Iout = similar(I1)
    _sub2ind!(Iout, inds, Iinds, I)
    Iout
end

function _sub2ind!(Iout, inds, Iinds, I)
    @_noinline_meta
    for i in Iinds
        # Iout[i] = _sub2ind(inds, map(Ij -> Ij[i], I)...)
        Iout[i] = sub2ind_vec(inds, i, I)
    end
    Iout
end

sub2ind_vec(inds, i, I) = (@_inline_meta; _sub2ind(inds, _sub2ind_vec(i, I...)...))
_sub2ind_vec(i, I1, I...) = (@_inline_meta; (I1[i], _sub2ind_vec(i, I...)...))
_sub2ind_vec(i) = ()