function get_sim_data_indirect(mechanism::Mechanism,env::Environment,Δt::Real;
                               relax_comp=false)
    vs = VariableSelector()
    add_var!(vs, :qnext, num_positions(mechanism))
    add_var!(vs, :vnext, num_velocities(mechanism))
    for i = 1:length(env.contacts)
        β_dim = size(env.contacts[i].obstacle.basis,2)
        add_var!(vs, Symbol("β", i), β_dim)
        add_var!(vs, Symbol("λ", i), 1)
        add_var!(vs, Symbol("c_n", i), 1)
        if relax_comp
            add_var!(vs, Symbol("slack", i), 3)
        end
    end

    cs = ConstraintSelector()
    add_eq!(cs, :kin, num_positions(mechanism))
    add_eq!(cs, :dyn, num_velocities(mechanism))
    for i = 1:length(env.contacts)
        β_dim = size(env.contacts[i].obstacle.basis,2)
        add_ineq!(cs, Symbol("β_pos", i), β_dim)
        add_ineq!(cs, Symbol("λ_pos", i), 1)
        add_ineq!(cs, Symbol("c_n_pos", i), 1)
        add_ineq!(cs, Symbol("ϕ_pos", i), 1)
        add_ineq!(cs, Symbol("fric_pos", i), β_dim)
        add_ineq!(cs, Symbol("cone_pos", i), 1)
        if relax_comp
            add_ineq!(cs, Symbol("slack_pos", i), 3)
            add_ineq!(cs, Symbol("ϕ_c_n_comp", i), 1)
            add_ineq!(cs, Symbol("fric_β_comp", i), β_dim)
            add_ineq!(cs, Symbol("cone_λ_comp", i), 1)
        else
            add_eq!(cs, Symbol("ϕ_c_n_comp", i), 1)
            add_eq!(cs, Symbol("fric_β_comp", i), β_dim)
            add_eq!(cs, Symbol("cone_λ_comp", i), 1)
        end
    end

    state_cache = [StateCache(mechanism) for n = 1:2]
    envj_cache = [EnvironmentJacobianCache(env) for n = 1:2]

    generate_solver_fn = :generate_solver_fn_sim_indirect
    extract_sol = :extract_sol_sim_indirect

    sim_data = SimData(mechanism,env,
                       state_cache,envj_cache,
                       Δt,vs,cs,generate_solver_fn,extract_sol,
                       [],[],[],[],[],[],[],[],[],1,[],[],[],[])

    sim_data
end

function extract_sol_sim_indirect(sim_data::SimData, results::AbstractArray{T,2}) where T
    vs = sim_data.vs
    relax_comp = haskey(vs.vars, :slack1)
    env = sim_data.env
    N = size(results,2)

    qtraj = Array{Array{Float64,1},1}(undef, 0)
    vtraj = Array{Array{Float64,1},1}(undef, 0)
    utraj = Array{Array{Float64,1},1}(undef, 0)
    contact_traj = Array{Array{Array{Float64,1},1},1}(undef, 0)
    slack_traj = Array{Array{Float64,1},1}(undef, 0)
    x = MechanismState(sim_data.mechanism)
    for n = 1:N
        set_configuration!(x, vs(results[:,n], Symbol("qnext")))
        normalize_configuration!(x)
        qnext = configuration(x)
        push!(qtraj, qnext)
        push!(vtraj, vs(results[:,n], Symbol("vnext")))
        contact_sol = Array{Array{Float64,1},1}(undef, 0)
        for i = 1:length(env.contacts)
            if relax_comp
                push!(slack_traj, vs(results[:,n], Symbol("slack", i)))
            end
            push!(contact_sol, vcat(vs(results[:,n], Symbol("c_n", i)),
                               vs(results[:,n], Symbol("β", i)),
                               vs(results[:,n], Symbol("λ", i))))
        end
        push!(contact_traj, contact_sol)
    end

    # some other usefull vectors
    ttraj = [(i-1)*sim_data.Δt for i = 1:N]
    qv_mat = vcat(hcat(qtraj...),hcat(vtraj...))

    qtraj, vtraj, utraj, contact_traj, slack_traj, ttraj, qv_mat, results
end

function generate_solver_fn_sim_indirect(sim_data,q0,v0,u0)
    x0 = sim_data.state_cache[1][Float64]
    Δt = sim_data.Δt
    vs = sim_data.vs
    cs = sim_data.cs

    relax_comp = haskey(vs.vars, :slack1)
    num_contacts = length(sim_data.env.contacts)
    num_vel = num_velocities(sim_data.mechanism)
    world = root_body(sim_data.mechanism)
    world_frame = default_frame(world)

    set_configuration!(x0, q0)
    set_velocity!(x0, v0)
    setdirty!(x0)
    H = mass_matrix(x0)

    function eval_obj(x::AbstractArray{T}) where T
        f = 0.

        if relax_comp
            for i = 1:num_contacts
                slack = vs(x, Symbol("slack", i))
                f += sum(slack)
            end
        end

        f
    end

    function eval_cons(x::AbstractArray{T}) where T
        xn = sim_data.state_cache[2][T]
        envj = sim_data.envj_cache[2][T]

        contact_bias = Vector{T}(undef, num_vel)
        g = Vector{T}(undef, cs.num_eqs + cs.num_ineqs) # TODO preallocate

        qnext = vs(x, :qnext)
        vnext = vs(x, :vnext)

        set_configuration!(xn, qnext)
        set_velocity!(xn, vnext)
        setdirty!(xn)

        normalize_configuration!(xn)

        config_derivative = configuration_derivative(xn) # TODO preallocate
        dyn_bias = dynamics_bias(xn) # TODO preallocate
        if (num_contacts > 0)
            contact_jacobian!(envj, xn)
            contact_τ_indirect!(contact_bias, sim_data, envj, x)
        end

        g[cs(:kin)] .= qnext .- q0 .- Δt .* config_derivative
        g[cs(:dyn)] .= H * (vnext - v0) .- Δt .* (u0 .- dyn_bias .- contact_bias)
        for i = 1:num_contacts
            cj = envj.contact_jacobians[i]
            c = cj.contact

            # TODO preallocate
            contact_v = cj.contact_rot' * point_jacobian(xn, path(sim_data.mechanism, world, c.body), transform(xn, c.point, world_frame)).J * vnext

            β = vs(x, Symbol("β", i))
            λ = vs(x, Symbol("λ", i))
            c_n = vs(x, Symbol("c_n", i))

            Dtv = c.obstacle.basis' * contact_v

            g[cs(Symbol("β_pos", i))] .= -β
            g[cs(Symbol("λ_pos", i))] .= -λ
            g[cs(Symbol("c_n_pos", i))] .= -c_n
            g[cs(Symbol("ϕ_pos", i))] .= -envj.contact_jacobians[i].ϕ
            # λe + D'*v >= 0
            g[cs(Symbol("fric_pos", i))] .= -(λ .+ Dtv)
            # μ*c_n - sum(β) >= 0
            g[cs(Symbol("cone_pos", i))] .= -(c.obstacle.μ .* c_n .- sum(β))
            if !relax_comp
                # dist * c_n = 0
                g[cs(Symbol("ϕ_c_n_comp", i))] .= envj.contact_jacobians[i].ϕ .* c_n
                # (λe + Dtv)' * β = 0
                g[cs(Symbol("fric_β_comp", i))] .= (λ .+ Dtv) .* β
                # (μ * c_n - sum(β)) * λ = 0
                g[cs(Symbol("cone_λ_comp", i))] .= (c.obstacle.μ .* c_n .- sum(β)) .* λ
            else
                slack = vs(x, Symbol("slack", i))
                g[cs(Symbol("slack_pos", i))] .= -slack
                g[cs(Symbol("ϕ_c_n_comp", i))] .= envj.contact_jacobians[i].ϕ .* c_n .- slack[1]
                g[cs(Symbol("fric_β_comp", i))] .= (λ .+ Dtv) .* β .- slack[2]
                g[cs(Symbol("cone_λ_comp", i))] .= (c.obstacle.μ .* c_n .- sum(β)) .* λ .- slack[3]
            end
        end

        g
    end

    generate_autodiff_solver_fn(eval_obj,eval_cons,cs.eqs,cs.ineqs,vs.num_vars)
end


function optimize!(structure::VerticalStructure, optimizer::AbstractStructuredOptimizer, x₀::AbstractVector)
    # Sanity checks
    supports_structure(optimizer, structure) || throw(UnsupportedStructure{typeof(optimizer), typeof(structure)}())
    check_loadable(optimizer, structure)
    # Load structure
    load_structure!(optimizer, structure, x₀)
    # Run structure-exploiting optimization procedure
    MOI.optimize!(optimizer)
    return nothing
end

function set_master_optimizer!(structure::VerticalStructure, optimizer)
    set_optimizer(structure.first_stage, optimizer)
    return nothing
end

function set_master_optimizer_attribute!(structure::VerticalStructure, attr::MOI.AbstractOptimizerAttribute, value)
    MOI.set(backend(structure.first_stage), attr, value)
    return nothing
end

function set_subproblem_optimizer!(structure::VerticalStructure, optimizer)
    set_optimizer!(scenarioproblems(structure), optimizer)
    return nothing
end

function set_subproblem_optimizer_attribute!(structure::VerticalStructure, attr::MOI.AbstractOptimizerAttribute, value)
    MOI.set(scenarioproblems(structure), attr, value)
    return nothing
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Junction <: Topology
        name::String
        elevation::Float64
        head::Union{Nothing, Float64}
        coordinates::Union{Nothing, NamedTuple{(:lat, :lon),Tuple{Float64,Float64}}}
        pattern_name::Union{Nothing, String}
        _forecasts::InfrastructureSystems.Forecasts
        internal::InfrastructureSystemsInternal
    end

A water-system Junction.

# Arguments
- `name::String`: the name of the junction
- `elevation::Float64`: elevation of junction
- `head::Union{Nothing, Float64}`: pressure head at junction
- `coordinates::Union{Nothing, NamedTuple{(:lat, :lon),Tuple{Float64,Float64}}}`: latitude and longitude coordinates of junction
- `pattern_name::Union{Nothing, String}`: name of head pattern array, for reservoirs only
- `_forecasts::InfrastructureSystems.Forecasts`
- `internal::InfrastructureSystemsInternal`: internal reference, do not modify
"""
mutable struct Junction <: Topology
    "the name of the junction"
    name::String
    "elevation of junction"
    elevation::Float64
    "pressure head at junction"
    head::Union{Nothing, Float64}
    "latitude and longitude coordinates of junction"
    coordinates::Union{Nothing, NamedTuple{(:lat, :lon),Tuple{Float64,Float64}}}
    "name of head pattern array, for reservoirs only"
    pattern_name::Union{Nothing, String}
    _forecasts::InfrastructureSystems.Forecasts
    "internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function Junction(name, elevation, head, coordinates, pattern_name, _forecasts=InfrastructureSystems.Forecasts(), )
    Junction(name, elevation, head, coordinates, pattern_name, _forecasts, InfrastructureSystemsInternal(), )
end

function Junction(; name, elevation, head, coordinates, pattern_name, _forecasts=InfrastructureSystems.Forecasts(), )
    Junction(name, elevation, head, coordinates, pattern_name, _forecasts, )
end

# Constructor for demo purposes; non-functional.
function Junction(::Nothing)
    Junction(;
        name="init",
        elevation=0.0,
        head=nothing,
        coordinates=nothing,
        pattern_name=nothing,
        _forecasts=InfrastructureSystems.Forecasts(),
    )
end

"""Get Junction name."""
get_name(value::Junction) = value.name
"""Get Junction elevation."""
get_elevation(value::Junction) = value.elevation
"""Get Junction head."""
get_head(value::Junction) = value.head
"""Get Junction coordinates."""
get_coordinates(value::Junction) = value.coordinates
"""Get Junction pattern_name."""
get_pattern_name(value::Junction) = value.pattern_name
"""Get Junction _forecasts."""
get__forecasts(value::Junction) = value._forecasts
"""Get Junction internal."""
get_internal(value::Junction) = value.internal


# Timers.

"""
Global vector of timers `(deadline, condition)` sorted by `deadline`.
"""
struct Timer
    deadline::Float64
    f::Any
end

Base.isless(a::Timer, b::Timer) = isless(a.deadline, b.deadline)

const timer_vector = Timer[]
const timer_lock = Threads.SpinLock()

function dump_timer_state(msg="")
    now = time()
    printerr("Timers $msg: ", [round(t.deadline - now; digits=2)
                               for t in timer_vector])
end

"""
    UnixIO.sleep(seconds)

Block the current task for a specified number of seconds.
(Waits for a timer based on `UnixIO.poll_wait()`).
"""
sleep(seconds::Float64) = wait_timeout(sleep_condition, seconds)

const sleep_condition = Base.ThreadSynchronizer()

"""
Restart the clock source responsible for calling `notify_timers()`.
The clock source will be sleeping based on `next_timer_deadline_ms()`,
if timer is registered with a nearer deadline, the clock source needs
to be restarted (it will call `next_timer_deadline_ms()` again and go
back to sleep).
"""
@db 3 function restart_timer_clock()
    wakeup_poll(poll_queue)
end

"""
    next_timer_deadline_ms(timeout_ms)

Time from now to the next timer deadline in ms.
If no timers due before `timeout_ms` return `timeout_ms`.
"""
function next_timer_deadline_ms(timeout_ms::Int)
    @dblock timer_lock begin
        deadline = next_timer_deadline()
        dt = deadline - time()
        if dt < (timeout_ms/1000)
            timeout_ms = max(0, round(Int, dt * 1000))
        end
    end
    return timeout_ms
end

function next_timer_deadline()
    assert_havelock(timer_lock)
    isempty(timer_vector) ? Inf : timer_vector[1].deadline
end

const inf_timer = Timer(Inf, nothing)

"""
    register_timer(f, deadline)

Register `f`` to be called at `time()` `deadline`.
"""
@db 5 function register_timer(f, deadline)
    if deadline == Inf
        return inf_timer
    end
    @db 5 return register_timer(Timer(deadline, f))
end

@db 3 function register_timer(t::Timer)
    i = 0
    @dblock timer_lock begin
        i = searchsortedfirst(timer_vector, t)
        insert!(timer_vector, i, t)
    end
    if i == 1                                    ;@db 3 "restart_timer_clock()"
        restart_timer_clock()
    end
    @db 3 return t
end

@db 3 function cancel_timer(t::Timer)
    @require t.deadline !=  Inf
    @dblock timer_lock begin
        if t.deadline >= next_timer_deadline()
            i = searchsortedfirst(timer_vector, t)
            if timer_vector[i] == t                              ;@db 3 "⏱ 🚫"
                deleteat!(timer_vector, i)
            end
        end
    end
end

"""
Cancel a timer `(deadline, condition)` without notifying `condition`.
"""
@db 5 function Base.close(t::Timer)
    if t.deadline != Inf
        cancel_timer(t)
    end
    nothing
end

"""
Notify timers whos `deadline` has passed and remove from vector.
"""
@db 6 function notify_timers(now = time())
    ready=[]
    @dblock timer_lock begin
        expired = 0
        for (i, t) in enumerate(timer_vector)
            if t.deadline > now
                break
            end
            push!(ready, t.f)
            expired = i
        end
        deleteat!(timer_vector, 1:expired)
    end
    for f in ready                                            ;@db 3 "⏱ -> $f"
        f()
    end
    nothing
end

"""
    wait_until(condition, [predicate=()->false], deadline)

If `predicate()` is false, repeatently `wait(condition)` until
`predicate()` is true (or until `deadline` is reached).
"""
@db 3 function wait_until(condition::Base.GenericCondition,
                          predicate,
                          deadline::Float64)
    if predicate()
        return
    end
    @dblock condition begin
        timer = register_timer(deadline) do
            @dblock condition notify(condition)
        end
        try
            while !predicate() && time() < deadline         ;@db 3 "waiting..."
                wait(condition)
            end
        finally
            close(timer)
        end
    end
    nothing
end

wait_until(c, d) = wait_until(c, ()->false, d)
@selfdoc wait_timeout(c, t) = wait_until(c, time() + Float64(t))

# Pretty Printing.

function dbshow(io::IO, t::Timer)
    print(io, "Timer(", debug_time(t.deadline),
                        " [", round(t.deadline - time(); digits=2), "])")
end

# End of file: timer.jl


using ChainRulesCore: NoTangent
using Base.Experimental: @opaque

struct ∂⃖rrule{N}; end
struct ∂⃖recurse{N}; end

include("recurse.jl")

function perform_optic_transform(@nospecialize(ff::Type{∂⃖recurse{N}}), @nospecialize(args)) where {N}
    @assert N >= 1

    # Check if we have an rrule for this function
    mthds = Base._methods_by_ftype(Tuple{args...}, -1, typemax(UInt))
    if length(mthds) != 1
        return :(throw(MethodError(ff, args)))
    end
    match = mthds[1]

    mi = Core.Compiler.specialize_method(match)
    ci = Core.Compiler.retrieve_code_info(mi)

    ci′ = copy(ci)
    ci′.edges = MethodInstance[mi]

    r = transform!(ci′, mi.def, length(args) - 1, match.sparams, N)
    if isa(r, Expr)
        return r
    end

    ci′.ssavaluetypes = length(ci′.code)
    ci′.method_for_inference_limit_heuristics = match.method
    ci′
end

# This relies on PartialStruct to infer well
struct Protected{N}
    a
end
(p::Protected)(args...) = getfield(p, :a)(args...)[1]
@Base.aggressive_constprop (::∂⃖{N})(p::Protected{N}, args...) where {N} = getfield(p, :a)(args...)
@Base.aggressive_constprop (::∂⃖{1})(p::Protected{1}, args...) = getfield(p, :a)(args...)
(::∂⃖{N})(p::Protected, args...) where {N} = error("TODO: Can we support this?")

struct OpticBundle{T}
    x::T
    clos
end
Base.getindex(o::OpticBundle, i::Int) = i == 1 ? o.x :
                                        i == 2 ? o.clos :
                                        throw(BoundsError(o, i))
Base.iterate(o::OpticBundle) = (o.x, nothing)
Base.iterate(o::OpticBundle, ::Nothing) = (o.clos, missing)
Base.iterate(o::OpticBundle, ::Missing) = nothing

# Desturucture using `getfield` rather than iterate to make
# inference happier
macro destruct(arg)
    @assert isexpr(arg, :(=))
    lhs = arg.args[1]
    @assert isexpr(lhs, :tuple)
    rhs = arg.args[2]
    s = gensym()
    quote
        $s = $(esc(rhs))
        $([:($(esc(arg)) = getfield($s, $i)) for (i, arg) in enumerate(lhs.args)]...)
        $s
    end
end

@eval function (::∂⃖{1})(::∂⃖{2}, args...)
    @destruct a, b = ∂⃖{3}()(args...)
    (a, $(Expr(:new, Protected{1}, :(Δ->begin
        x1, x2 = b(Δ)
        (x1, $(Expr(:new, Protected{1}, :((Δ...)->begin
            x3, x4 = x2(Δ...)
            (x3, $(Expr(:new, Protected{1}, :(Δ->begin
                x5, x6 = x4(Δ)
                x5, Δ->begin
                    (x7, x8) = x6(Δ...)
                    return (x8, x7)
                end
            end)))), ((x7, x8),)->begin
                (x9, x10) = x8(x7)
                (x10, x9...)
            end
        end)))), ((x9, x10),)->begin
            (x11, x12) = x10(x9...)
            (x12, x11)
        end
    end)))), ((x11, x12),)->begin
        (NoTangent(), x12(x11)...)
    end
end

struct ∂⃖weaveInnerOdd{N, O}; b̄; end
@Base.aggressive_constprop function (w::∂⃖weaveInnerOdd{N, N})(Δ) where {N}
    @destruct c, c̄ = w.b̄(Δ...)
    return (c̄, c)
end
@Base.aggressive_constprop function (w::∂⃖weaveInnerOdd{N, O})(Δ) where {N, O}
    @destruct c, c̄ = w.b̄(Δ...)
    return (c̄, c), ∂⃖weaveInnerEven{plus1(N), O}()
end
struct ∂⃖weaveInnerEven{N, O}; end
@Base.aggressive_constprop function (w::∂⃖weaveInnerEven{N, O})(Δ′, x...) where {N, O}
    @destruct y, ȳ  = Δ′(x...)
    return y, ∂⃖weaveInnerOdd{plus1(N), O}(ȳ)
end

struct ∂⃖weaveOuterOdd{N, O}; end
@Base.aggressive_constprop function (w::∂⃖weaveOuterOdd{N, N})((Δ′′, Δ′′′)) where {N}
    return (NoTangent(), Δ′′′(Δ′′)...)
end
@Base.aggressive_constprop function (w::∂⃖weaveOuterOdd{N, O})((Δ′′, Δ′′′)) where {N, O}
    @destruct α, ᾱ = Δ′′′(Δ′′)
    return (NoTangent(), α...), ∂⃖weaveOuterEven{plus1(N), O}(ᾱ)
end
struct ∂⃖weaveOuterEven{N, O}; ᾱ end
@Base.aggressive_constprop function (w::∂⃖weaveOuterEven{N, O})(Δ⁴...) where {N, O}
    return w.ᾱ(Base.tail(Δ⁴)...), ∂⃖weaveOuterOdd{plus1(N), O}()
end

function (::∂⃖{N})(::∂⃖{1}, args...) where {N}
    @destruct (a, ā) = ∂⃖{plus1(N)}()(args...)
    let O = c_order(N)
        (a, Protected{N}(@opaque Δ->begin
                (b, b̄) = ā(Δ)
                b, ∂⃖weaveInnerOdd{1, O}(b̄)
            end
        )), ∂⃖weaveOuterOdd{1, O}()
    end
end

function (::∂⃖{N})(::∂⃖{M}, args...) where {N, M}
    # TODO
    @destruct (a, b) = ∂⃖{N+M}()(args...)
    error("Not implemented yet ($N, $M)")
end

macro OpticBundle(a, b)
    aa = gensym()
    esc(quote
        $aa = $a
        $(Expr(:new, :(OpticBundle{typeof($aa)}), aa, b))
    end)
end

# ∂⃖rrule has a 4-recurrence - we model this as 4 separate structs that we
# cycle between. N.B.: These names match the names that these variables
# have in Snippet 19 of the terminology guide. They are probably not ideal,
# but if you rename them here, please update the terminology guide also.

struct ∂⃖rruleA{N, O}; ∂; ȳ; ȳ̄ ; end
struct ∂⃖rruleB{N, O}; ᾱ; ȳ̄ ; end
struct ∂⃖rruleC{N, O}; ȳ̄ ; Δ′′′; β̄ ; end
struct ∂⃖rruleD{N, O}; γ̄; β̄ ; end

@Base.aggressive_constprop function (a::∂⃖rruleA{N, O})(Δ) where {N, O}
    # TODO: Is this unthunk in the right place
    @destruct (α, ᾱ) = a.∂(a.ȳ, unthunk(Δ))
    (α, ∂⃖rruleB{N, O}(ᾱ, a.ȳ̄))
end

@Base.aggressive_constprop function (b::∂⃖rruleB{N, O})(Δ′...) where {N, O}
    @destruct ((Δ′′′, β), β̄) = b.ᾱ(Δ′)
    (β, ∂⃖rruleC{N, O}(b.ȳ̄, Δ′′′, β̄))
end

@Base.aggressive_constprop function (c::∂⃖rruleC{N, O})(Δ′′) where {N, O}
    @destruct (γ, γ̄) = c.ȳ̄((Δ′′, c.Δ′′′))
    (Base.tail(γ), ∂⃖rruleD{N, O}(γ̄, c.β̄))
end

@Base.aggressive_constprop function (d::∂⃖rruleD{N, O})(Δ⁴...) where {N, O}
    (δ₁, δ₂), δ̄  = d.γ̄(ZeroTangent(), Δ⁴...)
    (δ₁, ∂⃖rruleA{N, O+1}(d.β̄ , δ₂, δ̄ ))
end

# Terminal cases
@Base.aggressive_constprop function (c::∂⃖rruleB{N, N})(Δ′...) where {N}
    @destruct (Δ′′′, β) = c.ᾱ(Δ′)
    (β, ∂⃖rruleC{N, N}(c.ȳ̄, Δ′′′, nothing))
end
@Base.aggressive_constprop (c::∂⃖rruleC{N, N})(Δ′′) where {N} =
    Base.tail(c.ȳ̄((Δ′′, c.Δ′′′)))
(::∂⃖rruleD{N, N})(Δ...) where {N} = error("Should not be reached")

# ∂⃖rrule
@Base.pure term_depth(N) = 2^(N-2)
function (::∂⃖rrule{N})(z, z̄) where {N}
    @destruct (y, ȳ) = z
    y, ∂⃖rruleA{term_depth(N), 1}(∂⃖{minus1(N)}(), ȳ, z̄)
end

function (::∂⃖{N})(f::Core.IntrinsicFunction, args...) where {N}
    # A few intrinsic functions are inserted by the compiler, so they need to
    # be handled here. Otherwise, we just throw an appropriate error.
    if f === Core.Intrinsics.not_int && length(args) == 1
        return f(args...), EvenOddOdd{1, c_order(N)}(
            Δ->(NoTangent(), NoTangent()),
            Δ->NoTangent())
    end

    error("Rewrite reached intrinsic function $f. Missing rule?")
end

# The static parameter on `f` disables the compileable_sig heuristic
function (::∂⃖{N})(f::T, args...) where {T, N}
    if N == 1
        # Base case (inlined to avoid ambiguities with manually specified
        # higher order rules)
        z = rrule(f, args...)
        if z === nothing
            return ∂⃖recurse{1}()(f, args...)
        end
        return z
    else
        ∂⃖p = ∂⃖{minus1(N)}()
        @destruct z, z̄ = ∂⃖p(rrule, f, args...)
        if z === nothing
            return ∂⃖recurse{N}()(f, args...)
        else
            return ∂⃖rrule{N}()(z, z̄)
        end
    end
end

@Base.pure function (::∂⃖{1})(::typeof(Core.apply_type), head, args...)
    return rrule(Core.apply_type, head, args...)
end

struct KwFunc{T,S}
    f::T
    kwf::S
    function KwFunc(f)
        kwf = Core.kwfunc(f)
        new{Core.Typeof(f), Core.Typeof(kwf)}(f, kwf)
    end
end
(kw::KwFunc)(args...) = kw.kwf(args...)

function ChainRulesCore.rrule(::typeof(Core.kwfunc), f)
    KwFunc(f), Δ->(NoTangent(), Δ)
end

function ChainRulesCore.rrule(::KwFunc, kwargs, f, args...)
    r = Core.kwfunc(rrule)(kwargs, rrule, f, args...)
    if r === nothing
        return nothing
    end
    x, back = r
    x, Δ->begin
        (NoTangent(), NoTangent(), back(Δ)...)
    end
end

@Base.aggressive_constprop function ChainRulesCore.rrule(::typeof(Core.getfield), s, field::Symbol)
    getfield(s, field), let P = typeof(s)
        @Base.aggressive_constprop Δ->begin
            nt = NamedTuple{(field,)}((Δ,))
            (NoTangent(), Tangent{P, typeof(nt)}(nt), NoTangent())
        end
    end
end

struct ∂⃖getfield{n, f}; end
@Base.aggressive_constprop function (::∂⃖getfield{n, f})(Δ) where {n,f}
    if @generated
        return Expr(:call, tuple, NoTangent(),
            Expr(:call, tuple, (i == f ? :(Δ) : ZeroTangent() for i = 1:n)...),
            NoTangent())
    else
        return (NoTangent(), ntuple(i->i == f ? Δ : ZeroTangent(), n), NoTangent())
    end
end

struct EvenOddEven{O, P, F, G}; f::F; g::G; end
EvenOddEven{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddEven{O, P, F, G}(f, g)
struct EvenOddOdd{O, P, F, G}; f::F; g::G; end
EvenOddOdd{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddOdd{O, P, F, G}(f, g)
@Base.aggressive_constprop (o::EvenOddOdd{O, P, F, G})(Δ) where {O, P, F, G} = (o.f(Δ), EvenOddEven{plus1(O), P, F, G}(o.f, o.g))
@Base.aggressive_constprop (e::EvenOddEven{O, P, F, G})(Δ...) where {O, P, F, G} = (e.g(Δ...), EvenOddOdd{plus1(O), P, F, G}(e.f, e.g))
@Base.aggressive_constprop (o::EvenOddOdd{O, O})(Δ) where {O} = o.f(Δ)

@Base.aggressive_constprop function (::∂⃖{N})(::typeof(Core.getfield), s, field::Int) where {N}
    getfield(s, field), EvenOddOdd{1, c_order(N)}(
        ∂⃖getfield{nfields(s), field}(),
        @Base.aggressive_constprop (_, Δ, _)->getfield(Δ, field))
end

@Base.aggressive_constprop function (::∂⃖{N})(::typeof(Base.getindex), s::Tuple, field::Int) where {N}
    getfield(s, field), EvenOddOdd{1, c_order(N)}(
        ∂⃖getfield{nfields(s), field}(),
        @Base.aggressive_constprop (_, Δ, _)->lifted_getfield(Δ, field))
end

function (::∂⃖{N})(::typeof(Core.getfield), s, field::Symbol) where {N}
    getfield(s, field), let P = typeof(s)
        EvenOddOdd{1, c_order(N)}(
            (@Base.aggressive_constprop Δ->begin
                nt = NamedTuple{(field,)}((Δ,))
                (NoTangent(), Tangent{P, typeof(nt)}(nt), NoTangent())
            end),
            (@Base.aggressive_constprop (_, Δs, _)->begin
                isa(Δs, Union{ZeroTangent, NoTangent}) ? Δs : getfield(ChainRulesCore.backing(Δs), field)
            end))
    end
end

# TODO: Temporary - make better
function (::∂⃖{N})(::typeof(Base.getindex), a::Array, inds...) where {N}
    getindex(a, inds...), let
        EvenOddOdd{1, c_order(N)}(
            (@Base.aggressive_constprop Δ->begin
                Δ isa AbstractZero && return (NoTangent(), Δ, map(Returns(Δ), inds)...)
                BB = zero(a)
                BB[inds...] = Δ
                (NoTangent(), BB, map(x->NoTangent(), inds)...)
            end),
            (@Base.aggressive_constprop (_, Δ, _)->begin
                getindex(Δ, inds...)
            end))
    end
end

struct tuple_back{M}; end
(::tuple_back)(Δ::Tuple) = Core.tuple(NoTangent(), Δ...)
(::tuple_back{N})(Δ::AbstractZero) where {N} = Core.tuple(NoTangent(), ntuple(i->Δ, N)...)
(::tuple_back{N})(Δ::Tangent) where {N} = Core.tuple(NoTangent(), ntuple(i->lifted_getfield(Δ, i), N)...)

function (::∂⃖{N})(::typeof(Core.tuple), args::Vararg{Any, M}) where {N, M}
    Core.tuple(args...),
        EvenOddOdd{1, c_order(N)}(
            tuple_back{M}(),
            (Δ...)->Core.tuple(Δ[2:end]...)
        )
end

struct UnApply{Spec, Types}; end
@generated function (::UnApply{Spec, Types})(Δ) where {Spec, Types}
    args = Any[NoTangent(), NoTangent(), :(Δ[1])]
    start = 2
    for (l, T) in zip(Spec, Types.parameters)
        if T <: Array
            push!(args, :([Δ[$(start:(start+l-1))]...]))
        else
            push!(args, :(Δ[$(start:(start+l-1))]))
        end
        start += l
    end
    :(Core.tuple($(args...)))
end

struct ApplyOdd{O, P}; u; ∂⃖f; end
struct ApplyEven{O, P}; u; ∂⃖∂⃖f; end
@Base.aggressive_constprop function (a::ApplyOdd{O, P})(Δ) where {O, P}
    r, ∂⃖∂⃖f = a.∂⃖f(Δ)
    (a.u(r), ApplyEven{plus1(O), P}(a.u, ∂⃖∂⃖f))
end
@Base.aggressive_constprop function (a::ApplyEven{O, P})(_, _, ff, args...) where {O, P}
    r, ∂⃖∂⃖∂⃖f = Core._apply_iterate(iterate, a.∂⃖∂⃖f, (ff,), args...)
    (r, ApplyOdd{plus1(O), P}(a.u, ∂⃖∂⃖∂⃖f))
end
@Base.aggressive_constprop function (a::ApplyOdd{O, O})(Δ) where {O}
    r = a.∂⃖f(Δ)
    a.u(r)
end

function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Union{Tuple, Vector, NamedTuple}...) where {N}
    @assert iterate === Base.iterate
    x, ∂⃖f = Core._apply_iterate(iterate, this, (f,), args...)
    return x, ApplyOdd{1, c_order(N)}(UnApply{map(length, args), typeof(args)}(), ∂⃖f)
end

@Base.pure c_order(N::Int) = 2^N - 1

@Base.pure function (::∂⃖{N})(::typeof(Core.apply_type), head, args...) where {N}
    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, c_order(N)}()
end

@Base.aggressive_constprop lifted_getfield(x, s) = getfield(x, s)
lifted_getfield(x::ZeroTangent, s) = ZeroTangent()
lifted_getfield(x::NoTangent, s) = NoTangent()

function lifted_getfield(x::Tangent, s)
    z = getfield(ChainRulesCore.backing(ChainRulesCore.canonicalize(x)), s)
    z
end

function lifted_getfield(x::Tangent{<:Tangent{T}}, s) where T
    bb = getfield(x.backing, 1)
    z = lifted_getfield(bb, s)
    z
end

ChainRulesCore.backing(::ZeroTangent) = ZeroTangent()
ChainRulesCore.backing(::NoTangent) = NoTangent()

function reload()
    Core.eval(Diffractor, quote
        function (ff::∂⃖recurse)(args...)
            $(Expr(:meta, :generated_only))
            $(Expr(:meta,
                    :generated,
                    Expr(:new,
                        Core.GeneratedFunctionStub,
                        :perform_optic_transform,
                        Any[:ff, :args],
                        Any[],
                        @__LINE__,
                        QuoteNode(Symbol(@__FILE__)),
                        true)))
        end
    end)
end
reload()


# This file is a part of Julia. License is MIT: https://julialang.org/license

module dSFMT

import Base: copy, copy!, ==

export DSFMT_state, dsfmt_get_min_array_size, dsfmt_get_idstring,
       dsfmt_init_gen_rand, dsfmt_init_by_array, dsfmt_gv_init_by_array,
       dsfmt_fill_array_close_open!, dsfmt_fill_array_close1_open2!,
       win32_SystemFunction036!

"Mersenne Exponent"
const MEXP = 19937
"DSFMT internal state array size of N 128-bit integers."
const N    = floor(Int, ((MEXP - 128) / 104 + 1))
"""Julia DSFMT state representation size counted in 32-bit integers.

Size: (DSFMT state array of Int128 + 1)*4 + Int32 index + Int32 padding
"""
const JN32 = (N+1)*4+1+1
"Jump polynomial for 10^20 steps for dSFMT with exponent 19937"
const JPOLY1e21  = "e172e20c5d2de26b567c0cace9e7c6cc4407bd5ffcd22ca59d37b73d54fdbd937cd3abc6f502e8c186dbd4f1a06b9e2b894f31be77424f94dddfd5a45888a84ca66eeeb242eefe6764ed859dafccae7a6a635b3a63fe9dfbbd5f2d3f2610d39388f53060e84edae75be4f4f2272c0f1f26d1231836ad040ab091550f8a3a5423fb3ab83e068fe2684057f15691c4dc757a3aee4bca8595bf1ad03500d9620a5dbe3b2d64380694895d2f379ca928238293ea267ce14236d5be816a61f018fe4f6bc3c9865f5d4d4186e320ab653d1f3c035ae83e2ad725648a67d3480331e763a1dcdfb5711b56796170b124f5febd723a664a2deefbfa9999d922a108b0e683582ae8d3baacb5bb56683405ea9e6e0d71ddb24b2229c72bb9d07061f2d1fa097ade823b607a2029d6e121ae09d93de01a154199e8e6a6e77c970bda72ba8079b2b3a15dd494a3188b1d94a25ae108a8a5bd0b050e6ce64a365a21420e07fdeebecae02eb68a4304b59283055d22c27d680ea35952834d828c9b9b9dd1a886b4f7fe82fe8f2a962e1e5390e563dc281c799aee2a441b7a813facb6ff5e94c059710dcfe7e6b1635e21ae0dc878dd5f7cc0e1101a74452495a67d23a2672c939f32c81d4a2611073990e92a084cc3a62fd42ee566f29d963a9cc5100ccd0a200f49ce0a74fa891efa1b974d342b7fedf9269e40d9b34e3c59c3d37201aecd5a04f4ae3d0c9a68c7ab78c662390e4cf36cb63ea3539c442efd0bf4aace4b8c8bde93c3d84b4d6290adfae1c5e3fcd457b6f3159e501f17b72ff6bc13d6bf61fbdafabefd16ac1dae0bca667e4e16a2b800732f1d0a9274c8a4c6cccd2db62fc275dc308c31c11cd6fda78de2f81f0e542b76b42b2cc09ed8f965d94c714c9918064f53af5379cfbbc31edf9cbce694f63a75f122048de6e57b094908f749661456813a908027f5d8397ab7962bf75ac779a3e1b7ae3fbc93397a67b486bb849befff1de6162ef2819715a88f41881e366ace692a900796a2806393898dd1750ac2b4ca3d34ca48942322fb6375f0c9a00c9701048ee8d7d7a17e11739177a7ad5027556e85835daf8594d84a97fe6621c0fce1495ae6ab8676cdc992d247acf5a4e5ec8c4755fde28117228d2c3ecf89edb91e93d949e2174924572265e36d176d082ed1be884e51d885ba3cda175c51edcee5042eaf519d292aa05aa4185b03858d710a9d0880b3d4e5111f858a52fe352cbe0a24f06a3d977ae2eb85e2a03a68131d0ab91dac4941067cf90ecd0fce156bcd40b8968cd4aa11e0b4353b14508d79d13ac00af4a4d452496b7f2393699889aa1e508427dbf0be3db91d955feb51e559af57640c6b3f9d5f95609852c28f9462a9869dd93acbdb1aafb2381ebb886a0b3fcec278f8bb0f62c23e157e49b89245b0881268ce594acbddd3605b9eaa77c9ff513e0dbad514914136d96fe2843fe2b4e886a0b718a9b8d1132132110618d0d3595da284cd2a4c9d09386199e4f4d7723983d3a374b51cf20dac5cabb4ff7e7197c2ebd9318463409baa583d6a6115c1b768282ff37b0fe152c97671e400d5ccba7d6875df0bf95c5d91257fedb124de393f31908d0e36251326aa29dd5be86291c80b4bf78f419ec151eeaeff643a58b48ab35ad2cd2c0b77b1965966ef3db6b6373cb2c4b590cef2f16f4d6f62f13a6cbf1a481565b5935edd4e76f7b6a8fd0d74bc336b40a803aec38125c006c877dfdcdb9ba2b7aecab5cafe6076e024c73e3567adf97f607a71d180402c22a20a8388f517484cc4198f97c2fe4f3407e0dc577e61f0f71354aa601cf4e3e42e1edd8722d50f5af3441f68caa568cc1c3a19956c1233f265bb47236afab24ee42b27b0042b90693d77c1923147360ae6503f6ba6abbc9dd52a7b4c36a3b6b55f6a80cfa7f101dd9f1bfc7d7eaf09a5d636b510228f245bfb37b4625025d2c911435cdf6f878113753e0804ab8ecab870ad733b9728d7636b17578b41239393e7de47cbce871137d2b61729dda67b2b84cd3363aad64c5dd5bd172f1f091305b1ff78982abe7dab1588036d097cf497e300e6c78a926048febd1b9462c07f5868928357b74297c87f503056b89f786d22a538b6702e290bca04639a0f1d0939b67f409e5e58e472a6a07fa543e2531c2567ec73c41f6769b6ba94c5aa0a030d006f5b6b1c5fb218b86a8f63a48bc867466f20f699859e87956f48a182d26ed451861dd21201ecc7239037ada67319bdf0849c387c73a110af798b4c5f9018bc97993e060ea2a2937fa2eb095d65ec07009fc407a350f1d6fb3c98a0a5f204be985b0cb6962f0eb7844a179c4598a92ea32d2d706c800034d2e960ded5b476d77073316b933fb3e6ba2f4f24a3b73a1e4d8ed1491d757ecf56fd72465dac0000736744d28d29073091587c8bccad302f7054e8a32bb8724974d9f3e449fc70b2a41f0008b548f717ac0a2c3a6580bfb50774933a578ad6acdcb89940bb406ea540893f097d8a88d1609ed605f25499de939083a0c8a7c6db462df5dfa06c298dd233e249433a54267d5cdc22e5524705b7d6b16b96bb2cb83e00cef62e21d91528a74cf95bfd1d391590c93f4058e9bb02656fd087a5b63d738d1c3b5cf533fd59c81cf9136bfcd3e955c19daf9906ef175791fde6a1d98155d7881e241c3522551cf9fcae42e1e46929ea39fd00943446823f9755085ccc8456a3090b73a3031a201d9c704a4ad4868dd9b6d06205560013973f60d637de2f18354bf4523d9d81dc2a7e78cd42c586364bbe0ed86fde0f081f801c1a4abb830839b7796d9a01f141bec8bd93144104c6dc59170162c0a5a639eb63a0a164970de50eb2e04f027394b26ed48d341f7851994df79d7cd663672a556f25e5e16a3adbe1003d631de938fabfed234df12b5ff3027f4a2da823834cb098e0f977a4eb9614579d5c7a1d400a1a933a657aef8ea1a66743d73b0cf37a7d64e9a63e4c7b09945f0db750b311b39783fb5ea216616751967d480a630d3da7c89d1c7beae20369137e96734a4cfedca56a7887f076fe4fe97534ad3e4f74d1a81750581a5ea214b440c7f30331ab86c257534c71175d1e731303a48b01c589fda4fb0d4368b4dd63d91204cb6fc389b2202aa94391907bfb72902a4031f5589ed5f391c2ce92aa998c200ba3c77d8bd747b9d0a29fa85cda3949a6d2bd0c3402e68f98fd451aa27b6c2dfd170e004577cbdb25e3a1b9852e9f66a370789c47bfce722446dade1b32ceae71ee0e1d96edf7ed08a93e3690056f46c3d8e63f88e53673ee71d72cfedbeba493ee91333120e09e9ce9f9c9a7a400f814ea618b1de48f9805e092f4e20f301fbb65caa83735a2a5c89befe4bce4116dca3688e1e14c6f09a945671dedbb5c0ba526842b6cae31d8b5ff978bae928a17a75c134630dd9de988f6ad3d89a071b33775a9660a40b48ec61ad3f93ac81cb1c65d8b0bab5c214786abd13cc10a8ea2e2a370e86e2fa1a372d83c9697b5e37b281e51507685f714fdaebe49ffc93a5582e1936eaee8e4140a4b72"

mutable struct DSFMT_state
    val::Vector{Int32}

    DSFMT_state(val::Vector{Int32} = zeros(Int32, JN32)) =
        new(length(val) == JN32 ? val : throw(DomainError()))
end

copy!(dst::DSFMT_state, src::DSFMT_state) = (copy!(dst.val, src.val); dst)
copy(src::DSFMT_state) = DSFMT_state(copy(src.val))

==(s1::DSFMT_state, s2::DSFMT_state) = s1.val == s2.val

function dsfmt_get_idstring()
    idstring = ccall((:dsfmt_get_idstring,:libdSFMT),
                     Ptr{UInt8},
                     ())
    return unsafe_string(idstring)
end

function dsfmt_get_min_array_size()
    min_array_size = ccall((:dsfmt_get_min_array_size,:libdSFMT),
                           Int32,
                           ())
end

const dsfmt_min_array_size = dsfmt_get_min_array_size()

function dsfmt_init_gen_rand(s::DSFMT_state, seed::UInt32)
    ccall((:dsfmt_init_gen_rand,:libdSFMT),
          Void,
          (Ptr{Void}, UInt32,),
          s.val, seed)
end

function dsfmt_init_by_array(s::DSFMT_state, seed::Vector{UInt32})
    ccall((:dsfmt_init_by_array,:libdSFMT),
          Void,
          (Ptr{Void}, Ptr{UInt32}, Int32),
          s.val, seed, length(seed))
end

function dsfmt_gv_init_by_array(seed::Vector{UInt32})
    ccall((:dsfmt_gv_init_by_array,:libdSFMT),
          Void,
          (Ptr{UInt32}, Int32),
          seed, length(seed))
end

function dsfmt_fill_array_close1_open2!(s::DSFMT_state, A::Ptr{Float64}, n::Int)
    @assert Csize_t(A) % 16 == 0 # the underlying C array must be 16-byte aligned
    @assert dsfmt_min_array_size <= n && iseven(n)
    ccall((:dsfmt_fill_array_close1_open2,:libdSFMT),
          Void,
          (Ptr{Void}, Ptr{Float64}, Int),
          s.val, A, n)
end

function dsfmt_fill_array_close_open!(s::DSFMT_state, A::Ptr{Float64}, n::Int)
    @assert Csize_t(A) % 16 == 0 # the underlying C array must be 16-byte aligned
    @assert dsfmt_min_array_size <= n && iseven(n)
    ccall((:dsfmt_fill_array_close_open,:libdSFMT),
          Void,
          (Ptr{Void}, Ptr{Float64}, Int),
          s.val, A, n)
end

# dSFMT jump
function dsfmt_jump(s::DSFMT_state, jp::AbstractString)
    val = s.val
    nval = length(val)
    index = val[nval - 1]
    work = zeros(UInt64, JN32 >> 1)
    dsfmt = Vector{UInt64}(nval >> 1)
    ccall(:memcpy, Ptr{Void}, (Ptr{UInt64}, Ptr{Int32}, Csize_t),
          dsfmt, val, (nval - 1) * sizeof(Int32))
    dsfmt[end] = UInt64(N*2)

    for c in jp
        bits = parse(UInt8,c,16)
        for j in 1:4
            (bits & 0x01) != 0x00 && dsfmt_jump_add!(work, dsfmt)
            bits = bits >> 0x01
            dsfmt_jump_next_state!(dsfmt)
        end
    end

    work[end] = index
    return DSFMT_state(reinterpret(Int32, work))
end

function dsfmt_jump_add!(dest::Vector{UInt64}, src::Vector{UInt64})
    dp = dest[end] >> 1
    sp = src[end] >> 1
    diff = ((sp - dp + N) % N)
    i = 1
    while i <= N-diff
        j = i*2-1
        p = j + diff*2
        dest[j]   ⊻= src[p]
        dest[j+1] ⊻= src[p+1]
        i += 1
    end
    while i <= N
        j = i*2-1
        p = j + (diff - N)*2
        dest[j]   ⊻= src[p]
        dest[j+1] ⊻= src[p+1]
        i += 1
    end
    dest[N*2+1] ⊻= src[N*2+1]
    dest[N*2+2] ⊻= src[N*2+2]
    return dest
end

function dsfmt_jump_next_state!(mts::Vector{UInt64})
    POS1 = 117
    SL1  = 19
    SR   = 12
    MSK1 = 0x000ffafffffffb3f
    MSK2 = 0x000ffdfffc90fffd

    idx = (mts[end] >> 1) % N

    a = idx*2+1
    b = ((idx + POS1) % N)*2+1
    u = N*2+1

    t0 = mts[a]
    t1 = mts[a+1]
    L0 = mts[u]
    L1 = mts[u+1]
    mts[u]   = xor(t0 << SL1, L1 >> 32, L1 << 32, mts[b])
    mts[u+1] = xor(t1 << SL1, L0 >> 32, L0 << 32, mts[b+1])
    mts[a]   = xor(mts[u]   >> SR, mts[u]   & MSK1, t0)
    mts[a+1] = xor(mts[u+1] >> SR, mts[u+1] & MSK2, t1)

    mts[end] = (mts[end] + 2) % (N*2)
    return mts
end

## Windows entropy

if is_windows()
    function win32_SystemFunction036!(a::Array)
        ccall((:SystemFunction036, :Advapi32), stdcall, UInt8, (Ptr{Void}, UInt32), a, sizeof(a))
    end
end

end # module


const M = 10
function generate_f(N, d, k; X=rand(N, d))
    f = rand(MvNormal(zeros(N), kernelmatrix(k, X; obsdim=1) + 1e-5I))
    return X, f
end
const VIs = ["AnalyticVI", "QuadratureVI", "MCIntegrationVI"]
const VIcodes = ["AVI", "QVI", "MCVI"]

function tests(model1, model2, X, f, y, problem)
    model1, state1 = train!(model1, X, y, 1)
    L = AGP.objective(model1, X, y)
    train!(model1, X, y, 5)
    @test testconv(model1, problem, X, f, y)
    @test all(proba_y(model1, X)[2] .> 0)
    train!(model2, X, y, 6)
    @test testconv(model2, problem, X, f, y)
    @test all(proba_y(model2, X)[2] .> 0)
end

function tests(model1::OnlineSVGP, model2, X, f, y, problem)
    state1 = nothing
    for (X_, y_) in eachbatch((X, y); obsdim=1, size=10)
        model1, state1 = train!(model1, X_, y_, state1; iterations=5)
        L = AGP.ELBO(model1, state1, X_, y_)
    end
    @test testconv(model1, problem, X, f, y)
    @test all(proba_y(model1, X)[2] .> 0)
    state2 = nothing
    for (X_, y_) in eachbatch((X, y); obsdim=1, size=10)
        model2, state2 = train!(model2, X_, y_, state2; iterations=5)
    end
    @test testconv(model2, problem, X, f, y)
    @test all(proba_y(model2, X)[2] .> 0)
end

function tests(
    model1::AbstractGPModel{T,<:BernoulliLikelihood{<:AGP.SVMLink}},
    model2,
    X,
    f,
    y,
    problem,
) where {T}
    model1, state1 = train!(model1, X, y, 1)
    @test_nowarn AGP.objective(model1, X, y)
    train!(model1, X, y, 5)
    @test testconv(model1, problem, X, f, y)
    @test all(proba_y(model1, X)[2] .> 0)
    model2, state2 = train!(model2, X, y, 6)
    @test testconv(model2, problem, X, f, y)
    @test all(proba_y(model2, X)[2] .> 0)
end

function tests_likelihood(
    l, # likelihood
    ltype, # likelihood types
    dict::Dict,
    floattypes,
    problem,
    n_latent,
    X,
    f,
    y,
    k,
)
    Z = inducingpoints(KmeansAlg(M), X)
    @testset "VGP" begin
        for floattype in floattypes
            dictvgp = dict["VGP"]
            for (name, code, inference) in
                zip(VIs, VIcodes, [AnalyticVI(), QuadratureVI(), MCIntegrationVI()])
                @testset "$name" begin
                    test_inference_VGP(
                        X,
                        y,
                        f,
                        k,
                        l,
                        ltype,
                        floattype,
                        n_latent,
                        problem,
                        inference;
                        valid=dictvgp[code],
                    )
                end
            end
        end # Loop on float types
    end # VGP
    @testset "OSVGP" begin
        dictosvgp = dict["OSVGP"]
        for floattype in floattypes
            @testset "AnalyticVI" begin
                if dictosvgp["AVI"]
                    model = OnlineSVGP(
                        k,
                        l,
                        AnalyticVI(),
                        AGP.InducingPoints.OIPS();
                        optimiser=false,
                        verbose=0,
                    )
                    @test eltype(model) == floattype
                    @test AGP.likelihood(model) isa ltype
                    @test AGP.inference(model) isa AnalyticVI
                    @test AGP.getf(model) isa NTuple{n_latent,AGP.OnlineVarLatent}
                    @test AGP.n_latent(model) == n_latent
                    model_opt = OnlineSVGP(
                        k,
                        l,
                        AnalyticVI(),
                        AGP.InducingPoints.OIPS();
                        optimiser=true,
                        verbose=0,
                    )
                    tests(model, model_opt, X, f, y, problem)
                else
                    @test_throws ErrorException OnlineSVGP(
                        k, l, AnalyticVI(), UniGrid(10); optimiser=false, verbose=0
                    )
                end
            end  # Analytic VI
            @testset "NumericalVI" begin
                if dictosvgp["QVI"]
                    model = OnlineSVGP(k, l, QuadratureVI(), UniGrid(10); optimiser=false)
                    @test eltype(model) == floattype
                    @test AGP.likelihood(model) isa ltype
                    @test AGP.inference(model) isa QuadratureVI
                    @test AGP.getf(model) isa NTuple{n_latent,AGP.OnlineVarLatent}
                    @test AGP.n_latent(model) == n_latent
                    model_opt = OnlineSVGP(
                        k, l, QuadratureVI(), UniGrid(10); optimiser=true, verbose=0
                    )
                    tests(model, model_opt, X, f, y, problem)
                else
                    @test_throws ErrorException OnlineSVGP(
                        k, l, QuadratureVI(), UniGrid(10); optimiser=false, verbose=0
                    )
                end
                if dictosvgp["MCVI"]
                    model = OnlineSVGP(
                        k, l, MCIntegrationVI(), UniGrid(10); optimiser=false
                    )
                    @test eltype(model) == floattype
                    @test AGP.likelihood(model) isa ltype
                    @test AGP.inference(model) isa MCIntegrationVI
                    @test AGP.getf(model) isa NTuple{n_latent,AGP.SparseVarLatent}
                    @test AGP.n_latent(model) == n_latent
                    model_opt = OnlineSVGP(
                        k, l, MCIntegrationVI(), UniGrid(10); optimiser=true, verbose=0
                    )
                    tests(model, model_opt, X, f, y, problem)
                else
                    @test_throws ErrorException OnlineSVGP(
                        k, l, MCIntegrationVI(), UniGrid(10); optimiser=false, verbose=0
                    )
                end
            end # Loop on Numerical VI
        end # Loop on float types
    end #Loop on Online SVGP
    @testset "SVGP" begin
        dictsvgp = dict["SVGP"]
        for floattype in floattypes
            for (name, code, inference, stoch_inference) in zip(
                VIs,
                VIcodes,
                [AnalyticVI(), QuadratureVI(), MCIntegrationVI()],
                [AnalyticSVI(10), QuadratureSVI(10), MCIntegrationSVI(10)],
            )
                @testset "$name" begin
                    test_inference_SVGP(
                        X,
                        y,
                        f,
                        k,
                        l,
                        ltype,
                        floattype,
                        Z,
                        n_latent,
                        problem,
                        inference,
                        stoch_inference;
                        valid=dictsvgp[code],
                    )
                end #
            end
        end # Loop on float types
    end # SVGP
    @testset "MCGP" begin
        dictmcgp = dict["MCGP"]
        for floattype in floattypes
            @testset "Gibbs Sampling" begin
                if dictmcgp["Gibbs"]
                    model = MCGP(X, y, k, l, GibbsSampling())
                    @test eltype(model) == floattype
                    @test AGP.likelihood(model) isa ltype
                    @test AGP.inference(model) isa GibbsSampling
                    @test AGP.getf(model) isa NTuple{n_latent,AGP.SampledLatent}
                    @test AGP.n_latent(model) == n_latent
                    samples = sample(model, 100; progress=false)
                    @test_broken samples2 = sample(model, 100; cat=true, progress=false)
                else
                    @test_throws ErrorException MCGP(X, y, k, l, GibbsSampling())
                end
                if dictmcgp["HMC"]
                    model = MCGP(X, y, k, l, HMCSampling())
                    @test eltype(model) == floattype
                    @test AGP.likelihood(model) isa ltype
                    @test AGP.inference(model) isa HMCSampling
                    @test AGP.getf(model) isa NTuple{n_latent,AGP.SampledLatent}
                    @test AGP.n_latent(model) == n_latent
                    samples = sample(model, 20; progress=false)
                    @test_broken samples2 = sample(model, 20; cat=true)
                else
                    @test_throws ErrorException MCGP(X, y, k, l, HMCSampling())
                end
            end # Gibbs Sampling
        end # Loop floattypes
    end # MCGP
end

function testconv(
    model::AbstractGPModel,
    problem_type::String,
    X::AbstractArray,
    f::AbstractArray,
    y::AbstractArray,
)
    μ, Σ = predict_f(model, X; cov=true, diag=false)
    μ, diagΣ = predict_f(model, X; cov=true, diag=true)
    y_pred = predict_y(model, X)
    py_pred = proba_y(model, X)
    if problem_type == "Regression"
        err = mean(abs.(y_pred - f))
        # @info "Regression Error" err
        return err < 1.5
    elseif problem_type == "Classification"
        err = mean(y_pred .!= y)
        # @info "Classification Error" err
        return err < 0.5
    elseif problem_type == "MultiClass"
        err = mean(y_pred .!= y)
        # @info "Multiclass Error" err
        return err < 0.9
    elseif problem_type == "Poisson" || problem_type == "NegBinomial"
        err = mean(abs.(y_pred - y))
        # @info "Event Error" err
        return err < 20.0
    else
        throw(ErrorException("Wrong problem description $problem_type"))
    end
end

function test_inference_VGP(
    X, y, f, k, l, ltype, floattype, n_latent, problem, inference; valid=false
)
    if valid
        model = VGP(X, y, k, l, inference; optimiser=false, verbose=0)
        @test eltype(model) == floattype
        @test AGP.likelihood(model) isa ltype
        @test AGP.inference(model) isa typeof(inference)
        @test AGP.getf(model) isa NTuple{n_latent,AGP.VarLatent}
        @test AGP.n_latent(model) == n_latent
        model_opt = VGP(X, y, k, l, inference; optimiser=true, verbose=0)
        tests(model, model_opt, X, f, y, problem)
    else
        @test_throws ErrorException VGP(X, y, k, l, inference; optimiser=false, verbose=0)
    end
end

function test_inference_SVGP(
    X,
    y,
    f,
    k,
    l,
    ltype,
    floattype,
    Z,
    n_latent,
    problem,
    inference,
    stoch_inference;
    valid=false,
)
    if valid
        model = SVGP(k, l, inference, Z; optimiser=false, verbose=0)
        @test eltype(model) == floattype
        @test AGP.likelihood(model) isa ltype
        @test AGP.inference(model) isa typeof(inference)
        @test AGP.n_latent(model) == n_latent
        model_opt = SVGP(k, l, inference, Z; optimiser=true, Zoptimiser=true, verbose=0)
        tests(model, model_opt, X, f, y, problem)
        @test AGP.getf(model) isa NTuple{n_latent,AGP.SparseVarLatent}

        model_svi = SVGP(k, l, stoch_inference, Z; optimiser=false, verbose=0)
        tests(model_svi, model, X, f, y, problem)
    else
        @test_throws ErrorException SVGP(k, l, inference, Z; optimiser=false, verbose=0)
    end
end


@testset "Hybrid DCPLossLess with BasicHybridDispatch formulation" begin
    device_model = DeviceModel(HybridSystem, BasicHybridDispatch)
    sys = PSB.build_system(PSITestSystems, "c_sys5_hybrid")

    # Parameters Testing
    model = DecisionModel(MockOperationProblem, DCPPowerModel, sys)
    mock_construct_device!(model, device_model)
    moi_tests(model, false, 672, 0, 384, 336, 192, true)
    psi_checkobjfun_test(model, GAEVF)
end

@testset "Hybrid ACPPowerModel with BasicHybridDispatch formulation" begin
    device_model = DeviceModel(HybridSystem, BasicHybridDispatch)
    sys = PSB.build_system(PSITestSystems, "c_sys5_hybrid")

    # No Parameters Testing
    model = DecisionModel(MockOperationProblem, ACPPowerModel, sys)
    mock_construct_device!(model, device_model)
    moi_tests(model, false, 864, 0, 672, 624, 288, true)
    psi_checkobjfun_test(model, GAEVF)
end

@testset "Hybrid DCPLossLess BasicHybridDispatch" begin
    device_model = DeviceModel(HybridSystem, StandardHybridDispatch)
    sys = PSB.build_system(PSITestSystems, "c_sys5_hybrid")

    # Parameters Testing
    model = DecisionModel(MockOperationProblem, DCPPowerModel, sys)
    mock_construct_device!(model, device_model)
    moi_tests(model, false, 672, 0, 384, 336, 192, true)
    psi_checkobjfun_test(model, GAEVF)
end

@testset "Hybrid ACPPowerModel with StandardHybridDispatch formulation" begin
    device_model = DeviceModel(HybridSystem, StandardHybridDispatch)
    sys = PSB.build_system(PSITestSystems, "c_sys5_hybrid")

    # No Parameters Testing
    model = DecisionModel(MockOperationProblem, ACPPowerModel, sys)
    mock_construct_device!(model, device_model)
    moi_tests(model, false, 864, 0, 672, 624, 288, true)
    psi_checkobjfun_test(model, GAEVF)
end


include(joinpath(@__DIR__, "fig1.jl"))
using FoodCachingExperiments, Random, FoodCachingModels, Statistics, DataFrames
# This is a hack, because the labeling in Clayton99C_exp1 is wrong
FoodCachingExperiments.EXPERIMENTS[:Clayton99C_exp1].tests["interactions"].labels[2] = "missing"
import FoodCachingExperiments: EXPERIMENTS, bsave, bload

function plot_sigresults(sigresults;
                         measure = :ratio_same_rejects_mean)
    tmp = combine(groupby(sigresults, [:model, :experiment]),
                  [:samebinarysignificance, :meansamebinarysigA,
                   :meansamebinarysig, :atleasttwo, :ratio_same_rejects] .=> mean,
                 )
    sort!(tmp, :experiment, by = x -> findfirst(==(x), experiment_order))
    sort!(tmp, :model, by = x -> findfirst(==(x), model_order))
    groups = groupby(tmp, :model)
    ys = [getproperty(g, measure)  for g in groups]
    ns = [sum(rejects(EXPERIMENTS[e])) for e in experiment_order
          if e ∉ (:Clayton0103, :Watanabe05)]
    pe = 1:11
    me = 12:22
    se = 23:27
    ce = 28
    reverse(x) = 32 - x
    xtick = reverse.([pe; me .+ 1; se .+ 2; ce .+ 3])
    f1 = @pgf Axis({xtick = xtick, xticklabels = toname.(groups[1].experiment),
               xticklabel_style = {rotate = 60, anchor = "east"},
               width = "16cm", height = "6cm",
               cycle_list = [
                             "{thick, morange, mark = square}",
                             "{thick, mred, mark = square}",
                             "{thick, mblue, mark = square}",
                             "{thick, mgreen, mark = square}",
                             "{thick, gray!60!black, mark = square}"
                            ],
              },
              [Plot(Coordinates(xtick, y)) for y in ys]...)

    f2 = @pgf Axis({cycle_list = [
                                 "{thick, mark = square*, dotted, mark size = 3,
                                 mark options = {solid, fill = white}}",
                                 "{thick, dashed, mark = diamond*,
                                   mark options = {solid, fill =  white}, mark size = 3}",
                                 "{thick, mark = triangle*, mark size = 3,
                                   mark options = {fill = white}}",
                                ],
                   ylabel = "avg. reproducibility [\\%]",
               xticklabel_style = {rotate = 60, anchor = "east"},
               height = "4.8cm", width = "3.8cm",
               ymax = 105,
               ytick = 0:20:100,
               legend_entries = {"satiety experiments", "memory experiments",
                                 "planning experiments"},
               "every axis legend/.append style" = {at = "{(-.4, 1)}",
                                                    anchor = "south west"},
               xtick = 1:5, font = raw"\scriptsize",
               xticklabels = ["\\color{morange}Planning-By-Replay",
                              "\\color{mred}Plastic Caching",
                              "\\color{mblue}No-Plasticity",
                              "\\color{mgreen}No-Plasticity-No-Memory",
                              "\\parbox{3.4cm}{\\color{gray!60!black}\\raggedleft No-Plasticity-No-Memory-\\\\No-Motivational-Control}"],
              },
                  [Plot(Coordinates(1:5, [100*sum(y[x] .* 1/length(x)) for y in ys]))
               for x in Iterators.reverse([pe, me, se])]...)
    (f1, f2)
end

rm_test_locals.(values(EXPERIMENTS))

sigresults1 = run_sigtests(best, Ninner = 1, Nouter = 50)
bsave(joinpath(datapath, "sigresults1"), Dict(:sigresults1 => sigresults1))
sigresults10 = run_sigtests(best, Ninner = 10, Nouter = 50)
bsave(joinpath(datapath, "sigresults10"), Dict(:sigresults10 => sigresults10))

p1indi, p1all = plot_sigresults(sigresults1)
p10indi, p10all = plot_sigresults(sigresults10)

pgfsave(joinpath(figpath, "sig_summary_all1.tikz"), p1all)
pgfsave(joinpath(figpath, "sig_summary_all10.tikz"), p10all)


"""
ULMFiT - Text Classifier
"""

cd(@__DIR__)
include("custom_layers.jl")     # including custome layers created for ULMFiT model
include("utils.jl")     # including utilities
include("fine_tune_lm.jl")        # including useful functions from fine-tuning file
# include("data_loaders.jl")      # including helper functions for IMDB data loading

mutable struct TextClassifier
    vocab::Vector
    rnn_layers::Flux.Chain
    linear_layers::Flux.Chain
end

function TextClassifier(lm::LanguageModel=load_model!(), clsfr_out_sz::Integer=1, clsfr_hidden_sz::Integer=50, clsfr_hidden_drop::Float64=0.0)
    return TextClassifier(
        lm.vocab,
        lm.layers[1:8],
        Chain(
            gpu(PooledDense(length(lm.lstm_layer3.layer.cell.h), clsfr_hidden_sz, relu)),
            gpu(BatchNorm(clsfr_hidden_sz, relu)),
            Dropout(clsfr_hidden_drop),
            gpu(Dense(clsfr_hidden_sz, clsfr_out_sz, sigmoid)),
        )
    )
end

Flux.@treelike TextClassifier

# Forward step for classifier
function forward(classifier::TextClassifier, gen; now_per_pass::Integer=32)
    X = take!(gen)
    l = length(X)
    maxpools, meanpools = [], []
    # Truncated Backprop through time
    for i=1:floor(l/now_per_pass)   # Tracking is swiched off inside this loop
        (i == 1 && l%now_per_pass != 0) ? (last_idx = l%now_per_pass) : (last_idx = now_per_pass)
        H = broadcast(x -> gpu(classifier.rnn_layers[1](indices(x, classifier.vocab, "_unk_"))), X[1:last_idx])
        H = broadcast(x -> Tracker.data(classifier.rnn_layers[2:end](x)), H)
        H = cat(H..., dims=3)
        push!(maxpools, maximum(H, dims=3)[:, :, 1])
        push!(meanpools, (sum(H, dims=3)/size(H, 3))[:, :, 1])
        X = X[last_idx+1:end]
        Flux.truncate!(classifier.rnn_layers)
        GC.gc()
    end
    # last part of the sequecnes in X - Tracking is swiched on
    H = broadcast(x -> classifier.rnn_layers[1](indices(x, classifier.vocab, "_unk_")), X)
    H = classifier.rnn_layers[2:end].(H)
    H = H[end:-1:1]
    append!(H, maxpools)
    append!(H, meanpools)
    H = classifier.linear_layers(H[end:-1:1])
    return H
end

# Loss function for classifier
function loss(classifier::TextClassifier, gen::Channel; now_per_pass::Integer=32)
    H = forward(classifier, gen, now_per_pass=now_per_pass)
    Y = gpu(take!(gen))
    l = crossentropy(H, Y)
    Flux.reset!(classifier.rnn_layers)
    return l
end

# Training of classifier
function train_classifier!(classifier::TextClassifier=TextClassifier(), classes::Integer=1, hidden_layer_size::Integer=50, data_loader::Channel=imdb_classifier_data;
    batchsize::Integer=64, bptt::Integer=70, gradient_clip::Float64=0.25, stlr_cut_frac::Float64=0.1, stlr_ratio::Number=32, stlr_η_max::Float64=0.01, epochs::Integer=1, checkpoint_itvl=5000)
    trainable = []
    append!(trainable, [classifier.rnn_layers[[1, 3, 5, 7]]...])
    append!(trainable, [classifier.linear_layers[[1, 4]]...])
    opts = [ADAM(0.001, (0.7, 0.99)) for i=1:length(trainable)]
    cut = num_of_iters * epochs * stlr_cut_frac
    gpu!.(classifier.rnn_layers)

    for epoch=1:epochs
        gen = data_loader()
        num_of_iters = take!(gen)
        T = num_of_iters-Int(floor((num_of_iters*2)/100))
        set_trigger!.(T, classifier.rnn_layers)
        for iter=1:num_of_iters
            l = loss(classifier, gen, now_per_pass = now_per_pass)

            # Slanted triangular learning rates
            t = iter + (epoch-1)*num_of_iters
            p_frac = (iter < cut) ? iter/cut : (1 - ((iter-cut)/(cut*(1/stlr_cut_frac-1))))
            ηL = stlr_η_max*((1+p_frac*(stlr_ratio-1))/stlr_ratio)

            # Gradual-unfreezing Step with discriminative fine-tuning
            unfreezed_layers, cur_opts = (epoch < length(trainable)) ? (trainable[end-epoch+1:end], opts[end-epoch+1:end]) : (trainable, opts)
            discriminative_step!(unfreezed_layers, ηL, l, gradient_clip, cur_opts)

            # ASGD Step
            asgd_step!.(iter, classifier.rnn_layers)

            reset_masks!.(classifier.rnn_layers)

            println("loss:", l, "epoch: ", epoch, "iteration: ", iter)
        end
    end
end


using Base64
using Sockets

function main()
    str_size = 131072
    tries = 8192
    str = repeat("a", str_size)
    str2 = base64encode(str)
    str3 = String(base64decode(str2))

    notify("Julia\t$(getpid())")

    t = time()
    s_encoded = 0
    for i = 1:tries
        s_encoded += length(base64encode(str))
    end
    t_encoded = time() - t

    t = time()
    s_decoded = 0
    for i = 1:tries
        s_decoded += length(base64decode(str2))
    end
    t_decoded = time() - t

    notify("stop")

    print("encode $(str[1:4])... to $(str2[1:4]): $s_encoded, $t_encoded\n")
    print("decode $(str2[1:4])... to $(str3[1:4]): $s_decoded, $t_decoded\n")
end

function notify(msg)
    try
        socket = connect("localhost", 9001)
        write(socket, msg)
        close(socket)
    catch
        # standalone usage
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    for (src, dst) in [["hello", "aGVsbG8="], ["world", "d29ybGQ="]]
        encoded = base64encode(src)
        if encoded != dst
            println(stderr, "$(encoded) != $(dst)")
            exit(1)
        end
        decoded = String(base64decode(dst))
        if decoded != src
            println(stderr, "$(decoded) != $(src)")
            exit(1)
        end
    end

    main()
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

"""
    Random

Support for generating random numbers. Provides [`rand`](@ref), [`randn`](@ref),
[`AbstractRNG`](@ref), [`MersenneTwister`](@ref), and [`RandomDevice`](@ref).
"""
module Random

include("DSFMT.jl")

using .DSFMT
using Base.GMP.MPZ
using Base.GMP: Limb
import SHA

using Base: BitInteger, BitInteger_types, BitUnsigned, require_one_based_indexing

import Base: copymutable, copy, copy!, ==, hash, convert,
             rand, randn, show

export rand!, randn!,
       randexp, randexp!,
       bitrand,
       randstring,
       randsubseq, randsubseq!,
       shuffle, shuffle!,
       randperm, randperm!,
       randcycle, randcycle!,
       AbstractRNG, MersenneTwister, RandomDevice, TaskLocalRNG, Xoshiro

## general definitions

"""
    AbstractRNG

Supertype for random number generators such as [`MersenneTwister`](@ref) and [`RandomDevice`](@ref).
"""
abstract type AbstractRNG end

Base.broadcastable(x::AbstractRNG) = Ref(x)

gentype(::Type{X}) where {X} = eltype(X)
gentype(x) = gentype(typeof(x))

### integers

# we define types which encode the generation of a specific number of bits
# the "raw" version means that the unused bits are not zeroed

abstract type UniformBits{T<:BitInteger} end

struct UInt10{T}    <: UniformBits{T} end
struct UInt10Raw{T} <: UniformBits{T} end

struct UInt23{T}    <: UniformBits{T} end
struct UInt23Raw{T} <: UniformBits{T} end

struct UInt52{T}    <: UniformBits{T} end
struct UInt52Raw{T} <: UniformBits{T} end

struct UInt104{T}    <: UniformBits{T} end
struct UInt104Raw{T} <: UniformBits{T} end

struct UInt2x52{T}    <: UniformBits{T} end
struct UInt2x52Raw{T} <: UniformBits{T} end

uint_sup(::Type{<:Union{UInt10,UInt10Raw}}) = UInt16
uint_sup(::Type{<:Union{UInt23,UInt23Raw}}) = UInt32
uint_sup(::Type{<:Union{UInt52,UInt52Raw}}) = UInt64
uint_sup(::Type{<:Union{UInt104,UInt104Raw}}) = UInt128
uint_sup(::Type{<:Union{UInt2x52,UInt2x52Raw}}) = UInt128

for UI = (:UInt10, :UInt10Raw, :UInt23, :UInt23Raw, :UInt52, :UInt52Raw,
          :UInt104, :UInt104Raw, :UInt2x52, :UInt2x52Raw)
    @eval begin
        $UI(::Type{T}=uint_sup($UI)) where {T} = $UI{T}()
        # useful for defining rand generically:
        uint_default(::$UI) = $UI{uint_sup($UI)}()
    end
end

gentype(::Type{<:UniformBits{T}}) where {T} = T

### floats

abstract type FloatInterval{T<:AbstractFloat} end

struct CloseOpen01{T<:AbstractFloat} <: FloatInterval{T} end # interval [0,1)
struct CloseOpen12{T<:AbstractFloat} <: FloatInterval{T} end # interval [1,2)

const FloatInterval_64 = FloatInterval{Float64}
const CloseOpen01_64   = CloseOpen01{Float64}
const CloseOpen12_64   = CloseOpen12{Float64}

CloseOpen01(::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen01{T}()
CloseOpen12(::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen12{T}()

gentype(::Type{<:FloatInterval{T}}) where {T<:AbstractFloat} = T

const BitFloatType = Union{Type{Float16},Type{Float32},Type{Float64}}

### Sampler

abstract type Sampler{E} end

gentype(::Type{<:Sampler{E}}) where {E} = E

# temporarily for BaseBenchmarks
RangeGenerator(x) = Sampler(default_rng(), x)

# In some cases, when only 1 random value is to be generated,
# the optimal sampler can be different than if multiple values
# have to be generated. Hence a `Repetition` parameter is used
# to choose the best one depending on the need.
const Repetition = Union{Val{1},Val{Inf}}

# these default fall-back for all RNGs would be nice,
# but generate difficult-to-solve ambiguities
# Sampler(::AbstractRNG, X, ::Val{Inf}) = Sampler(X)
# Sampler(::AbstractRNG, ::Type{X}, ::Val{Inf}) where {X} = Sampler(X)

"""
    Sampler(rng, x, repetition = Val(Inf))

Return a sampler object that can be used to generate random values from `rng` for `x`.

When `sp = Sampler(rng, x, repetition)`, `rand(rng, sp)` will be used to draw random values,
and should be defined accordingly.

`repetition` can be `Val(1)` or `Val(Inf)`, and should be used as a suggestion for deciding
the amount of precomputation, if applicable.

[`Random.SamplerType`](@ref) and [`Random.SamplerTrivial`](@ref) are default fallbacks for
*types* and *values*, respectively. [`Random.SamplerSimple`](@ref) can be used to store
pre-computed values without defining extra types for only this purpose.
"""
Sampler(rng::AbstractRNG, x, r::Repetition=Val(Inf)) = Sampler(typeof_rng(rng), x, r)
Sampler(rng::AbstractRNG, ::Type{X}, r::Repetition=Val(Inf)) where {X} =
    Sampler(typeof_rng(rng), X, r)

typeof_rng(rng::AbstractRNG) = typeof(rng)

# this method is necessary to prevent rand(rng::AbstractRNG, X) from
# recursively constructing nested Sampler types.
Sampler(T::Type{<:AbstractRNG}, sp::Sampler, r::Repetition) =
    throw(MethodError(Sampler, (T, sp, r)))

# default shortcut for the general case
Sampler(::Type{RNG}, X) where {RNG<:AbstractRNG} = Sampler(RNG, X, Val(Inf))
Sampler(::Type{RNG}, ::Type{X}) where {RNG<:AbstractRNG,X} = Sampler(RNG, X, Val(Inf))

#### pre-defined useful Sampler types

"""
    SamplerType{T}()

A sampler for types, containing no other information. The default fallback for `Sampler`
when called with types.
"""
struct SamplerType{T} <: Sampler{T} end

Sampler(::Type{<:AbstractRNG}, ::Type{T}, ::Repetition) where {T} = SamplerType{T}()

Base.getindex(::SamplerType{T}) where {T} = T

struct SamplerTrivial{T,E} <: Sampler{E}
    self::T
end

"""
    SamplerTrivial(x)

Create a sampler that just wraps the given value `x`. This is the default fall-back for
values.
The `eltype` of this sampler is equal to `eltype(x)`.

The recommended use case is sampling from values without precomputed data.
"""
SamplerTrivial(x::T) where {T} = SamplerTrivial{T,gentype(T)}(x)

Sampler(::Type{<:AbstractRNG}, x, ::Repetition) = SamplerTrivial(x)

Base.getindex(sp::SamplerTrivial) = sp.self

# simple sampler carrying data (which can be anything)
struct SamplerSimple{T,S,E} <: Sampler{E}
    self::T
    data::S
end

"""
    SamplerSimple(x, data)

Create a sampler that wraps the given value `x` and the `data`.
The `eltype` of this sampler is equal to `eltype(x)`.

The recommended use case is sampling from values with precomputed data.
"""
SamplerSimple(x::T, data::S) where {T,S} = SamplerSimple{T,S,gentype(T)}(x, data)

Base.getindex(sp::SamplerSimple) = sp.self

# simple sampler carrying a (type) tag T and data
struct SamplerTag{T,S,E} <: Sampler{E}
    data::S
    SamplerTag{T}(s::S) where {T,S} = new{T,S,gentype(T)}(s)
end

#### helper samplers

# TODO: make constraining constructors to enforce that those
# types are <: Sampler{T}

##### Adapter to generate a randome value in [0, n]

struct LessThan{T<:Integer,S} <: Sampler{T}
    sup::T
    s::S    # the scalar specification/sampler to feed to rand
end

function rand(rng::AbstractRNG, sp::LessThan)
    while true
        x = rand(rng, sp.s)
        x <= sp.sup && return x
    end
end

struct Masked{T<:Integer,S} <: Sampler{T}
    mask::T
    s::S
end

rand(rng::AbstractRNG, sp::Masked) = rand(rng, sp.s) & sp.mask

##### Uniform

struct UniformT{T} <: Sampler{T} end

uniform(::Type{T}) where {T} = UniformT{T}()

rand(rng::AbstractRNG, ::UniformT{T}) where {T} = rand(rng, T)

### machinery for generation with Sampler

# This describes how to generate random scalars or arrays, by generating a Sampler
# and calling rand on it (which should be defined in "generation.jl").
# NOTE: this section could be moved into a separate file when more containers are supported.

#### scalars

rand(rng::AbstractRNG, X)                                           = rand(rng, Sampler(rng, X, Val(1)))
# this is needed to disambiguate
rand(rng::AbstractRNG, X::Dims)                                     = rand(rng, Sampler(rng, X, Val(1)))
rand(rng::AbstractRNG=default_rng(), ::Type{X}=Float64) where {X} = rand(rng, Sampler(rng, X, Val(1)))::X

rand(X)                   = rand(default_rng(), X)
rand(::Type{X}) where {X} = rand(default_rng(), X)

#### arrays

rand!(A::AbstractArray{T}, X) where {T}             = rand!(default_rng(), A, X)
rand!(A::AbstractArray{T}, ::Type{X}=T) where {T,X} = rand!(default_rng(), A, X)

rand!(rng::AbstractRNG, A::AbstractArray{T}, X) where {T}             = rand!(rng, A, Sampler(rng, X))
rand!(rng::AbstractRNG, A::AbstractArray{T}, ::Type{X}=T) where {T,X} = rand!(rng, A, Sampler(rng, X))

function rand!(rng::AbstractRNG, A::AbstractArray{T}, sp::Sampler) where T
    for i in eachindex(A)
        @inbounds A[i] = rand(rng, sp)
    end
    A
end

rand(r::AbstractRNG, dims::Integer...) = rand(r, Float64, Dims(dims))
rand(                dims::Integer...) = rand(Float64, Dims(dims))

rand(r::AbstractRNG, X, dims::Dims)  = rand!(r, Array{gentype(X)}(undef, dims), X)
rand(                X, dims::Dims)  = rand(default_rng(), X, dims)

rand(r::AbstractRNG, X, d::Integer, dims::Integer...) = rand(r, X, Dims((d, dims...)))
rand(                X, d::Integer, dims::Integer...) = rand(X, Dims((d, dims...)))
# note: the above methods would trigger an ambiguity warning if d was not separated out:
# rand(r, ()) would match both this method and rand(r, dims::Dims)
# moreover, a call like rand(r, NotImplementedType()) would be an infinite loop

rand(r::AbstractRNG, ::Type{X}, dims::Dims) where {X} = rand!(r, Array{X}(undef, dims), X)
rand(                ::Type{X}, dims::Dims) where {X} = rand(default_rng(), X, dims)

rand(r::AbstractRNG, ::Type{X}, d::Integer, dims::Integer...) where {X} = rand(r, X, Dims((d, dims...)))
rand(                ::Type{X}, d::Integer, dims::Integer...) where {X} = rand(X, Dims((d, dims...)))

# SamplerUnion(X, Y, ...}) == Union{SamplerType{X}, SamplerType{Y}, ...}
SamplerUnion(U...) = Union{Any[SamplerType{T} for T in U]...}
const SamplerBoolBitInteger = SamplerUnion(Bool, BitInteger_types...)

include("Xoshiro.jl")
include("RNGs.jl")
include("generation.jl")
include("normal.jl")
include("misc.jl")
include("XoshiroSimd.jl")

## rand & rand! & seed! docstrings

"""
    rand([rng=default_rng()], [S], [dims...])

Pick a random element or array of random elements from the set of values specified by `S`;
`S` can be

* an indexable collection (for example `1:9` or `('x', "y", :z)`),
* an `AbstractDict` or `AbstractSet` object,
* a string (considered as a collection of characters), or
* a type: the set of values to pick from is then equivalent to `typemin(S):typemax(S)` for
  integers (this is not applicable to [`BigInt`](@ref)), to ``[0, 1)`` for floating
  point numbers and to ``[0, 1)+i[0, 1)`` for complex floating point numbers;

`S` defaults to [`Float64`](@ref).
When only one argument is passed besides the optional `rng` and is a `Tuple`, it is interpreted
as a collection of values (`S`) and not as `dims`.

!!! compat "Julia 1.1"
    Support for `S` as a tuple requires at least Julia 1.1.

# Examples
```julia-repl
julia> rand(Int, 2)
2-element Array{Int64,1}:
 1339893410598768192
 1575814717733606317

julia> using Random

julia> rand(MersenneTwister(0), Dict(1=>2, 3=>4))
1=>2

julia> rand((2, 3))
3

julia> rand(Float64, (2, 3))
2×3 Array{Float64,2}:
 0.999717  0.0143835  0.540787
 0.696556  0.783855   0.938235
```

!!! note
    The complexity of `rand(rng, s::Union{AbstractDict,AbstractSet})`
    is linear in the length of `s`, unless an optimized method with
    constant complexity is available, which is the case for `Dict`,
    `Set` and dense `BitSet`s. For more than a few calls, use `rand(rng,
    collect(s))` instead, or either `rand(rng, Dict(s))` or `rand(rng,
    Set(s))` as appropriate.
"""
rand

"""
    rand!([rng=default_rng()], A, [S=eltype(A)])

Populate the array `A` with random values. If `S` is specified
(`S` can be a type or a collection, cf. [`rand`](@ref) for details),
the values are picked randomly from `S`.
This is equivalent to `copyto!(A, rand(rng, S, size(A)))`
but without allocating a new array.

# Examples
```jldoctest
julia> rng = MersenneTwister(1234);

julia> rand!(rng, zeros(5))
5-element Vector{Float64}:
 0.5908446386657102
 0.7667970365022592
 0.5662374165061859
 0.4600853424625171
 0.7940257103317943
```
"""
rand!

"""
    seed!([rng=default_rng()], seed) -> rng
    seed!([rng=default_rng()]) -> rng

Reseed the random number generator: `rng` will give a reproducible
sequence of numbers if and only if a `seed` is provided. Some RNGs
don't accept a seed, like `RandomDevice`.
After the call to `seed!`, `rng` is equivalent to a newly created
object initialized with the same seed.

If `rng` is not specified, it defaults to seeding the state of the
shared task-local generator.

# Examples
```julia-repl
julia> Random.seed!(1234);

julia> x1 = rand(2)
2-element Vector{Float64}:
 0.32597672886359486
 0.5490511363155669

julia> Random.seed!(1234);

julia> x2 = rand(2)
2-element Vector{Float64}:
 0.32597672886359486
 0.5490511363155669

julia> x1 == x2
true

julia> rng = Xoshiro(1234); rand(rng, 2) == x1
true

julia> Xoshiro(1) == Random.seed!(rng, 1)
true

julia> rand(Random.seed!(rng), Bool) # not reproducible
true

julia> rand(Random.seed!(rng), Bool) # not reproducible either
false

julia> rand(Xoshiro(), Bool) # not reproducible either
true
```
"""
seed!(rng::AbstractRNG, ::Nothing) = seed!(rng)

end # module


using Printf
import MathOptInterface

const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities

# supported GAMS command line options
const MOI_SUPPORTED_CLOPTIONS = (
   "reslim",
   "iterlim",
   "holdfixed",
   "nodlim",
   "optca",
   "optcr",
   "solver",
   "threads",
   "trace",
   "traceopt",
   "license",
   "logoption",
   "lp",
   "mip",
   "rmip",
   "nlp",
   "dnlp",
   "cns",
   "minlp",
   "rminlp",
   "qcp",
   "miqcp",
   "rmiqcp"
)

# supported model types
const MOI_SUPPORTED_MODEL_TYPES = (
   GAMS.MODEL_TYPE_LP,
   GAMS.MODEL_TYPE_MIP,
   GAMS.MODEL_TYPE_RMIP,
   GAMS.MODEL_TYPE_NLP,
   GAMS.MODEL_TYPE_MINLP,
   GAMS.MODEL_TYPE_RMINLP,
   GAMS.MODEL_TYPE_QCP,
   GAMS.MODEL_TYPE_MIQCP,
   GAMS.MODEL_TYPE_RMIQCP,
   GAMS.MODEL_TYPE_MCP,
   GAMS.MODEL_TYPE_MPEC,
)

mutable struct VariableInfo
   name::String
   type::GAMSVarType
   lower_bound::Union{Float64, Nothing}
   lower_bound_dual_start::Union{Float64, Nothing}
   upper_bound::Union{Float64, Nothing}
   upper_bound_dual_start::Union{Float64, Nothing}
   start::Union{Float64, Nothing}
end

VariableInfo() = VariableInfo("", VARTYPE_FREE, nothing, nothing, nothing,
   nothing, nothing)

mutable struct ConstraintInfo{F, S}
   name::String
   func::F
   set::S
   dual_start::Union{Nothing, Float64}
end

ConstraintInfo(func, set) = ConstraintInfo("", func, set, nothing)

mutable struct Optimizer <: MOI.AbstractOptimizer
   gamswork::Union{GAMSWorkspace, Nothing}

   # problem attributes
   model_type::GAMSModelType
   n_binary::Int
   n_integer::Int
   n_semicont::Int
   n_semiint::Int
   m::Int
   m_lin::Int
   m_quad::Int
   m_nonlin::Int
   sense::MOI.OptimizationSense
   objective::Union{MOI.SingleVariable, MOI.ScalarAffineFunction{Float64}, MOI.ScalarQuadraticFunction{Float64}, Nothing}
   objvar::Bool
   variable_info::Vector{VariableInfo}
   linear_le_constraints::Vector{ConstraintInfo{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}}
   linear_ge_constraints::Vector{ConstraintInfo{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}}
   linear_eq_constraints::Vector{ConstraintInfo{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}}
   quadratic_le_constraints::Vector{ConstraintInfo{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}}
   quadratic_ge_constraints::Vector{ConstraintInfo{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}}
   quadratic_eq_constraints::Vector{ConstraintInfo{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}}
   sos1_constraints::Vector{ConstraintInfo{MOI.VectorOfVariables, MOI.SOS1{Float64}}}
   sos2_constraints::Vector{ConstraintInfo{MOI.VectorOfVariables, MOI.SOS2{Float64}}}
   nlp_data::Union{MOI.NLPBlockData, Nothing}
   complementarity_constraints::Vector{ConstraintInfo{MOI.VectorAffineFunction{Float64}, MOI.Complements}}

   # parameters
   sysdir::Union{String, Nothing}
   workdir::Union{String, Nothing}
   user_model_type::GAMSModelType
   gams_options::Dict{String, Any}
   solver_options::Dict{String, Any}

   # solution attributes
   solve_time::Float64
   solve_status::GAMSSolveStatus
   model_status::GAMSModelStatus
   sol::Union{GAMSSolution, Nothing}
   obj::Float64
   obj_est::Float64
end

function Optimizer(workspace::Union{Nothing, GAMSWorkspace} = nothing)
   gams_options = Dict{String, Any}()
   gams_options["threads"] = 1
   return Optimizer(workspace, MODEL_TYPE_UNDEFINED, 0, 0, 0, 0, 0, 0, 0, 0,
                    MOI.FEASIBILITY_SENSE, nothing, true, [], [], [], [], [], [],
                    [], [], [], nothing, [], nothing, nothing, MODEL_TYPE_UNDEFINED,
                    gams_options, Dict{String, Any}(), NaN, SOLVE_STATUS_UNDEFINED,
                    MODEL_STATUS_UNDEFINED, nothing, NaN, NaN)
end

MOI.get(::Optimizer, ::MOI.SolverName) = "GAMS"

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true
MOI.supports(::Optimizer, ::MOI.RawParameter) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.SingleVariable}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}) = true
MOI.supports(::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}) = true
MOI.supports(::Optimizer, ::MOI.VariableName, ::Type{MOI.VariableIndex}) = true
MOI.supports(::Optimizer, ::MOI.ConstraintName, ::Type{MOI.ConstraintIndex}) = true
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::Optimizer, ::MOI.NumberOfThreads) = true
MOI.supports(::Optimizer, ::MOI.NLPBlock) = true
MOI.supports(::Optimizer, ::MOI.NLPBlockDual) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.LessThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.GreaterThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.EqualTo{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.ZeroOne}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.Integer}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.Semicontinuous{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.SingleVariable}, ::Type{MOI.Semiinteger{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.VectorOfVariables}, ::Type{MOI.SOS1{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.VectorOfVariables}, ::Type{MOI.SOS2{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarAffineFunction{Float64}}, ::Type{MOI.LessThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarAffineFunction{Float64}}, ::Type{MOI.GreaterThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarAffineFunction{Float64}}, ::Type{MOI.EqualTo{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarQuadraticFunction{Float64}}, ::Type{MOI.LessThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarQuadraticFunction{Float64}}, ::Type{MOI.GreaterThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.ScalarQuadraticFunction{Float64}}, ::Type{MOI.EqualTo{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{MOI.VectorAffineFunction{Float64}}, ::Type{MOI.Complements}) = true
MOI.supports(::Optimizer, ::MOI.ConstraintDualStart, ::MOI.ConstraintIndex{MOI.SingleVariable, MOI.GreaterThan{Float64}}, ::Union{Real, Nothing}) = true
MOI.supports(::Optimizer, ::MOI.ConstraintDualStart, ::MOI.ConstraintIndex{MOI.SingleVariable, MOI.LessThan{Float64}}, ::Union{Real, Nothing}) = true
MOI.supports(::Optimizer, ::MOI.ConstraintDualStart, ::MOI.ConstraintIndex{MOI.SingleVariable, MOI.EqualTo{Float64}}, ::Union{Real, Nothing}) = true

MOIU.supports_default_copy_to(model::Optimizer, copy_names::Bool) = !copy_names

function MOI.copy_to(
   model::Optimizer,
   src::MOI.ModelLike;
   copy_names = false
)
   return MOIU.default_copy_to(model, src, copy_names)
end

function MOI.empty!(
   model::Optimizer
)
   model.model_type = MODEL_TYPE_UNDEFINED
   model.n_binary = 0
   model.n_integer = 0
   model.n_semicont = 0
   model.n_semiint = 0
   model.m = 0
   model.m_lin = 0
   model.m_quad = 0
   model.m_nonlin = 0
   model.sense = MOI.FEASIBILITY_SENSE
   model.objective = nothing
   empty!(model.variable_info)
   empty!(model.linear_le_constraints)
   empty!(model.linear_ge_constraints)
   empty!(model.linear_eq_constraints)
   empty!(model.quadratic_le_constraints)
   empty!(model.quadratic_ge_constraints)
   empty!(model.quadratic_eq_constraints)
   empty!(model.sos1_constraints)
   empty!(model.sos2_constraints)
   model.nlp_data = nothing
   empty!(model.complementarity_constraints)
   model.solve_time = NaN
   model.solve_status = SOLVE_STATUS_UNDEFINED
   model.model_status = MODEL_STATUS_UNDEFINED
   model.sol = nothing
   model.obj = NaN
   model.obj_est = NaN
end

function MOI.is_empty(
   model::Optimizer
)
   return model.model_type == MODEL_TYPE_UNDEFINED
      model.n_binary == 0 &&
      model.n_integer == 0 &&
      model.n_semicont == 0 &&
      model.n_semiint == 0 &&
      model.m == 0 &&
      model.m_lin == 0 &&
      model.m_quad == 0 &&
      model.m_nonlin == 0 &&
      model.sense == MOI.FEASIBILITY_SENSE &&
      isnothing(model.objective) &&
      isempty(model.variable_info) &&
      isempty(model.linear_le_constraints) &&
      isempty(model.linear_ge_constraints) &&
      isempty(model.linear_eq_constraints) &&
      isempty(model.quadratic_le_constraints) &&
      isempty(model.quadratic_ge_constraints) &&
      isempty(model.quadratic_eq_constraints) &&
      isempty(model.sos1_constraints) &&
      isempty(model.sos2_constraints) &&
      isnothing(model.nlp_data) &&
      isempty(model.complementarity_constraints)
end

has_start(var::VariableInfo) = ! isnothing(var.start)
has_upper_bound(var::VariableInfo) = ! isnothing(var.upper_bound)
has_lower_bound(var::VariableInfo) = ! isnothing(var.lower_bound)
is_fixed(var::VariableInfo) = has_lower_bound(var) && has_upper_bound(var) && var.lower_bound == var.upper_bound

has_upper_bound(model::Optimizer, vi::MOI.VariableIndex) = has_upper_bound(model.variable_info[vi.value])
has_lower_bound(model::Optimizer, vi::MOI.VariableIndex) = has_lower_bound(model.variable_info[vi.value])
is_fixed(model::Optimizer, vi::MOI.VariableIndex) = is_fixed(model.variable_info[vi.value])

offset_linear_le(model::Optimizer) = 0
offset_linear_ge(model::Optimizer) = length(model.linear_le_constraints)
offset_linear_eq(model::Optimizer) = offset_linear_ge(model) + length(model.linear_ge_constraints)
offset_quadratic_le(model::Optimizer) = offset_linear_eq(model) + length(model.linear_eq_constraints)
offset_quadratic_ge(model::Optimizer) = offset_quadratic_le(model) + length(model.quadratic_le_constraints)
offset_quadratic_eq(model::Optimizer) = offset_quadratic_ge(model) + length(model.quadratic_ge_constraints)
offset_nonlin(model::Optimizer) = offset_quadratic_eq(model) + length(model.quadratic_eq_constraints)
offset_complementarity(model::Optimizer) = 0

function MOI.set(
   model::Optimizer,
   ::MOI.NLPBlock,
   nlp_data::MOI.NLPBlockData
)
   model.nlp_data = nlp_data
   MOI.initialize(model.nlp_data.evaluator, [:ExprGraph])
   return
end

include("options.jl")
include("variables.jl")
include("objective.jl")
include("constraints.jl")
include("translate.jl")
include("solve.jl")


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct TGTypeI <: TurbineGov
        R::Float64
        Ts::Float64
        Tc::Float64
        T3::Float64
        T4::Float64
        T5::Float64
        P_min::Float64
        P_max::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of a Turbine Governor Type I.

# Arguments
- `R::Float64`: Droop parameter, validation range: (0, nothing)
- `Ts::Float64`: Governor time constant, validation range: (0, nothing)
- `Tc::Float64`: Servo time constant, validation range: (0, nothing)
- `T3::Float64`: Transient gain time constant, validation range: (0, nothing)
- `T4::Float64`: Power fraction time constant, validation range: (0, nothing)
- `T5::Float64`: Reheat time constant, validation range: (0, nothing)
- `P_min::Float64`: Min Power into the Governor, validation range: (0, nothing)
- `P_max::Float64`: Max Power into the Governor, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct TGTypeI <: TurbineGov
    "Droop parameter"
    R::Float64
    "Governor time constant"
    Ts::Float64
    "Servo time constant"
    Tc::Float64
    "Transient gain time constant"
    T3::Float64
    "Power fraction time constant"
    T4::Float64
    "Reheat time constant"
    T5::Float64
    "Min Power into the Governor"
    P_min::Float64
    "Max Power into the Governor"
    P_max::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function TGTypeI(R, Ts, Tc, T3, T4, T5, P_min, P_max, ext=Dict{String, Any}(), )
    TGTypeI(R, Ts, Tc, T3, T4, T5, P_min, P_max, ext, [:x_g1, :x_g2, :x_g3], 3, InfrastructureSystemsInternal(), )
end

function TGTypeI(; R, Ts, Tc, T3, T4, T5, P_min, P_max, ext=Dict{String, Any}(), )
    TGTypeI(R, Ts, Tc, T3, T4, T5, P_min, P_max, ext, )
end

# Constructor for demo purposes; non-functional.
function TGTypeI(::Nothing)
    TGTypeI(;
        R=0,
        Ts=0,
        Tc=0,
        T3=0,
        T4=0,
        T5=0,
        P_min=0,
        P_max=0,
        ext=Dict{String, Any}(),
    )
end

"""Get TGTypeI R."""
get_R(value::TGTypeI) = value.R
"""Get TGTypeI Ts."""
get_Ts(value::TGTypeI) = value.Ts
"""Get TGTypeI Tc."""
get_Tc(value::TGTypeI) = value.Tc
"""Get TGTypeI T3."""
get_T3(value::TGTypeI) = value.T3
"""Get TGTypeI T4."""
get_T4(value::TGTypeI) = value.T4
"""Get TGTypeI T5."""
get_T5(value::TGTypeI) = value.T5
"""Get TGTypeI P_min."""
get_P_min(value::TGTypeI) = value.P_min
"""Get TGTypeI P_max."""
get_P_max(value::TGTypeI) = value.P_max
"""Get TGTypeI ext."""
get_ext(value::TGTypeI) = value.ext
"""Get TGTypeI states."""
get_states(value::TGTypeI) = value.states
"""Get TGTypeI n_states."""
get_n_states(value::TGTypeI) = value.n_states
"""Get TGTypeI internal."""
get_internal(value::TGTypeI) = value.internal


"""
    AutoDiff.Jacobian(polar, func::Function, variable::AbstractVariable)

Instantiate a Jacobian AD factory for constraint function
`func`, w.r.t. state ``x`` (if `variable=State()`) or control
``u`` (if `variable=Control()`).

The coloring is done using Jacobian's expressions from MATPOWER.

### Examples

```julia
julia> Jacx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
```
"""
function AutoDiff.Jacobian(
    polar::PolarForm{T, VI, VT, MT}, func, variable,
) where {T, VI, VT, MT}
    @assert is_constraint(func)

    if isa(polar.device, CPU)
        SMT = SparseMatrixCSC{Float64,Int}
        A = Vector
    elseif isa(polar.device, CUDADevice)
        SMT = CUSPARSE.CuSparseMatrixCSR{Float64}
        A = CUDA.CuVector
    end

    pf = polar.network
    nbus = PS.get(pf, PS.NumberOfBuses())
    if isa(variable, State)
        map = VI(polar.mapx)
    elseif isa(variable, Control)
        map = VI(polar.mapu)
    end

    nmap = length(map)

    # Sparsity pattern
    J = jacobian_sparsity(polar, func, variable)

    # Coloring
    coloring = AutoDiff.SparseDiffTools.matrix_colors(J)
    ncolor = size(unique(coloring),1)

    # TODO: clean
    nx = 2 * nbus
    x = VT(zeros(Float64, nx))
    m = size(J, 1)

    # Move Jacobian to the GPU
    J = convert(SMT, J)

    # Seedings
    t1s{N} = ForwardDiff.Dual{Nothing,Float64, N} where N
    t1sx = A{t1s{ncolor}}(x)
    t1sF = A{t1s{ncolor}}(zeros(Float64, m))
    t1sseeds = AutoDiff.init_seed(coloring, ncolor, nmap)

    # Move the seeds over to the device, if necessary
    gput1sseeds = A{ForwardDiff.Partials{ncolor,Float64}}(t1sseeds)
    compressedJ = MT(zeros(Float64, ncolor, m))

    # Views
    varx = view(x, map)
    t1svarx = view(t1sx, map)

    return AutoDiff.Jacobian{typeof(func), VI, VT, MT, SMT, typeof(gput1sseeds), typeof(t1sx), typeof(varx), typeof(t1svarx)}(
        func, variable, J, compressedJ, coloring,
        gput1sseeds, t1sF, x, t1sx, map, varx, t1svarx
    )
end

"""
    AutoDiff.jacobian!(polar::PolarForm, jac::AutoDiff.Jacobian, buffer)

Update the sparse Jacobian entries `jacobian.J` using AutoDiff.
No allocations are taking place in this function.

* `polar::PolarForm`: polar formulation, stores all parameters.
* `jac::AutoDiff.Jacobian`: AutoDiff Factory with Jacobian to update.
* `buffer::PolarNetworkState`: store current values for network's variables.

"""
function AutoDiff.jacobian!(polar::PolarForm, jac::AutoDiff.Jacobian, buffer)
    nbus = get(polar, PS.NumberOfBuses())
    type = jac.var
    if isa(type, State)
        copyto!(jac.x, 1, buffer.vmag, 1, nbus)
        copyto!(jac.x, nbus+1, buffer.vang, 1, nbus)
        jac.t1sx .= jac.x
        jac.t1sF .= 0.0
    elseif isa(type, Control)
        copyto!(jac.x, 1, buffer.vmag, 1, nbus)
        copyto!(jac.x, nbus+1, buffer.pinj, 1, nbus)
        jac.t1sx .= jac.x
        jac.t1sF .= 0.0
    end

    AutoDiff.seed!(jac.t1sseeds, jac.varx, jac.t1svarx)

    if isa(type, State)
        jac.func(
            polar,
            jac.t1sF,
            view(jac.t1sx, 1:nbus),
            view(jac.t1sx, nbus+1:2*nbus),
            buffer.pinj,
            buffer.qinj,
        )
    elseif isa(type, Control)
        jac.func(
            polar,
            jac.t1sF,
            view(jac.t1sx, 1:nbus),
            buffer.vang,
            view(jac.t1sx, nbus+1:2*nbus),
            buffer.qinj,
        )
    end

    AutoDiff.getpartials_kernel!(jac.compressedJ, jac.t1sF)
    AutoDiff.uncompress_kernel!(jac.J, jac.compressedJ, jac.coloring)
    return jac.J
end

# Handle properly constant Jacobian case
function AutoDiff.ConstantJacobian(polar::PolarForm, func::Function, variable::Union{State,Control})
    @assert is_constraint(func)

    if isa(polar.device, CPU)
        SMT = SparseMatrixCSC{Float64,Int}
    elseif isa(polar.device, CUDADevice)
        SMT = CUSPARSE.CuSparseMatrixCSR{Float64}
    end

    nbus = get(polar, PS.NumberOfBuses())
    vmag = ones(nbus)
    vang = ones(nbus)
    V = vmag .* exp.(im .* vang)
    # Evaluate Jacobian with MATPOWER
    J = matpower_jacobian(polar, variable, func, V)
    # Move Jacobian to the GPU
    if isa(polar.device, CUDADevice) && iszero(J)
        # CUSPARSE does not support zero matrix. Return nothing instead.
        J = nothing
    else
        J = convert(SMT, J)
    end
    return AutoDiff.ConstantJacobian(J)
end

function AutoDiff.jacobian!(polar::PolarForm, jac::AutoDiff.ConstantJacobian, buffer)
    return jac.J
end

function AutoDiff.Hessian(polar::PolarForm{T, VI, VT, MT}, func) where {T, VI, VT, MT}
    @assert is_constraint(func)

    if isa(polar.device, CPU)
        A = Vector
    elseif isa(polar.device, CUDADevice)
        A = CUDA.CuVector
    end

    pf = polar.network
    nbus = PS.get(pf, PS.NumberOfBuses())
    n_cons = size_constraint(polar, func)

    map = VI(polar.hessianstructure.map)
    nmap = length(map)

    x = VT(zeros(Float64, 3*nbus))

    t1s{N} = ForwardDiff.Dual{Nothing,Float64, N} where N
    t1sx = A{t1s{1}}(x)
    t1sF = A{t1s{1}}(zeros(Float64, n_cons))
    host_t1sseeds = Vector{ForwardDiff.Partials{1,Float64}}(undef, nmap)
    t1sseeds = A{ForwardDiff.Partials{1,Float64}}(undef, nmap)
    varx = view(x, map)
    t1svarx = view(t1sx, map)
    VHP = typeof(host_t1sseeds)
    VP = typeof(t1sseeds)
    VD = typeof(t1sx)
    adj_t1sx = similar(t1sx)
    adj_t1sF = similar(t1sF)
    buffer = AutoDiff.TapeMemory(polar, func, VD; with_stack=false)
    return AutoDiff.Hessian{typeof(func), VI, VT, VHP, VP, VD, typeof(varx), typeof(t1svarx), typeof(buffer)}(
        func, host_t1sseeds, t1sseeds, x, t1sF, adj_t1sF, t1sx, adj_t1sx, map, varx, t1svarx, buffer,
    )
end

function _init_seed_hessian!(dest, tmp, v::AbstractArray, nmap)
    @inbounds for i in 1:nmap
        dest[i] = ForwardDiff.Partials{1, Float64}(NTuple{1, Float64}(v[i]))
    end
    return
end
function _init_seed_hessian!(dest, tmp, v::CUDA.CuArray, nmap)
    hostv = Array(v)
    @inbounds Threads.@threads for i in 1:nmap
        tmp[i] = ForwardDiff.Partials{1, Float64}(NTuple{1, Float64}(hostv[i]))
    end
    copyto!(dest, tmp)
    return
end

# λ' * H * v
function AutoDiff.adj_hessian_prod!(
    polar, H::AutoDiff.Hessian, hv, buffer, λ, v,
)
    @assert length(hv) == length(v)
    nbus = get(polar, PS.NumberOfBuses())
    x = H.x
    ntgt = length(v)
    t1sx = H.t1sx
    adj_t1sx = H.∂t1sx
    t1sF = H.t1sF
    adj_t1sF = H.∂t1sF
    # Move data
    copyto!(x, 1, buffer.vmag, 1, nbus)
    copyto!(x, nbus+1, buffer.vang, 1, nbus)
    copyto!(x, 2*nbus+1, buffer.pinj, 1, nbus)
    # Init dual variables
    t1sx .= H.x
    adj_t1sx .= 0.0
    t1sF .= 0.0
    adj_t1sF .= λ
    # Seeding
    nmap = length(H.map)

    # Init seed
    _init_seed_hessian!(H.t1sseeds, H.host_t1sseeds, v, nmap)
    AutoDiff.seed!(H.t1sseeds, H.varx, H.t1svarx)

    adjoint!(
        polar, H.buffer,
        t1sF, adj_t1sF,
        view(t1sx, 1:nbus), view(adj_t1sx, 1:nbus),                   # vmag
        view(t1sx, nbus+1:2*nbus), view(adj_t1sx, nbus+1:2*nbus),     # vang
        view(t1sx, 2*nbus+1:3*nbus), view(adj_t1sx, 2*nbus+1:3*nbus), # pinj
    )

    AutoDiff.getpartials_kernel!(hv, adj_t1sx, H.map)
    return nothing
end

# Adjoint's structure
"""
    AdjointStackObjective{VT}

An object for storing the adjoint stack in the adjoint objective computation.

"""
struct AdjointStackObjective{VT<:AbstractVector}
    ∇fₓ::VT
    ∇fᵤ::VT
    ∂pg::VT
    ∂vm::VT
    ∂va::VT
    ∂pinj::VT
    jvₓ::VT
    jvᵤ::VT
end

function AdjointStackObjective(polar::PolarForm{T, VI, VT, MT}) where {T, VI, VT, MT}
    nbus = get(polar, PS.NumberOfBuses())
    return AdjointStackObjective{VT}(
        xzeros(VT, get(polar, NumberOfState())),
        xzeros(VT, get(polar, NumberOfControl())),
        xzeros(VT, get(polar, PS.NumberOfGenerators())),
        xzeros(VT, nbus),
        xzeros(VT, nbus),
        xzeros(VT, nbus),
        xzeros(VT, get(polar, NumberOfState())),
        xzeros(VT, get(polar, NumberOfControl())),
    )
end

# Adjoint's stack for Polar
struct AdjointPolar{VT} <: AutoDiff.AbstractAdjointStack{VT}
    ∂vm::VT
    ∂va::VT
    ∂pinj::VT
    ∂qinj::VT
    ∂x::VT
    ∂u::VT
end

function AdjointPolar{VT}(nx::Int, nu::Int, nbus::Int) where {VT}
    return AdjointPolar{VT}(
        xzeros(VT, nbus),
        xzeros(VT, nbus),
        xzeros(VT, nbus),
        xzeros(VT, nbus),
        xzeros(VT, nx),
        xzeros(VT, nu),
    )
end

function reset!(adj::AdjointPolar)
    adj.∂vm .= 0.0
    adj.∂va .= 0.0
    adj.∂pinj .= 0.0
    adj.∂x .= 0.0
    adj.∂u .= 0.0
end

# Stack constructor for each constraint
function AdjointPolar(polar::PolarForm{T, VI, VT, MT}) where {T, VI, VT, MT}
    nbus = get(polar, PS.NumberOfBuses())
    nx = get(polar, NumberOfState())
    nu = get(polar, NumberOfControl())
    return AdjointPolar{VT}(nx, nu, nbus)
end

struct FullSpaceJacobian{Jacx,Jacu}
    x::Jacx
    u::Jacu
end

struct FullSpaceHessian{SpMT}
    xx::SpMT
    xu::SpMT
    uu::SpMT
end

struct ConstraintsJacobianStorage{SpMT}
    Jx::SpMT
    Ju::SpMT
    constraints_ad::Vector{FullSpaceJacobian}
end

function ConstraintsJacobianStorage(polar::PolarForm{T, VI, VT, MT}, constraints::Vector{Function}) where {T, VI, VT, MT}
    if isa(polar.device, CPU)
        SpMT = SparseMatrixCSC{Float64, Int}
    elseif isa(polar.device, CUDADevice)
        SpMT = CUSPARSE.CuSparseMatrixCSR{Float64}
    end

    SparseCPU = SparseMatrixCSC{Float64, Int}

    Jx = SparseCPU[]
    Ju = SparseCPU[]
    # Build global Jacobian on the CPU
    for cons in constraints
        push!(Jx, jacobian_sparsity(polar, cons, State()))
        push!(Ju, jacobian_sparsity(polar, cons, Control()))
    end
    gJx = convert(SpMT, vcat(Jx...))
    gJu = convert(SpMT, vcat(Ju...))

    # Build AD
    cons_ad = FullSpaceJacobian[]
    for cons in constraints
        jac_ad_x = _build_jacobian(polar, cons, State())
        jac_ad_u = _build_jacobian(polar, cons, Control())
        push!(cons_ad, FullSpaceJacobian(jac_ad_x, jac_ad_u))
    end

    return ConstraintsJacobianStorage{SpMT}(
        gJx,
        gJu,
        cons_ad,
    )
end

function update_full_jacobian!(
    polar::PolarForm,
    cons_jac::ConstraintsJacobianStorage{SpMT},
    buffer::PolarNetworkState
) where {SpMT}
    shift = 0
    for ad in cons_jac.constraints_ad
        # Update Jacobian
        Jx = AutoDiff.jacobian!(polar, ad.x, buffer)::SpMT
        Ju = AutoDiff.jacobian!(polar, ad.u, buffer)::Union{Nothing, SpMT}
        # Copy back results
        _transfer_sparse!(cons_jac.Jx, Jx, shift)
        if !isnothing(Ju)
            _transfer_sparse!(cons_jac.Ju, Ju, shift)
        end

        shift += size(Jx, 1)
    end
    return
end

struct HessianStorage{VT,Hess1,Hess2}
    state::Hess1
    obj::Hess2
    constraints::Vector{AutoDiff.Hessian}
    # Adjoints
    z::VT
    ψ::VT
    tmp_tgt::VT
    tmp_hv::VT
end

function HessianStorage(polar::PolarForm{T, VI, VT, MT}, constraints::Vector{Function}) where {T, VI, VT, MT}
    nx = get(polar, NumberOfState())
    nu = get(polar, NumberOfControl())

    Hstate = AutoDiff.Hessian(polar, power_balance)
    Hobj = AutoDiff.Hessian(polar, active_power_generation)
    Hcons = AutoDiff.Hessian[]
    for cons in constraints
        push!(Hcons, _build_hessian(polar, cons))
    end

    z = xzeros(VT, nx)
    ψ = xzeros(VT, nx)

    tgt = xzeros(VT, nx+nu)
    hv = xzeros(VT, nx+nu)

    return HessianStorage{VT, typeof(Hstate), typeof(Hobj)}(Hstate, Hobj, Hcons, z, ψ, tgt, hv)
end