const PLAYER_ID = 1
const OPPONENT_ID = 2

# @todo each game needs a action_space_size (see paper)
#action_space_size(game::AbstractEnv) = game.action_space_size
# @todo each game needs a discount (see paper)

# @todo change to "get_to_play"
to_play(game::AbstractEnv) = to_play(length(game.history) + 1)
to_play(position::Integer) = isodd(position) ? PLAYER_ID : OPPONENT_ID

const PLAYER_1 = 1
const PLAYER_2 = 2
OVERRIDE_MSG = "This should overridden in concrete implemention of AbstractState"

is_terminal(::AbstractEnv)                         = @assert false OVERRIDE_MSG
state(::AbstractEnv)                               = @assert false OVERRIDE_MSG
current_image(::AbstractEnv)                       = @assert false OVERRIDE_MSG
next_to_play(::AbstractEnv)                        = @assert false OVERRIDE_MSG
legal_actions(::AbstractEnv)                       = @assert false OVERRIDE_MSG
action_space_size(::AbstractEnv)                   = @assert false OVERRIDE_MSG
append_observation!(::AbstractEnv)                 = @assert false OVERRIDE_MSG
islegal(::AbstractEnv, ::Integer)                  = @assert false OVERRIDE_MSG
apply!(::AbstractEnv, ::Integer)                   = @assert false OVERRIDE_MSG
action_image(::Type{<:AbstractEnv}, ::Integer)     = @assert false OVERRIDE_MSG
make_image(::AbstractEnv, ::Integer)               = @assert false OVERRIDE_MSG
get_state(::AbstractEnv, ::Integer)                = @assert false OVERRIDE_MSG
make_image_label(::AbstractEnv, ::Integer)         = @assert false OVERRIDE_MSG
reward_matrix(::AbstractEnv)                       = @assert false OVERRIDE_MSG
policy_matrix(::AbstractEnv)                       = @assert false OVERRIDE_MSG
value_matrix(::AbstractEnv, ::Integer; td_steps)   = @assert false OVERRIDE_MSG
to_play_at(::AbstractEnv, ::Integer)               = @assert false OVERRIDE_MSG
discount(::AbstractEnv)                            = @assert false OVERRIDE_MSG
getstatesplus(::AbstractEnv)                       = @assert false OVERRIDE_MSG
getstates(::AbstractEnv)                           = @assert false OVERRIDE_MSG
getconfig(::Type{<:AbstractEnv})                   = @assert false OVERRIDE_MSG
gettestfn(::Type{<:AbstractEnv})                   = @assert false OVERRIDE_MSG
state_type(::Type{<:AbstractEnv})                  = @assert false OVERRIDE_MSG
policy_testset(::Type{<:AbstractEnv})              = @assert false OVERRIDE_MSG
gridsvg(::AbstractEnv, ::Integer)                  = @assert false OVERRIDE_MSG

state_type(x::AbstractEnv) = state_type(typeof(x))
#Base.length(x::AbstractEnv) = length(history(x)) + 1
#Base.getindex(x::AbstractEnv, i::Int) = get_state(x, i)
#Base.iterate(x::AbstractEnv, state_idx=1) = state_idx > length(x) ? nothing : (x[state_idx], state_idx+1)

function cartesian_coords_to_action(
    ::Type{<:AbstractEnv},
    x::Int,
    y::Int)
    @assert false OVERRIDE_MSG
end

function apply!(game::AbstractEnv, actions::Vector{<:Integer})
    actions .|> (x) -> apply!(game, x)
end

uniform_policy(game::AbstractEnv) = ones(Float32, action_space_size(game)) ./ action_space_size(game)
action_history(g::AbstractEnv) = ActionHistory(g)
action_history(g::AbstractEnv, idx) = ActionHistory(g, idx)
history(game::AbstractEnv) = game.history

# {Start} A1 => (R=0.0, V=0.011154757, π=[0.64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.36])


# define the type of data

mutable struct fixedData
    # static network data
    baseMVA :: Float64
    bType :: Dict{Int64,Any}

    IDList :: Array{Int64,1}
    genIDList :: Array{Int64,1}
    brList :: Array{Any,1}
    brRev :: Dict{Any,Any}

    Loc :: Dict{Int64,Any}
    LocRev :: Dict{Int64,Any}
    Vmax :: Dict{Int64,Any}
    Vmin :: Dict{Int64,Any}
    Pmax :: Dict{Int64,Any}
    Pmin :: Dict{Int64,Any}
    Qmax :: Dict{Int64,Any}
    Qmin :: Dict{Int64,Any}
    gs :: Dict{Int64,Any}
    bs :: Dict{Int64,Any}
    Vmag :: Dict{Int64,Any}
    Vang :: Dict{Int64,Any}
    Pd :: Dict{Int64,Any}
    Qd :: Dict{Int64,Any}
    Pg :: Dict{Int64,Any}
    Qg :: Dict{Int64,Any}
    RU :: Dict{Int64,Any}
    RD :: Dict{Int64,Any}

    g :: Dict{Tuple{Int64,Int64,Int64},Any}
    b :: Dict{Tuple{Int64,Int64,Int64},Any}
    bc :: Dict{Tuple{Int64,Int64,Int64},Any}
    θmax :: Dict{Tuple{Int64,Int64,Int64},Any}
    θmin :: Dict{Tuple{Int64,Int64,Int64},Any}
    rateA :: Dict{Tuple{Int64,Int64,Int64},Any}
    τ1 :: Dict{Tuple{Int64,Int64,Int64},Any}
    τ2 :: Dict{Tuple{Int64,Int64,Int64},Any}
    σ :: Dict{Tuple{Int64,Int64,Int64},Any}

    cp :: Dict{Any,Any}
    cq :: Dict{Any,Any}
    cz :: Any

    busInd :: Dict{Any,Any}
    branchDict1 :: Dict{Any,Any}
    branchDict2 :: Dict{Any,Any}
    connectPair :: Array{Any,1}
    connectDict :: Dict{Any,Any}
    kpDict :: Dict{Any,Any}

end

struct costDataType
    model :: Int64
    upCost :: Float64
    downCost :: Float64
    n :: Int64
    params :: Array{Float64,1}
end

struct probDistrn
    # probability distribution of disruption timing
    tDistrn :: Dict{Int64,Any}

    # probability distribution of disruption magnitude
    # ωDistrn: keys are the indices, values are probability
    # ωDict: keys are the indices, values are lists of broken components
    ωDistrn :: Dict{Any,Any}
    ωDict :: Dict{Any,Any}
    ωτ :: Dict{Any,Any}
end

mutable struct batteryData
    # battery information: charging/discharging factor, capacity
    IDList :: Array{Any,1}
    Loc :: Dict{Int64,Any}
    bInv :: Dict{Int64,Any}
    ηα :: Dict{Int64,Any}
    ηβ :: Dict{Int64,Any}
    cap :: Dict{Int64,Any}
    cost :: Dict{Int64,Any}
    uCap :: Dict{Int64,Any}
end

struct solData
    # solution information: generaion level, storage level, maximum generation output
    sp :: Dict{Any,Any}
    sq :: Dict{Any,Any}
    w :: Dict{Any,Any}
    u :: Dict{Any,Any}
    lp :: Dict{Any,Any}
    lq :: Dict{Any,Any}
    zp :: Dict{Any,Any}
end

struct demandData
    # demand data: active demand and reactive demand
    T :: Int64
    pd :: Dict{Any,Any}
    qd :: Dict{Any,Any}
end

struct cutData
    # cut data: slope for sp, w and u and the intercept
    solStatus :: Any
    λ :: Dict{Any,Any}
    γ :: Dict{Any,Any}
    μ :: Dict{Any,Any}
    vhat :: Float64
end


include("dm.jl")
include("tau.jl")

function _sim!(sys::System, pathEnsemble::Ensemble, i::Int64)
    for k in eachindex(pathEnsemble)
        _directMethod!(pathEnsemble[k], sys.model.F, sys.model.ν, sys.times[i], sys.model.nr)
    end
end

function _sim!(sys::System, pathEnsemble::Ensemble, i::Int64, tauVar::TauVar)
    for k in eachindex(pathEnsemble)
        _optimisedTauLeap!(sys.model, pathEnsemble[k], sys.times[i], sys.routine.tauOpt, tauVar)
    end
end


abstract type AbstractCut end

struct FeasibilityCut{T, VT<:AbstractVector{T}} <: AbstractCut
    a::VT
    β::T
end

abstract type AbstractOptimalityCut <: AbstractCut end

struct MultiOptimalityCut{T, VT<:AbstractVector{T}} <: AbstractCut
    a::VT
    β::T
end

struct AveragedOptimalityCut{T, VT<:AbstractVector{T}} <: AbstractCut
    a::VT
    β::T
end

function SOI.addcut!(sp::StochasticProgram, node, pool::SOI.AbstractSolutionPool, to::TimerOutput, ztol)
    if SOI.allfeasible(pool)
        gencut(SOI.get(sp, CutGenerator(), node), sp, node, pool, to::TimerOutput, ztol)
    else
        gencut(FeasibilityCutGenerator(), sp, node, pool, to::TimerOutput, ztol)
    end
end

function SOI.addcut!(sp::StochasticProgram, tr::Transition, cut::FeasibilityCut)
    # coef is a ray
    # so alpha * coef is also valid for any alpha >= 0.
    # Hence coef might have very large coefficients and alter
    # the numerial accuracy of the master's solver.
    # We scale it to avoid this issue
    scaling = max(abs(cut.β), maximum(abs, cut.a))
    addcut(nodedata(sp, SOI.get(sp, SOI.Target(), tr)).fcuts, cut.a/scaling, sign(cut.β), nodedata(sp, SOI.get(sp, SOI.Source(), tr)).nlds)
end
function SOI.addcut!(sp::StochasticProgram, tr::Transition, cut::MultiOptimalityCut)
    addcut(nodedata(sp, SOI.get(sp, SOI.Target(), tr)).ocuts, cut.a, cut.β, nodedata(sp, SOI.get(sp, SOI.Source(), tr)).nlds)
end
function SOI.addcut!(sp::StochasticProgram, node, cut::AveragedOptimalityCut)
    addcut(nodedata(sp, node).nlds.localOC, cut.a, cut.β, nodedata(sp, node).nlds)
end

function SOI.applycuts!(sp::StochasticProgram, node)
    applycut(FeasibilityCutGenerator(), sp, node)
    applycut(SOI.get(sp, CutGenerator(), node), sp, node)
end

function SOI.applycuts!(sp::StochasticProgram, tr::Transition, ::Type{<:FeasibilityCut})
    apply!(nodedata(sp, SOI.get(sp, SOI.Target(), tr)).fcuts)
end
function SOI.applycuts!(sp::StochasticProgram, tr::Transition, ::Type{<:MultiOptimalityCut})
    apply!(nodedata(sp, SOI.get(sp, SOI.Target(), tr)).ocuts)
end
function SOI.applycuts!(sp::StochasticProgram, node, ::Type{<:AveragedOptimalityCut})
    apply!(nodedata(sp, node).nlds.localOC)
end


@testset "inset.jl" begin
    @testset "InSet()" begin
        trailer = SeaPearl.Trailer()
        a = SeaPearl.IntSetVar(2, 5, "a", trailer)
        x = SeaPearl.IntVar(0, 3, "x", trailer)

        constraint = SeaPearl.InSet(x, a, trailer)

        @test constraint in a.onDomainChange
        @test constraint in x.onDomainChange
        @test constraint.active.value
    end

    @testset "propagate!(::InSet)" begin
        trailer = SeaPearl.Trailer()
        a = SeaPearl.IntSetVar(2, 5, "a", trailer)
        x = SeaPearl.IntVar(0, 3, "x", trailer)
        constraint = SeaPearl.InSet(x, a, trailer)

        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        @test SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test length(x.domain) == 2
        @test !(0 in x.domain)
        @test !(1 in x.domain)
        @test 2 in x.domain
        @test constraint.active.value
        @test prunedDomains == SeaPearl.CPModification("x" => [0, 1])

        SeaPearl.assign!(x, 3)

        @test SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test SeaPearl.required_values(a.domain) == Set{Int}([3])
        @test SeaPearl.possible_not_required_values(a.domain) == Set{Int}([2, 4, 5])
        @test !constraint.active.value

        # Infeasibility
        trailer = SeaPearl.Trailer()
        a = SeaPearl.IntSetVar(2, 5, "a", trailer)
        x = SeaPearl.IntVar(0, 1, "x", trailer)
        constraint = SeaPearl.InSet(x, a, trailer)

        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        @test !SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        #Infeasibility with excluding
        trailer = SeaPearl.Trailer()
        a = SeaPearl.IntSetVar(2, 5, "a", trailer)
        x = SeaPearl.IntVar(2, 3, "x", trailer)
        constraint = SeaPearl.InSet(x, a, trailer)

        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        SeaPearl.exclude!(a.domain, 2)
        SeaPearl.exclude!(a.domain, 3)

        @test !SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        # Deactivation
        trailer = SeaPearl.Trailer()
        a = SeaPearl.IntSetVar(2, 5, "a", trailer)
        x = SeaPearl.IntVar(2, 3, "x", trailer)
        constraint = SeaPearl.InSet(x, a, trailer)

        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        SeaPearl.require!(a.domain, 2)
        SeaPearl.require!(a.domain, 3)

        @test SeaPearl.propagate!(constraint, toPropagate, prunedDomains)
        @test !constraint.active.value
    end
end


module ScoringTests
  using Base.Test
  include("../src/scoring.jl")

  # Offensive scoring
  @test score_off() == 0
  @test score_off(rush_yds=1) == 0
  @test score_off(rush_yds=10) == 1
  @test score_off(rush_td=1) == 6
  @test score_off(rec_yds=1) == 0
  @test score_off(rec_yds=10) == 1
  @test score_off(rec_td=1) == 6
  @test score_off(pass_yds=10) == 0
  @test score_off(pass_yds=25) == 1
  @test score_off(pass_td=1) == 4
  @test score_off(pass_int=1) == -2
  @test score_off(sack=1) == 0
  @test score_off(fum_l=1) == -2
  @test score_off(pr_td=1) == 6
  @test score_off(kr_td=1) == 6

  @test score_off(rush_yds=10, rush_td=1, rec_yds=10, rec_td=1, pass_yds=25, pass_td=1, pass_int=1, sack=1, fum_l=1, pr_td=1, kr_td=1) == (1 + 6 + 1 + 6 + 1 + 4 + -2 + 0 + -2 + 6 + 6)

  # Kicker scoring
  @test score_k() == 0
  @test score_k(xp_made=1) == 1
  @test score_k(xp_att=1) == -1
  @test score_k(fg_made=1) == 0
  @test score_k(fg_att=1) == -1
  @test score_k(f0019=1) == 3
  @test score_k(f2029=1) == 3
  @test score_k(f3039=1) == 3
  @test score_k(f4049=1) == 4
  @test score_k(f50=1) == 5

  @test score_k(xp_made=1, xp_att=2, fg_made=1, fg_att=2, f0019=1, f2029=1, f3039=1, f4049=1, f50=1) == (1 + -1 + -1 + 3 + 3 + 3 + 4 + 5)

  # Defense / Special Teams scoring
  @test score_def_st() == 0
  @test score_def_st(games=1) == 5
  @test score_def_st(games=2) == 10
  @test score_def_st(papg=1) == 0
  @test score_def_st(games=1, papg=0) == 5
  @test score_def_st(games=1, papg=6) == 4
  @test score_def_st(games=1, papg=13) == 3
  @test score_def_st(games=1, papg=17) == 1
  @test score_def_st(games=1, papg=27) == 0
  @test score_def_st(games=1, papg=34) == -1
  @test score_def_st(games=1, papg=45) == -3
  @test score_def_st(games=1, papg=46) == -5
  @test score_def_st(kr_td=1) == 6
  @test score_def_st(pr_td=1) == 6
  @test score_def_st(int_td=1) == 6
  @test score_def_st(fum_td=1) == 6
  @test score_def_st(sack=1) == 1
  @test score_def_st(sfty=1) == 2

  @test score_def_st(games=1, papg=1, kr_td=1, pr_td=1, int_td=1, fum_td=1, sack=1, sfty=1) == (4 + 6 + 6 + 6 + 6 + 1 + 2)

end


@testset "Datastructures" begin
    # Test the types and the sizes of the moments
    # For 64 bit floating numbers, one dimensional
    @testset "One dimensional moments" begin
        n = 10
        typelist = [Float16 Float32 Float64 Real]
        for type in typelist
            moments = JuSwalbe.Macroquant{Vector{type},Vector{type}}(zeros(type,n),zeros(type,n),zeros(type,n),zeros(type,n))
            @test isa(moments, JuSwalbe.Macroquant)
            test_dict = struct2dict(moments)
            for key in keys(test_dict)
                @test isa(test_dict[key], Array{type,1})
                @test size(test_dict[key], 1) == n
            end
        end
    end
    @testset "One dimensional forces" begin
        n = 10
        typelist = [Float16 Float32 Float64 Real]
        for type in typelist
            forces = JuSwalbe.Forces{Vector{type}}(zeros(type,n),zeros(type,n),zeros(type,n),zeros(type,n))
            @test isa(forces, JuSwalbe.Forces)
            test_dict = struct2dict(forces)
            for key in keys(test_dict)
                @test isa(test_dict[key], Array{type,1})
                @test size(test_dict[key], 1) == n
            end
        end
    end
    # For 32 bit floating numbers, one dimensional
    @testset "Two dimensional moments" begin
        m,n = 10,8
        typelist = [Float16 Float32 Float64 Real]
        for type in typelist
            twovec = JuSwalbe.Twovector{Matrix{type}}(zeros(type, (m,n)), zeros(type, (m,n)))
            moments = JuSwalbe.Macroquant{Matrix{type},JuSwalbe.Twovector{Matrix{type}}}(zeros(type, (m,n)),twovec ,zeros(type, (m,n)),zeros(type, (m,n)))
            @test isa(moments, JuSwalbe.Macroquant)
            test_dict = struct2dict(moments)
            for key in keys(test_dict)
                if key == :velocity
                    @test isa(test_dict[key], JuSwalbe.Twovector)
                    @test isa(test_dict[key].x, Matrix{type})
                    @test isa(test_dict[key].y, Matrix{type})
                    @test size(test_dict[key].x) == (m,n)
                    @test size(test_dict[key].y) == (m,n)
                else
                    @test isa(test_dict[key], Array{type,2})
                    @test size(test_dict[key], 1) == m
                    @test size(test_dict[key], 2) == n
                end
            end
        end
    end
    @testset "Two dimensional forces" begin
        n, m = 10, 6
        typelist = [Float16 Float32 Float64 Real]
        for type in typelist
            twovec = JuSwalbe.Twovector{Matrix{type}}(zeros(type, (n,m)),zeros(type, (n,m)))
            forces = JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{type}}}(twovec, twovec, twovec, twovec)
            @test isa(forces, JuSwalbe.Forces)
            test_dict = struct2dict(forces)
            for key in keys(test_dict)
                @test isa(test_dict[key], JuSwalbe.Twovector)
                @test isa(test_dict[key].x, Matrix{type})
                @test isa(test_dict[key].y, Matrix{type})
                @test size(test_dict[key].x) == (n,m)
                @test size(test_dict[key].y) == (n,m)
            end
        end
    end
end


using Distributed: @everywhere, @spawnat, @distributed
@everywhere include("gru_dynamic_struct.jl")

const window_size = 50

create_model_definitions(layer) =
begin
    eval(Meta.parse("@everywhere " * "struct Model\n" * *(["l$i::Layer\n" for i in 1:length(layers) +1]...) * "end"))
    eval(Meta.parse("@everywhere " * "(model::Model)(io) =\n" * "begin\n" * *(["io = model.l$i(io)\n" for i in 1:length(layers) +1]...) * "end"))
end

soften = arr -> (begin
    new_arr = []
    soft = softmax(arr[1:Int(length(arr)/4)])
    for i in 1:Int(length(arr)/4)
        push!(new_arr, soft[i])
    end
    for i in Int(length(arr)/4):Int(length(arr)*3/4)
        push!(new_arr, sigm(arr[i]))
    end
new_arr
end)

prop(model, x) =
begin
    # for mfield in fieldnames(Model)
    #     layer = getfield(model, mfield)
    #     layer.state = zeros(1, length(layer.bs))
    # end
[soften(model(t)) for t in x]
end

cross_entropy(out, label) =
begin
    - sum(label .* log.(out))
end

train!(model, data, lr, batch_size) =
begin
    results = []
    for i in 1:trunc(Int, length(data)/batch_size)
        res = @distributed (vcat) for sequence in data[(i-1)*batch_size+1:i*batch_size]
            d = @diff (begin
                    sequence_parts = batchify(sequence, window_size)
                    part_losses = []
                    for part in sequence_parts
                        x = sequence[1:end-1]
                        y = sequence[2:end]
                        o = prop(model, x)
                        l = [cross_entropy(o_t, y_t) for (o_t, y_t) in zip(o, y)]
                        push!(part_losses, sum(l))
                    end
                sum(part_losses) end)
            grads = []
            for mfield in fieldnames(Model)
                layer = getfield(model, mfield)
                for lfield in fieldnames(Layer)
                    push!(grads, grad(d, getfield(layer, lfield)))
                end
            end
            @spawnat 1 print("/")
            grads, value(d)
        end
        push!(results, res)
    end

    loss = 0.0
    for (g,l) in results
        loss += l
        i = 0
        for mfield in fieldnames(Model)
            layer = getfield(model, mfield)
            for lfield in fieldnames(Layer)
                i +=1
                setfield!(layer, lfield, Param(getfield(layer, lfield) - g[i] .* lr))
            end
        end
    end
    @show loss
end


#TODO: Make inmmutable later. Requires refactor of ThreePhase Fault callbacks
mutable struct SimulationInputs
    sys::PSY.System
    injectors_data::Vector{<:PSY.StaticInjection}
    counts::Base.ImmutableDict{Symbol, Int}
    Ybus::SparseMatrixCSC{Complex{Float64}, Int64}
    dyn_lines::Bool
    voltage_buses_ix::Vector{Int}
    current_buses_ix::Vector{Int}
    global_index::Dict{String, Dict{Symbol, Int64}}
    total_shunts::Dict{Int64, Float64}
    global_vars::Dict{Symbol, Number}
    lookup::Dict{Int, Int}
    DAE_vector::Vector{Bool}
    aux_arrays::Dict{Int, Vector}
    tspan::NTuple{2, Float64}
end

function SimulationInputs(;
    sys::PSY.System,
    dynamic_injectors = Vector{PSY.StaticInjection}(),
    counts::Base.ImmutableDict{Symbol, Int} = Base.ImmutableDict{Symbol, Int}(),
    Ybus::SparseMatrixCSC{Complex{Float64}, Int64} = SparseMatrixCSC{
        Complex{Float64},
        Int64,
    }(zeros(1, 1)),
    dyn_lines::Bool = false,
    voltage_buses_ix::Vector{Int} = Vector{Int}(),
    current_buses_ix::Vector{Int} = Vector{Int}(),
    global_index::Dict{String, Dict{Symbol, Int64}} = Dict{String, Dict{Symbol, Int64}}(),
    total_shunts::Dict{Int64, Float64} = Dict{Int64, Float64}(),
    global_vars::Dict{Symbol, Number} = Dict{Symbol, Number}(),
    lookup::Dict{Int, Int} = Dict{Int, Int}(),
    DAE_vector::Vector{Bool} = Vector{Bool}(),
    aux_arrays::Dict{Int, Vector} = Dict{Int, Vector}(),
    tspan::NTuple{2, Float64} = (0.0, 0.0),
)
    return SimulationInputs(
        sys,
        dynamic_injectors,
        counts,
        Ybus,
        dyn_lines,
        voltage_buses_ix,
        current_buses_ix,
        global_index,
        total_shunts,
        global_vars,
        lookup,
        DAE_vector,
        aux_arrays,
        tspan,
    )
end

function build!(inputs::SimulationInputs)
    sys = get_system(inputs)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    DAE_vector = collect(falses(n_buses * 2))
    global_state_index = MAPPING_DICT()
    state_space_ix = [n_buses * 2]
    current_buses_ix = collect(1:n_buses)
    static_bus_var_count = 2 * length(current_buses_ix)
    voltage_buses_ix = Vector{Int}()
    total_states = 0
    first_dyn_branch_point = -1
    branches_n_states = 0
    global_vars = Dict{Symbol, Number}(
        :ω_sys => 1.0,
        :ω_sys_index => -1, #To define 0 if infinite source, bus_number otherwise,
    )
    total_shunts = Dict{Int, Float64}()
    found_ref_bus = false
    sys_basepower = PSY.get_base_power(sys)

    inputs.Ybus, inputs.lookup = _get_Ybus(sys)
    dyn_branches = PSY.get_components(PSY.DynamicBranch, sys)

    if !(isempty(dyn_branches))
        inputs.dyn_lines = true
        first_dyn_branch_point = state_space_ix[1] + 1
        for br in dyn_branches
            arc = PSY.get_arc(br)
            n_states = PSY.get_n_states(br)
            from_bus_number = PSY.get_number(arc.from)
            to_bus_number = PSY.get_number(arc.to)
            bus_ix_from = get_lookup(inputs)[from_bus_number]
            bus_ix_to = get_lookup(inputs)[to_bus_number]
            merge!(
                +,
                total_shunts,
                Dict(
                    bus_ix_from => 1 / PSY.get_b(br).from,
                    bus_ix_to => 1 / PSY.get_b(br).to,
                ),
            )
            push!(voltage_buses_ix, bus_ix_from, bus_ix_to)
            DAE_vector[bus_ix_from] = DAE_vector[bus_ix_from + n_buses] = true
            DAE_vector[bus_ix_to] = DAE_vector[bus_ix_to + n_buses] = true
            DAE_vector = push!(DAE_vector, collect(trues(n_states))...)
            total_states += n_states
            _add_states_to_global!(global_state_index, state_space_ix, br)
        end

        for (ix, val) in enumerate(DAE_vector[1:n_buses])
            if val
                #This V_ix should be V_number.
                global_state_index["V_$(ix)"] = Dict(:R => ix, :I => ix + n_buses)
                total_states += 2
                static_bus_var_count -= 2
                push!(voltage_buses_ix, ix)
                @assert static_bus_var_count >= 0
            end
        end
        branches_n_states = state_space_ix[1] - n_buses * 2
    else
        inputs.dyn_lines = false
        @debug("System doesn't contain Dynamic Branches")
    end
    unique!(voltage_buses_ix)
    sources = PSY.get_components(PSY.Source, sys)
    for s in sources
        btype = PSY.get_bustype(PSY.get_bus(s))
        if (btype == PSY.BusTypes.REF) && found_ref_bus
            throw(IS.ConflictingInputsError("The system can't have more than one source or generator in the REF Bus"))
        end
        btype != PSY.BusTypes.REF && continue
        global_vars[:ω_sys_index] = 0 #To define 0 if infinite source, bus_number otherwise,
        found_ref_bus = true
    end

    dynamic_injection = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x -> PSY.get_dynamic_injector(x) !== nothing,
    )

    if isempty(dynamic_injection)
        error("System doesn't contain any DynamicInjection devices")
    end

    for d in dynamic_injection
        @debug PSY.get_name(d)
        dynamic_device = PSY.get_dynamic_injector(d)
        if !(:states in fieldnames(typeof(dynamic_device)))
            continue
        end
        device_bus = PSY.get_bus(d)
        btype = PSY.get_bustype(device_bus)
        if (btype == PSY.BusTypes.REF) && found_ref_bus
            throw(IS.ConflictingInputsError("The system can't have more than one source or generator in the REF Bus"))
        end
        _make_device_index!(d)
        device_n_states = PSY.get_n_states(dynamic_device)
        DAE_vector = push!(DAE_vector, collect(trues(device_n_states))...)
        total_states += device_n_states
        _add_states_to_global!(global_state_index, state_space_ix, dynamic_device)
        push!(inputs.injectors_data, d)

        btype != PSY.BusTypes.REF && continue
        if typeof(dynamic_device) <: PSY.DynamicGenerator
            ω_ix = global_state_index[PSY.get_name(d)][:ω]
        elseif typeof(dynamic_device) <: PSY.DynamicInverter
            #TO DO: Make it general for cases when ω is not a state (droop)!
            ω_ix = global_state_index[PSY.get_name(d)][:ω_oc]
        end
        global_vars[:ω_sys_index] = ω_ix #To define 0 if infinite source, bus_number otherwise,
        found_ref_bus = true
    end
    injection_n_states = state_space_ix[1] - branches_n_states - n_buses * 2
    @assert total_states == state_space_ix[1] - static_bus_var_count
    @debug total_states
    setdiff!(current_buses_ix, voltage_buses_ix)

    inputs.global_vars = global_vars
    inputs.DAE_vector = DAE_vector
    inputs.global_index = global_state_index
    inputs.voltage_buses_ix = voltage_buses_ix
    inputs.current_buses_ix = current_buses_ix
    inputs.total_shunts = total_shunts
    inputs.counts = Base.ImmutableDict(
        :total_states => total_states,
        :injection_n_states => injection_n_states,
        :branches_n_states => branches_n_states,
        :first_dyn_injection_pointer => 2 * n_buses + branches_n_states + 1,
        :first_dyn_branch_point => first_dyn_branch_point,
        :total_variables => total_states + static_bus_var_count,
        :bus_count => n_buses,
    )

    @assert get_ω_sys(inputs) != -1
    return
end

get_system(inputs::SimulationInputs) = inputs.sys
get_injectors_data(inputs::SimulationInputs) = inputs.injectors_data
get_counts(inputs::SimulationInputs) = inputs.counts
get_voltage_buses_ix(inputs::SimulationInputs) = inputs.voltage_buses_ix
get_current_buses_ix(inputs::SimulationInputs) = inputs.current_buses_ix
get_global_index(inputs::SimulationInputs) = inputs.global_index
get_Ybus(inputs::SimulationInputs) = inputs.Ybus
get_total_shunts(inputs::SimulationInputs) = inputs.total_shunts
get_global_vars(inputs::SimulationInputs) = inputs.global_vars
get_dyn_lines(inputs::SimulationInputs) = inputs.dyn_lines
get_lookup(inputs::SimulationInputs) = inputs.lookup
get_DAE_vector(inputs::SimulationInputs) = inputs.DAE_vector
get_aux_arrays(inputs::SimulationInputs) = inputs.aux_arrays
get_tspan(inputs::SimulationInputs) = inputs.tspan

get_injection_pointer(inputs::SimulationInputs) =
    get_counts(inputs)[:first_dyn_injection_pointer]
get_branches_pointer(inputs::SimulationInputs) = get_counts(inputs)[:first_dyn_branch_point]
get_n_injection_states(inputs::SimulationInputs) = get_counts(inputs)[:injection_n_states]
get_n_branches_states(inputs::SimulationInputs) = get_counts(inputs)[:branches_n_states]
get_system_state_count(inputs::SimulationInputs) = get_counts(inputs)[:total_states]
get_variable_count(inputs::SimulationInputs) = get_counts(inputs)[:total_variables]
get_device_index(inputs::SimulationInputs, device::D) where {D <: PSY.DynamicInjection} =
    get_global_index(inputs)[device.name]
get_bus_count(inputs::SimulationInputs) = get_counts(inputs)[:bus_count]
get_ω_sys(inputs::SimulationInputs) = get_global_vars(inputs)[:ω_sys]


export AbstractSolution,
       ReachSolution,
       CheckSolution,
       solve,
       project

"""
    AbstractSolution

Abstract type representing the solution of a rechability problem.
"""
abstract type AbstractSolution end

"""
    ReachSolution{S<:LazySet} <: AbstractSolution

Type that wraps a the solution of a reachability problem as a sequence of
lazy sets, and a dictionary of options.

### Fields

- `Xk`       -- the list of reachable states
- `options`  -- the dictionary of options

### Notes

If the solution has been projected in 2D, the sequence `Xk` is an array
of polygons in constraint representation. In high-dimensions this is a sequence
of cartesian product arrays of low-dimensional sets.
"""
struct ReachSolution{S<:LazySet} <: AbstractSolution
  Xk::Vector{S}
  options::Options
end
# constructor with no options
ReachSolution(Xk::Vector{S}) where {S<:LazySet} =
    ReachSolution{S}(Xk, Options())

"""
    CheckSolution

Type that wraps a the solution of a property checking problem, which is just the
answer if the property is satisfied.

### Fields

- `satisfied` -- is the property satisfied?
- `violation` -- step at which the property is violated (-1 otherwise)
- `options`   -- the dictionary of options
"""
struct CheckSolution <: AbstractSolution
  satisfied::Bool
  violation::Int
  options::Options
end
# constructor with no options
CheckSolution(satisfied::Bool, violation::Int) =
    CheckSolution(satisfied, violation, Options())

"""
    solve(system, options)  or  solve(system, :key1 => val1, [...], keyK => valK)

Solves a reachability problem s.t. the given options.
If some options are not defined, we may fall back to default values.

### Input

- `system`  -- a (discrete or continuoues) system specification
- `options` -- options for solving the problem

### Output

A solution object whose content depends on the input options.

### Notes

To see all available input options, see
`keys(Reachability.available_keywords.dict)`.
"""
function solve(system::InitialValueProblem,
               options_input::Options)::AbstractSolution

    # ==========
    # Dimensions
    # ==========
    options_input.dict[:n] = statedim(system)

    # =======
    # Options
    # =======
    # solver-specific options (adds default values for unspecified options)
    options = validate_solver_options_and_add_default_values!(options_input)

    # ===================
    # Time discretization
    # ===================
    if system isa InitialValueProblem{<:AbstractContinuousSystem}
        info("Time discretization...")
        tic()
        Δ = discretize(
            system,
            options[:δ],
            approx_model=options[:approx_model],
            pade_expm=options[:pade_expm],
            lazy_expm=options[:lazy_expm_discretize],
            lazy_sih=options[:lazy_sih]
            )
        tocc()
    else
        Δ = system
    end

    # ==============
    # Transformation
    # ==============
    if options[:coordinate_transformation] != ""
        info("Transformation...")
        tic()
        (Δ, transformation_matrix) = transform(Δ, options[:coordinate_transformation])
        tocc()
    else
        transformation_matrix = nothing
    end

    # ===================================
    # Sparse/dense/lazy matrix conversion
    # ===================================
    A = Δ.s.A
    create_new_system = false
    if !options[:lazy_expm] && options[:lazy_expm_discretize]
        # convert SparseMatrixExp to eplicit matrix
        info("Making lazy matrix exponential explicit...")
        tic()
        n = options.dict[:n]
        if options[:assume_sparse]
            B = sparse(Int[], Int[], eltype(A)[], n, n)
        else
            B = Matrix{eltype(A)}(n, n)
        end
        for i in 1:n
            B[i, :] = get_row(A, i)
        end
        A = B
        create_new_system = true
        tocc()
    end
    if options[:assume_sparse]
        if A isa SparseMatrixExp
            # ignore this case
        elseif !method_exists(sparse, Tuple{typeof(A)})
            info("`assume_sparse` option cannot be applied to a matrix of " *
                 "type $(typeof(A)) and will be ignored")
        elseif !(A isa AbstractSparseMatrix)
            # convert to sparse matrix
            A = sparse(A)
            create_new_system = true
        end
    end
    if create_new_system
        # set new matrix
        if method_exists(inputset, Tuple{typeof(Δ.s)})
            Δ = DiscreteSystem(A, Δ.x0, inputset(Δ))
        else
            Δ = DiscreteSystem(A, Δ.x0)
        end
    end

    if options[:mode] == "reach"
        # ================================
        # Input -> Output variable mapping
        # ================================
        options.dict[:inout_map] =
            inout_map_reach(options[:partition], options[:blocks], options[:n])

        # ============================
        # Reachable states computation
        # ============================
        info("Reachable States Computation...")
        tic()
        output_function = options[:project_reachset] ?
            nothing :
            options[:projection_matrix]
        Rsets = reach(
            Δ,
            options[:N];
            algorithm=options[:algorithm],
            ε_init=options[:ε_init],
            set_type_init=options[:set_type_init],
            ε_iter=options[:ε_iter],
            set_type_iter=options[:set_type_iter],
            assume_sparse=options[:assume_sparse],
            assume_homogeneous=options[:assume_homogeneous],
            lazy_X0=options[:lazy_X0],
            blocks=options[:blocks],
            partition=options[:partition],
            block_types_init=options[:block_types_init],
            block_types_iter=options[:block_types_iter],
            template_directions_init=options[:template_directions_init],
            template_directions_iter=options[:template_directions_iter],
            lazy_inputs_interval=options[:lazy_inputs_interval],
            output_function=output_function
            )
        info("- Total")
        tocc()

        # ==========
        # Projection
        # ==========
        if options[:project_reachset] || options[:projection_matrix] != nothing
            info("Projection...")
            tic()
            RsetsProj = project(Rsets, options;
                                transformation_matrix=transformation_matrix)
            tocc()
            return ReachSolution(RsetsProj, options)
        end

        return ReachSolution(Rsets, options)

    elseif options[:mode] == "check"
        # ============================================
        # Input -> Output variable mapping in property
        # ============================================
        options.dict[:property] = inout_map_property(options[:property],
            options[:partition], options[:blocks], options[:n])

        # =================
        # Property checking
        # =================
        info("Property Checking...")
        tic()
        answer = check_property(
            Δ,
            options[:N];
            algorithm=options[:algorithm],
            ε_init=options[:ε_init],
            set_type_init=options[:set_type_init],
            ε_iter=options[:ε_iter],
            set_type_iter=options[:set_type_iter],
            assume_sparse=options[:assume_sparse],
            assume_homogeneous=options[:assume_homogeneous],
            lazy_X0=options[:lazy_X0],
            blocks=options[:blocks],
            partition=options[:partition],
            block_types_init=options[:block_types_init],
            block_types_iter=options[:block_types_iter],
            template_directions_init=options[:template_directions_init],
            template_directions_iter=options[:template_directions_iter],
            eager_checking=options[:eager_checking],
            property=options[:property],
            lazy_inputs_interval=options[:lazy_inputs_interval]
            )
        info("- Total")
        tocc()

        if answer == 0
            info("The property is satisfied!")
            return CheckSolution(true, -1, options)
        else
            info("The property may be violated at index $answer," *
                " (time point $(answer * options[:δ]))!")
            return CheckSolution(false, answer, options)
        end
    else
        error("Unsupported mode.")
    end
end

solve(system::AbstractSystem, options::Pair{Symbol,<:Any}...) =
    solve(system, Options(Dict{Symbol,Any}(options)))

"""
    project(Rsets, options; [transformation_matrix])

Projects a sequence of sets according to the settings defined in the options.

### Input

- `Rsets`   -- solution of a reachability problem
- `options` -- options structure
- `transformation_matrix` -- (optional, default: nothing) matrix implementing
                              the transformation)

### Notes

A projection matrix can be given in the options structure, or passed as a
dictionary entry.
"""
function project(Rsets::Vector{<:LazySet}, options::Options;
                 transformation_matrix=nothing)
    plot_vars = copy(options[:plot_vars])
    for i in 1:length(plot_vars)
        if plot_vars[i] != 0
            plot_vars[i] = options[:inout_map][plot_vars[i]]
        end
    end
    reduced_n = sum(x -> x != 0, options[:inout_map])
    output_function = !options[:project_reachset]
    RsetsProj = project_reach(plot_vars,
                              reduced_n,
                              options[:δ],
                              Rsets,
                              options[:algorithm],
                              ε=options[:ε_proj],
                              set_type=options[:set_type_proj],
                              transformation_matrix=transformation_matrix,
                              projection_matrix=options[:projection_matrix],
                              output_function=output_function
                             )
end

project(reach_sol::AbstractSolution) = project(reach_sol.Xk, reach_sol.options)

project(Rsets::Vector{<:LazySet}, options::Pair{Symbol,<:Any}...) =
    project(Rsets, Options(Dict{Symbol,Any}(options)))


export DrillReward,
    flow, dflow!, flowdψ

# -------------------------------------------
# abstract types
# -------------------------------------------

# generic functions to access coefs
idx_drill(d) = OneTo(_nparm(d))
theta_drill(d, theta) = view(theta, idx_drill(d))

# @deprecate length(f::AbstractPayoffFunction) _nparm(f)

_nparm(m::AbstractDrillModel) = _nparm(reward(m))
coefnames(x::AbstractDrillModel) = coefnames(reward(x))

# -------------------------------------------
# Generic functions
# -------------------------------------------

NotDefinedError(m) = throw(error("Not defined for $m"))

# navigate through state space
num_choices(m::AbstractStateSpace, i...) = NotDefinedError(m)
actionspace(m::AbstractStateSpace, i...) = NotDefinedError(m)
_Dgt0(      m::AbstractStateSpace, i...) = NotDefinedError(m)

# navigate through model primitives
reward(    m::AbstractDrillModel) = NotDefinedError(m)
statespace(m::AbstractDrillModel) = NotDefinedError(m)

# From model to things about states
@inline actionspace(m::AbstractDrillModel, i...) = actionspace(statespace(m), i...)
@inline num_choices(m::AbstractDrillModel, i...) = num_choices(statespace(m), i...)
@inline _Dgt0(      m::AbstractDrillModel, i...) = _Dgt0(statespace(m), i...)

@inline _ψ(    m::AbstractDrillModel, state, s::SimulationDraw) = _Dgt0(m, state) ? _ψ2(s) : _ψ1(s)
@inline _dψdθρ(m::AbstractDrillModel, state, s::SimulationDraw) = _Dgt0(m, state) ? azero(s) : _dψ1dθρ(s)

# -------------------------------------------
# some functions to look at stuff
# -------------------------------------------

@inline actionspace(obs::ObservationDrill) = actionspace(_model(obs), _x(obs))
@inline statespace( obs::ObservationDrill) = statespace(_model(obs))
@inline reward(     obs::ObservationDrill) = reward(_model(obs))

@inline _sgnext(            obs::ObservationDrill) = _sgnext(statespace(obs), _x(obs), _y(obs))
@inline _sgnext(d::Integer, obs::ObservationDrill) = _sgnext(statespace(obs), _x(obs), d)
@inline _Dgt0(obs::ObservationDrill) = _Dgt0(statespace(obs), _x(obs))
@inline _d(   obs::ObservationDrill) = _y(obs)

@inline _ψ(    obs::ObservationDrill, ψ::Float64) = ψ
@inline _ψ(    obs::ObservationDrill, s::SimulationDraw) = _ψ(_model(obs), _x(obs), s)
@inline _dψdθρ(obs::ObservationDrill, s::SimulationDraw) = _dψdθρ(_model(obs), _x(obs), s)

@inline expires_today(            obs::ObservationDrill) = expires_today(statespace(obs), _x(obs), _y(obs))
@inline expires_today(d::Integer, obs::ObservationDrill) = expires_today(statespace(obs), _x(obs), d)

# -------------------------------------------
# Payoffs...
# -------------------------------------------

# @inline flow(        d, obs, theta, s) = flow(  reward(_model(obs)),           d,   obs, theta, s)
# @inline dflow!(grad, d, obs, theta, s) = dflow!(reward(_model(obs)), grad,     d,   obs, theta, s)
# @inline flowdψ(grad, d, obs, theta, s) = flowdψ(reward(_model(obs)), grad,     d,   obs, theta, s)
# @inline flow(           obs, theta, s) = flow(  reward(_model(obs)),       _y(obs), obs, theta, s)
# @inline dflow!(grad,    obs, theta, s) = dflow!(reward(_model(obs)), grad, _y(obs), obs, theta, s)
# @inline flowdψ(grad,    obs, theta, s) = flowdψ(reward(_model(obs)), grad, _y(obs), obs, theta, s)
# @deprecate dflowdψ(args...) flowdψ(args...)

@inline flow!(grad, obs, θ, s, dograd) = flow!(grad, reward(_model(obs)), _y(obs), obs, θ, s, dograd)

function dflow(x::AbstractPayoffFunction, d, obs, theta, s)
    length(theta) == _nparm(x) || throw(DimensionMismatch())
    grad = zero(theta)
    dflow!(grad, x, d, obs, theta, s, true)
    return grad
end

function flow(x::AbstractPayoffFunction, d, obs, theta, s)
    grad = Vector{eltype(theta)}(undef, 0)
    return flow!(grad, x, d, obs, theta, s, false)
end

@inline function full_payoff!(grad, d, obs, theta, sim, dograd)
    return flow!(grad, reward(_model(obs)), d, obs, theta, sim, dograd)
end

@inline function full_payoff(d, obs, theta, sim)
    grad = Vector{eltype(theta)}(undef, 0)
    full_payoff!(grad, d, obs, theta, sim, false)
end

# Check Derivatives
#-------------------------------------------------------------

"check gradient of reward function"
function check_flow_grad(m, d, obs, theta, sim)

    # compute analytic
    grad = ones(eltype(theta), length(theta))
    flow!(grad, m, d, obs, theta, sim, true)

    # check finite difference for dθ
    f(x) = flow!(grad, m, d, obs, x, sim, false)
    @test grad ≈ Calculus.gradient(f, theta)

    # check finite difference for dψ
    fdpsi(x) = flow!(grad, m, d, obs, theta, SimulationDrawFD(sim, x), false)
    @test flowdψ(m, d, obs, theta, sim) ≈ Calculus.derivative(fdpsi, 0.0)
end

function check_flow_grad(m, d, obs, theta)
    u, v = randn(2)
    simdraw(x) = SimulationDraw(u, v, theta_drill_ρ(m, x))
    sim = simdraw(theta)
    check_flow_grad(m,d,obs,theta,sim)
end

# @deprecate flow(x::AbstractDrillModel) reward(x)


## Simulation function
# dispatch the simulate method to be problem specific

function AutomotiveDrivingModels.simulate!{S,D,I,A,M<:AutomotiveDrivingModels.DriverModel, C<:Tuple{Vararg{Any}}}(
    ::Type{A},
    rec::EntityQueueRecord{S, D, I},
    scene::EntityFrame{S, D, I},
    env::CrosswalkEnv,
    models::Dict{I, M},
    nticks::Int,
    rng::AbstractRNG,
    n_ped::Int64,
    callbacks::C
    )

    empty!(rec)
    update!(rec, scene)

    # potential early out right off the bat
    if _run_callbacks(callbacks, rec, env, models, 0)
        return rec
    end

    actions = Array{A}(length(scene))

    # next_time = -log(rand(rng))/env.params.ped_rate
    # time_precision = Int(abs(log10(1/rec.timestep)))
    # next_time = round(next_time, time_precision)
    # println("First pedestrian Spawning time: $next_time")
    for tick in 1 : nticks
        # time = tick*rec.timestep
        # # println("time: $time")
        # if time ≈ next_time && scene.n < n_ped + 1
        #     # println("Spawning new pedestrians")
        #     next_time = time + max(-log(rand(rng))/env.params.ped_rate, rec.timestep)
        #     next_time = round(next_time, time_precision)
        #     # println("Next pedestrian at $(next_time)")
        #     new_ped = initial_pedestrian(scene, env, rng)
        #     models[new_ped.id] = ConstantPedestrian(dt = config.sim_dt)#TODO parameterized
        #     push!(scene, new_ped)
        #     actions = Array{A}(length(scene))
        # end
        if rand(config.rng) < env.params.ped_rate
            # println("Spawning new pedestrians")
            new_ped = initial_pedestrian(scene, env, config.rng)
            models[new_ped.id] = ConstantPedestrian(dt = config.sim_dt)#TODO parameterized
            push!(scene, new_ped)
            actions = Array{Any}(length(scene))
        end

        get_actions!(actions, scene, env.roadway, models)
        tick!(scene, env.roadway, actions, rec.timestep)
        update!(rec, scene)

        if _run_callbacks(callbacks, rec, env, models, tick)
            break
        end
    end
    return rec
end

function AutomotiveDrivingModels.simulate!{S,D,I,M<:AutomotiveDrivingModels.DriverModel, C<:Tuple{Vararg{Any}}}(
    rec::EntityQueueRecord{S, D, I},
    scene::EntityFrame{S, D, I},
    env::CrosswalkEnv,
    models::Dict{I, M},
    nticks::Int,
    rng::AbstractRNG,
    n_ped::Int64,
    callbacks::C
    )

    return simulate!(Any, rec, scene, env, models, nticks, rng, n_ped, callbacks)
end

function _run_callbacks{S,D,I,M<:DriverModel,C<:Tuple{Vararg{Any}}}(callbacks::C, rec::EntityQueueRecord{S,D,I}, env::CrosswalkEnv, models::Dict{I,M}, tick::Int)
    isdone = false
    for callback in callbacks
        isdone |= run_callback(callback, rec, env, models, tick)
    end
    return isdone
end

"""
    initial_pedestrian(scene::Scene, env::CrosswalkEnv, scene::Scene, rng::AbstractRNG)
Create a new pedestrian entity at its initial state
"""
function initial_pedestrian(scene::Scene, env::CrosswalkEnv, rng::AbstractRNG)
    # lateral position
    d_lat = Distributions.Uniform(env.params.roadway_length - env.params.crosswalk_width + 2, env.params.roadway_length + env.params.crosswalk_width-1)
    x0 = 0.5*rand(rng, d_lat)
    # x0 =  0.5*env.params.roadway_length

    # vertical position
    y0 = -env.params.crosswalk_length/4

    # orientation
    θ = π/2

    #velocity
    v0 = rand(rng, Distributions.Uniform(0., env.params.ped_max_speed))
    cw_roadway = Roadway([RoadSegment(2, [env.crosswalk])]);
    ped_initialstate = VehicleState(VecSE2(x0, y0, θ), env.crosswalk, cw_roadway, v0);
    # ped_initialstate = VehicleState(VecSE2(x0, y0, θ), env.crosswalk, env.roadway, v0)

    # new id, increment last id
    id = scene.n + 2

    return Vehicle(ped_initialstate,  VehicleDef(AgentClass.PEDESTRIAN, 1.0, 1.0), id)
end

function initial_scene(models::Dict{Int64, DriverModel}, env::CrosswalkEnv, config::EvalConfig)
    scene = Scene()
    n_ped = rand(config.rng, 0:10)
    n_ped = 0
    ego = models[1].ego.state
    pomdp = models[1].updater.problem
    b0 = initialstate(pomdp, ego)
    for i = 1:n_ped
        ped_state = rand(config.rng, b0).ped
        d_lat = Distributions.Uniform(env.params.roadway_length - env.params.crosswalk_width + 2, env.params.roadway_length + env.params.crosswalk_width-1)
        x0 = 0.5*rand(config.rng, d_lat)
        y0 = ped_state.posG.y
        cw_roadway = Roadway([RoadSegment(2, [env.crosswalk])])
        ped_state = VehicleState(VecSE2(x0, y0, ped_state.posG.θ), env.crosswalk, cw_roadway, ped_state.v)

        if !off_the_grid(pomdp, ped_state)
            id = scene.n + 2
            new_ped = Vehicle(ped_state, PEDESTRIAN_DEF, id)
            models[new_ped.id] = ConstantPedestrian(dt = config.sim_dt)
            push!(scene, new_ped)
        end
    end
    # println("Starting with ", scene.n, " pedestrians on the scene")
    return scene
end

# function initial_scene(models::Dict{Int64, DriverModel}, env::CrosswalkEnv,
#                        config::EvalConfig)
#
#     nticks = rand(config.rng, 0:200)
#     rec = SceneRecord(nticks+1, config.sim_dt)
#     scene = Scene()
#     empty!(rec)
#     update!(rec, scene)
#     actions = Array{Any}(length(scene))
#
#     for tick in 1 : nticks
#         time = tick*rec.timestep
#         # println("time: $time")
#         if rand(config.rng) < env.params.ped_rate
#             # println("Spawning new pedestrians")
#             new_ped = initial_pedestrian(scene, env, config.rng)
#             models[new_ped.id] = ConstantPedestrian(dt = config.sim_dt)#TODO parameterized
#             push!(scene, new_ped)
#             actions = Array{Any}(length(scene))
#         end
#
#         get_actions!(actions, scene, env.roadway, models)
#         tick!(scene, env.roadway, actions, rec.timestep)
#         update!(rec, scene)
#     end
#     @assert findfirst(EGO_ID, scene) == nothing
#     return rec[0]
# end


module TimeSteppers

export
    AdamsBashforthTimeStepper,
    time_step!,
    compute_w_from_continuity!

using Oceananigans: device

using GPUifyLoops: @launch, @loop, @unroll

import Oceananigans: TimeStepper

using Oceananigans: AbstractGrid, Model, Tendencies, tracernames,
                    @hascuda, CPU, GPU, launch_config, datatuples, datatuple,
                    @loop_xyz,

                    buoyancy_perturbation,
                    x_f_cross_U, y_f_cross_U, z_f_cross_U,

                    fill_halo_regions!, apply_z_bcs!, apply_y_bcs!,

                    time_to_run

using Oceananigans.SurfaceWaves: x_curl_Uˢ_cross_U, y_curl_Uˢ_cross_U, z_curl_Uˢ_cross_U,
                                 ∂t_uˢ, ∂t_vˢ, ∂t_wˢ

@hascuda using CUDAnative, CUDAdrv, CuArrays

using Oceananigans.Operators
using Oceananigans.Solvers
using Oceananigans.Diagnostics
using Oceananigans.OutputWriters

using Oceananigans.Solvers: solve_poisson_3d!, PoissonBCs, PPN, PNN
using Oceananigans.Diagnostics: run_diagnostic
using Oceananigans.OutputWriters: write_output

using ..TurbulenceClosures: calculate_diffusivities!, ∂ⱼ_2ν_Σ₁ⱼ, ∂ⱼ_2ν_Σ₂ⱼ, ∂ⱼ_2ν_Σ₃ⱼ, ∇_κ_∇c

"""
    TimeStepper(name, args...)

Returns a timestepper with name `name`, instantiated with `args...`.

Example
=======

julia> stepper = TimeStepper(:AdamsBashforth, Float64, CPU(), grid, tracernames)
"""
function TimeStepper(name::Symbol, args...)
    fullname = Symbol(name, :TimeStepper)
    return eval(Expr(:call, fullname, args...))
end

# Fallback
TimeStepper(stepper, args...) = stepper

"""Returns the arguments passed to boundary conditions functions."""
boundary_condition_function_arguments(model) =
    (model.clock.time, model.clock.iteration, datatuple(model.velocities),
     datatuple(model.tracers), model.parameters)

#####
##### Time-stepping kernels/function that are independent of the TimeStepper
#####

"""
    time_step!(model; Nt, Δt, kwargs...)

Step forward `model` `Nt` time steps with step size `Δt`.

The kwargs are passed to the `time_step!` function specific to `model.timestepper`.
"""
time_step!(model; Nt, Δt, kwargs...) = time_step!(model, Nt, Δt; kwargs...)

function time_step!(model, Nt, Δt; kwargs...)

    if model.clock.iteration == 0
        [ run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ write_output(model, out)    for out  in values(model.output_writers) ]
    end

    for n in 1:Nt
        time_step!(model, Δt; kwargs...)

        [ time_to_run(model.clock, diag) && run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ time_to_run(model.clock, out) && write_output(model, out) for out in values(model.output_writers) ]
    end

    return nothing
end

"""
    time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)

Perform precomputations necessary for an explicit timestep or substep.
"""
function time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)

    fill_halo_regions!(merge(velocities, tracers), model.boundary_conditions.solution, model.architecture,
                       model.grid, boundary_condition_function_arguments(model)...)

    calculate_diffusivities!(diffusivities, model.architecture, model.grid, model.closure, model.buoyancy,
                             velocities, tracers)

    # Diffusivities share bcs with pressure:
    fill_halo_regions!(diffusivities, model.boundary_conditions.pressure, model.architecture, model.grid)

    @launch(device(model.architecture), config=launch_config(model.grid, :xy),
            update_hydrostatic_pressure!(pressures.pHY′, model.grid, model.buoyancy, tracers))

    fill_halo_regions!(pressures.pHY′, model.boundary_conditions.pressure, model.architecture, model.grid)

    return nothing
end

"""
    calculate_tendencies!(diffusivities, pressures, velocities, tracers, model)

Calculate the interior and boundary contributions to tendency terms without the
contribution from non-hydrostatic pressure.
"""
function calculate_tendencies!(tendencies, velocities, tracers, pressures, diffusivities, model)

    calculate_interior_source_terms!(tendencies, model.architecture, model.grid, model.coriolis, model.buoyancy,
                                     model.surface_waves, model.closure, velocities, tracers, pressures.pHY′,
                                     diffusivities, model.forcing, model.parameters, model.clock.time)

    calculate_boundary_source_terms!(tendencies, model.boundary_conditions.solution, model.architecture,
                                     model.grid, boundary_condition_function_arguments(model)...)

    return nothing
end

"""
    calculate_pressure_correction!(nonhydrostatic_pressure, Δt, tendencies, velocities, model)

Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `velocities`, and step size `Δt`.
"""
function calculate_pressure_correction!(nonhydrostatic_pressure, Δt, tendencies, velocities, model)
    velocity_tendencies = (u=tendencies.u, v=tendencies.v, w=tendencies.w)

    velocity_tendency_boundary_conditions = (u=model.boundary_conditions.tendency.u,
                                             v=model.boundary_conditions.tendency.v,
                                             w=model.boundary_conditions.tendency.w)

    fill_halo_regions!(velocity_tendencies, velocity_tendency_boundary_conditions, model.architecture, model.grid)

    @launch(device(model.architecture), config=launch_config(model.grid, :xyz),
            calculate_poisson_right_hand_side!(model.poisson_solver.storage, model.architecture, model.grid,
                                               model.poisson_solver.bcs, velocities, tendencies, Δt))

    solve_for_pressure!(nonhydrostatic_pressure, model.architecture, model.grid, model.poisson_solver,
                        model.poisson_solver.storage)

    fill_halo_regions!(nonhydrostatic_pressure, model.boundary_conditions.pressure, model.architecture, model.grid)

    return nothing
end

calculate_pressure_correction!(::Nothing, args...) = nothing

"""
    calculate_explicit_substep!(tendencies, velocities, tracers, pressures, diffusivities, model)

Calculate the initial and explicit substep of the two-step fractional step method with pressure correction.
"""
function calculate_explicit_substep!(tendencies, velocities, tracers, pressures, diffusivities, model)
    time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)
    calculate_tendencies!(tendencies, velocities, tracers, pressures, diffusivities, model)
    return nothing
end

"""
    complete_pressure_correction_step!(velocities, Δt, tracers, pressures, tendencies, model)

After calculating the pressure correction, complete the pressure correction step by updating
the velocity and tracer fields.
"""
function complete_pressure_correction_step!(velocities, Δt, tracers, pressures, tendencies, model)
    update_solution!(velocities, tracers, model.architecture, model.grid, Δt, tendencies, pressures.pNHS)

    velocity_boundary_conditions = (u=model.boundary_conditions.solution.u,
                                    v=model.boundary_conditions.solution.v,
                                    w=model.boundary_conditions.solution.w)

    # Recompute vertical velocity w from continuity equation to ensure incompressibility
    fill_halo_regions!(velocities, velocity_boundary_conditions, model.architecture, model.grid,
                       boundary_condition_function_arguments(model)...)

    compute_w_from_continuity!(model)

    return nothing
end

include("kernels.jl")
include("adams_bashforth.jl")

end # module


# KKT system updates -------------------------------------------------------
# Set diagonal
function set_aug_diagonal!(kkt::AbstractKKTSystem, ips::InteriorPointSolver)
    kkt.pr_diag .= ips.zl./(ips.x.-ips.xl) .+ ips.zu./(ips.xu.-ips.x)
    fill!(kkt.du_diag, 0.0)
end
function set_aug_diagonal!(kkt::SparseUnreducedKKTSystem, ips::InteriorPointSolver)
    kkt.pr_diag .= 0.0
    kkt.du_diag .= 0.0
    kkt.l_lower .= .-sqrt.(ips.zl_r)
    kkt.u_lower .= .-sqrt.(ips.zu_r)
    kkt.l_diag  .= ips.xl_r .- ips.x_lr
    kkt.u_diag  .= ips.x_ur .- ips.xu_r
end

# Robust restoration
function set_aug_RR!(kkt::AbstractKKTSystem, ips::InteriorPointSolver, RR::RobustRestorer)
    kkt.pr_diag .= ips.zl./(ips.x.-ips.xl) .+ ips.zu./(ips.xu.-ips.x) .+ RR.zeta.*RR.D_R.^2
    kkt.du_diag .= .-RR.pp./RR.zp .- RR.nn./RR.zn
end
function set_aug_RR!(kkt::SparseUnreducedKKTSystem, ips::InteriorPointSolver, RR::RobustRestorer)
    kkt.pr_diag.= RR.zeta.*RR.D_R.^2
    kkt.du_diag.= .-RR.pp./RR.zp.-RR.nn./RR.zn
    kkt.l_lower.=.-sqrt.(ips.zl_r)
    kkt.u_lower.=.-sqrt.(ips.zu_r)
    kkt.l_diag .= ips.xl_r .- ips.x_lr
    kkt.u_diag .= ips.x_ur .- ips.xu_r
end

# Set RHS
function set_aug_rhs!(ips::InteriorPointSolver, kkt::AbstractKKTSystem, c)
    primal(ips.p) .= .-ips.f.+ips.mu./(ips.x.-ips.xl).-ips.mu./(ips.xu.-ips.x).-ips.jacl
    dual(ips.p)   .= .-c
end

function set_aug_rhs!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem, c)
    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu.-ips.jacl
    dual(ips.p) .= .-c
    dual_lb(ips.p) .= (ips.xl_r-ips.x_lr).*kkt.l_lower .+ ips.mu./kkt.l_lower
    dual_ub(ips.p) .= (ips.xu_r-ips.x_ur).*kkt.u_lower .- ips.mu./kkt.u_lower
end

function set_aug_rhs_ifr!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem,c)
    primal(ips._w1) .= 0.0
    dual(ips._w1) .= .-c
    dual_lb(ips._w1) .= 0.0
    dual_ub(ips._w1) .= 0.0
end

# Set RHS RR
function set_aug_rhs_RR!(
    ips::InteriorPointSolver, kkt::AbstractKKTSystem, RR::RobustRestorer, rho,
)
    primal(ips.p) .= .-RR.f_R.-ips.jacl.+RR.mu_R./(ips.x.-ips.xl).-RR.mu_R./(ips.xu.-ips.x)
    dual(ips.p) .= .-ips.c.+RR.pp.-RR.nn.+(RR.mu_R.-(rho.-ips.l).*RR.pp)./RR.zp.-(RR.mu_R.-(rho.+ips.l).*RR.nn)./RR.zn
end

# Finish
function finish_aug_solve!(ips::InteriorPointSolver, kkt::AbstractKKTSystem, mu)
    dual_lb(ips.d) .= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
    dual_ub(ips.d) .= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
end

function finish_aug_solve!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem, mu)
    dual_lb(ips.d) .*= .-kkt.l_lower
    dual_ub(ips.d) .*= kkt.u_lower
    dual_lb(ips.d) .= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
    dual_ub(ips.d) .= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
end

# Initial
function set_initial_rhs!(ips::InteriorPointSolver, kkt::AbstractKKTSystem)
    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu
    dual(ips.p) .= 0.0
end
function set_initial_rhs!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem)
    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu
    dual(ips.p) .= 0.0
    dual_lb(ips.p) .= 0.0
    dual_ub(ips.p) .= 0.0
end

# Set ifr
function set_aug_rhs_ifr!(ips::InteriorPointSolver, kkt::AbstractKKTSystem)
    primal(ips._w1) .= 0.0
    dual(ips._w1) .= .-ips.c
end

# Finish RR
function finish_aug_solve_RR!(dpp,dnn,dzp,dzn,l,dl,pp,nn,zp,zn,mu_R,rho)
    dpp .= (mu_R.+pp.*dl.-(rho.-l).*pp)./zp
    dnn .= (mu_R.-nn.*dl.-(rho.+l).*nn)./zn
    dzp .= (mu_R.-zp.*dpp)./pp.-zp
    dzn .= (mu_R.-zn.*dnn)./nn.-zn
end

# Kernel functions ---------------------------------------------------------
is_valid(val::Real) = !(isnan(val) || isinf(val))
function is_valid(vec::AbstractArray)
    @inbounds for i=1:length(vec)
        is_valid(vec[i]) || return false
    end
    return true
end
is_valid(args...) = all(is_valid(arg) for arg in args)

function get_varphi(obj_val,x_lr,xl_r,xu_r,x_ur,mu)
    varphi = obj_val
    @simd for i=1:length(x_lr)
        @inbounds xll = x_lr[i]-xl_r[i]
        xll < 0 && return Inf
        varphi -= mu*log(xll)
    end
    @simd for i=1:length(x_ur)
        @inbounds xuu = xu_r[i]-x_ur[i]
        xuu < 0 && return Inf
        varphi -= mu*log(xuu)
    end
    return varphi
end
get_inf_pr(c) = norm(c,Inf)
function get_inf_du(f,zl,zu,jacl,sd)
    inf_du = 0.
    @simd for i=1:length(f)
        @inbounds inf_du = max(inf_du,abs(f[i]-zl[i]+zu[i]+jacl[i]))
    end
    return inf_du/sd
end
function get_inf_compl(x_lr,xl_r,zl_r,xu_r,x_ur,zu_r,mu,sc)
    inf_compl = 0.
    @simd for i=1:length(x_lr)
        @inbounds inf_compl = max(inf_compl,abs((x_lr[i]-xl_r[i])*zl_r[i]-mu))
    end
    @simd for i=1:length(x_ur)
        @inbounds inf_compl = max(inf_compl,abs((xu_r[i]-x_ur[i])*zu_r[i]-mu))
    end
    return inf_compl/sc
end
function get_varphi_d(f,x,xl,xu,dx,mu)
    varphi_d = 0.
    @simd for i=1:length(f)
        @inbounds varphi_d += (f[i] - mu/(x[i]-xl[i]) + mu/(xu[i]-x[i])) *dx[i]
    end
    return varphi_d
end
function get_alpha_max(x,xl,xu,dx,tau)
    alpha_max = 1.
    @simd for i=1:length(x)
        @inbounds dx[i]<0 && (alpha_max=min(alpha_max,(-x[i]+xl[i])*tau/dx[i]))
        @inbounds dx[i]>0 && (alpha_max=min(alpha_max,(-x[i]+xu[i])*tau/dx[i]))
    end
    return alpha_max
end
function get_alpha_z(zl_r,zu_r,dzl,dzu,tau)
    alpha_z = 1.
    @simd for i=1:length(zl_r)
        @inbounds dzl[i]<0 && (alpha_z=min(alpha_z,-zl_r[i]*tau/dzl[i]))
     end
    @simd for i=1:length(zu_r)
        @inbounds dzu[i]<0 && (alpha_z=min(alpha_z,-zu_r[i]*tau/dzu[i]))
    end
    return alpha_z
end
function get_obj_val_R(p,n,D_R,x,x_ref,rho,zeta)
    obj_val_R = 0.
    @simd for i=1:length(p)
        @inbounds obj_val_R += rho*(p[i]+n[i]) .+ zeta/2*D_R[i]^2*(x[i]-x_ref[i])^2
    end
    return obj_val_R
end
get_theta(c) = norm(c,1)
function get_theta_R(c,p,n)
    theta_R = 0.
    @simd for i=1:length(c)
        @inbounds theta_R += abs(c[i]-p[i]+n[i])
    end
    return theta_R
end
function get_inf_pr_R(c,p,n)
    inf_pr_R = 0.
    @simd for i=1:length(c)
        @inbounds inf_pr_R = max(inf_pr_R,abs(c[i]-p[i]+n[i]))
    end
    return inf_pr_R
end
function get_inf_du_R(f_R,l,zl,zu,jacl,zp,zn,rho,sd)
    inf_du_R = 0.
    @simd for i=1:length(zl)
        @inbounds inf_du_R = max(inf_du_R,abs(f_R[i]-zl[i]+zu[i]+jacl[i]))
    end
    @simd for i=1:length(zp)
        @inbounds inf_du_R = max(inf_du_R,abs(rho-l[i]-zp[i]))
    end
    @simd for i=1:length(zn)
        @inbounds inf_du_R = max(inf_du_R,abs(rho+l[i]-zn[i]))
    end
    return inf_du_R/sd
end
function get_inf_compl_R(x_lr,xl_r,zl_r,xu_r,x_ur,zu_r,pp,zp,nn,zn,mu_R,sc)
    inf_compl_R = 0.
    @simd for i=1:length(x_lr)
        @inbounds inf_compl_R = max(inf_compl_R,abs((x_lr[i]-xl_r[i])*zl_r[i]-mu_R))
    end
    @simd for i=1:length(xu_r)
        @inbounds inf_compl_R = max(inf_compl_R,abs((xu_r[i]-x_ur[i])*zu_r[i]-mu_R))
    end
    @simd for i=1:length(pp)
        @inbounds inf_compl_R = max(inf_compl_R,abs(pp[i]*zp[i]-mu_R))
    end
    @simd for i=1:length(nn)
        @inbounds inf_compl_R = max(inf_compl_R,abs(nn[i]*zn[i]-mu_R))
    end
    return inf_compl_R/sc
end
function get_alpha_max_R(x,xl,xu,dx,pp,dpp,nn,dnn,tau_R)
    alpha_max_R = 1.
    @simd for i=1:length(x)
        @inbounds dx[i]<0 && (alpha_max_R=min(alpha_max_R,(-x[i]+xl[i])*tau_R/dx[i]))
        @inbounds dx[i]>0 && (alpha_max_R=min(alpha_max_R,(-x[i]+xu[i])*tau_R/dx[i]))
    end
    @simd for i=1:length(pp)
        @inbounds dpp[i]<0 && (alpha_max_R=min(alpha_max_R,-pp[i]*tau_R/dpp[i]))
    end
    @simd for i=1:length(nn)
        @inbounds dnn[i]<0 && (alpha_max_R=min(alpha_max_R,-nn[i]*tau_R/dnn[i]))
    end
    return alpha_max_R
end
function get_alpha_z_R(zl_r,zu_r,dzl,dzu,zp,dzp,zn,dzn,tau_R)
    alpha_z_R = 1.
    @simd for i=1:length(zl_r)
        @inbounds dzl[i]<0 && (alpha_z_R=min(alpha_z_R,-zl_r[i]*tau_R/dzl[i]))
    end
    @simd for i=1:length(zu_r)
        @inbounds dzu[i]<0 && (alpha_z_R=min(alpha_z_R,-zu_r[i]*tau_R/dzu[i]))
    end
    @simd for i=1:length(zp)
        @inbounds dzp[i]<0 && (alpha_z_R=min(alpha_z_R,-zp[i]*tau_R/dzp[i]))
    end
    @simd for i=1:length(zn)
        @inbounds dzn[i]<0 && (alpha_z_R=min(alpha_z_R,-zn[i]*tau_R/dzn[i]))
    end
    return alpha_z_R
end
function get_varphi_R(obj_val,x_lr,xl_r,xu_r,x_ur,pp,nn,mu_R)
    varphi_R = obj_val
    @simd for i=1:length(x_lr)
        @inbounds xll = x_lr[i]-xl_r[i]
        xll < 0 && return Inf
        varphi_R -= mu_R*log(xll)
    end
    @simd for i=1:length(x_ur)
        @inbounds xuu = xu_r[i]-x_ur[i]
        xuu < 0 && return Inf
        varphi_R -= mu_R*log(xuu)
    end
    @simd for i=1:length(pp)
        @inbounds pp[i] < 0 && return Inf
        @inbounds varphi_R -= mu_R*log(pp[i])
    end
    @simd for i=1:length(pp)
        @inbounds nn[i] < 0 && return Inf
        @inbounds varphi_R -= mu_R*log(nn[i])
    end
    return varphi_R
end
function get_F(c,f,zl,zu,jacl,x_lr,xl_r,zl_r,xu_r,x_ur,zu_r,mu)
    F = 0.
    for i=1:length(c)
        @inbounds F = max(F,c[i])
    end
    for i=1:length(f)
        @inbounds F = max(F,f[i]-zl[i]+zu[i]+jacl[i])
    end
    for i=1:length(x_lr)
        x_lr[i] >= xl_r[i] || return Inf
        zl_r[i] >= 0       || return Inf
        @inbounds F = max(F,(x_lr[i]-xl_r[i])*zl_r[i]-mu)
    end
    for i=1:length(x_ur)
        xu_r[i] >= x_ur[i] || return Inf
        zu_r[i] >= 0       || return Inf
        @inbounds F = max(F,(xu_r[i]-xu_r[i])*zu_r[i]-mu)
    end
    return F
end
function get_varphi_d_R(f_R,x,xl,xu,dx,pp,nn,dpp,dnn,mu_R,rho)
    varphi_d = 0.
    @simd for i=1:length(x)
        @inbounds varphi_d += (f_R[i] - mu_R/(x[i]-xl[i]) + mu_R/(xu[i]-x[i])) *dx[i]
    end
    @simd for i=1:length(pp)
        @inbounds varphi_d += (rho - mu_R/pp[i]) *dpp[i]
    end
    @simd for i=1:length(nn)
        @inbounds varphi_d += (rho - mu_R/nn[i]) *dnn[i]
    end
    return varphi_d
end
function initialize_variables!(x,xl,xu,bound_push,bound_fac)
    @inbounds @simd for i=1:length(x)
        if xl[i]!=-Inf && xu[i]!=Inf
            x[i]=min(xu[i]-min(bound_push*max(1,abs(xu[i])),bound_fac*(xu[i]-xl[i])),
                     max(xl[i]+min(bound_push*max(1,abs(xl[i])),bound_fac*(xu[i]-xl[i])),x[i]))
        elseif xl[i]!=-Inf && xu[i]==Inf
            x[i]=max(xl[i]+bound_push*max(1,abs(xl[i])),x[i])
        elseif xl[i]==-Inf && xu[i]!=Inf
            x[i]=min(xu[i]-bound_push*max(1,abs(xu[i])),x[i])
        end
    end
end

function adjust_boundary!(x_lr,xl_r,x_ur,xu_r,mu)
    adjusted = 0
    c1 = eps(Float64)*mu
    c2= eps(Float64)^(3/4)
    @simd for i=1:length(xl_r)
        @inbounds x_lr[i]-xl_r[i] < c1 && (xl_r[i] -= c2*max(1,abs(x_lr[i]));adjusted+=1)
    end
    @simd for i=1:length(xu_r)
        @inbounds xu_r[i]-x_ur[i] < c1 && (xu_r[i] += c2*max(1,abs(x_ur[i]));adjusted+=1)
    end
    return adjusted
end
function get_rel_search_norm(x,dx)
    rel_search_norm = 0.
    @simd for i=1:length(x)
        @inbounds rel_search_norm = max(rel_search_norm,abs(dx[i])/(1. +abs(x[i])))
    end
    return rel_search_norm
end

# Utility functions
get_sd(l,zl_r,zu_r,s_max) =
    max(s_max,(norm(l,1)+norm(zl_r,1)+norm(zu_r,1)) / max(1,(length(l)+length(zl_r)+length(zu_r))))/s_max
get_sc(zl_r,zu_r,s_max) =
    max(s_max,(norm(zl_r,1)+norm(zu_r,1)) / max(1,length(zl_r)+length(zu_r)))/s_max
get_mu(mu,mu_min,mu_linear_decrease_factor,mu_superlinear_decrease_power,tol) =
    max(mu_min, max(tol/10,min(mu_linear_decrease_factor*mu,mu^mu_superlinear_decrease_power)))
get_tau(mu,tau_min)=max(tau_min,1-mu)
function get_alpha_min(theta,varphi_d,theta_min,gamma_theta,gamma_phi,alpha_min_frac,del,s_theta,s_phi)
    if varphi_d<0
        if theta<=theta_min
            return alpha_min_frac*min(
                gamma_theta,gamma_phi*theta/(-varphi_d),
                del*theta^s_theta/(-varphi_d)^s_phi)
        else
            return alpha_min_frac*min(
                gamma_theta,gamma_phi*theta/(-varphi_d))
        end
    else
        return alpha_min_frac*gamma_theta
    end
end
is_switching(varphi_d,alpha,s_phi,del,theta,s_theta) = varphi_d < 0 && alpha*(-varphi_d)^s_phi > del*theta^s_theta
is_armijo(varphi_trial,varphi,eta_phi,alpha,varphi_d) = varphi_trial <= varphi + eta_phi*alpha*varphi_d
is_sufficient_progress(theta_trial,theta,gamma_theta,varphi_trial,varphi,gamma_phi,has_constraints) =
    (has_constraints && ((theta_trial<=(1-gamma_theta)*theta+10*eps(Float64)*abs(theta))) ||
    ((varphi_trial<=varphi-gamma_phi*theta +10*eps(Float64)*abs(varphi))))
augment_filter!(filter,theta,varphi,gamma_theta) = push!(filter,((1-gamma_theta)*theta,varphi-gamma_theta*theta))
function is_filter_acceptable(filter,theta,varphi)
    !isnan(theta) || return false
    !isinf(theta) || return false
    !isnan(varphi) || return false
    !isinf(varphi) || return false

    for (theta_F,varphi_F) in filter
        theta <= theta_F || varphi <= varphi_F || return false
    end
    return true
end
is_barr_obj_rapid_increase(varphi,varphi_trial,obj_max_inc) =
    varphi_trial >= varphi && log(10,varphi_trial-varphi) > obj_max_inc + max(1.,log(10,abs(varphi)))
reset_bound_dual!(z,x,mu,kappa_sigma) = (z.=max.(min.(z,kappa_sigma.*mu./x),mu/kappa_sigma./x))
reset_bound_dual!(z,x1,x2,mu,kappa_sigma) = (z.=max.(min.(z,(kappa_sigma*mu)./(x1.-x2)),(mu/kappa_sigma)./(x1.-x2)))
function get_ftype(filter,theta,theta_trial,varphi,varphi_trial,switching_condition,armijo_condition,
                   theta_min,obj_max_inc,gamma_theta,gamma_phi,has_constraints)
    is_filter_acceptable(filter,theta_trial,varphi_trial) || return " "
    !is_barr_obj_rapid_increase(varphi,varphi_trial,obj_max_inc) || return " "

    if theta <= theta_min && switching_condition
        armijo_condition && return "f"
    else
        is_sufficient_progress(
            theta_trial,theta,gamma_theta,varphi_trial,varphi,gamma_phi,has_constraints) && return "h"
    end

    return " "
end

# fixed variable treatment ----------------------------------------------------
function _get_fixed_variable_index(
    mat::SparseMatrixCSC{Tv,Ti1}, ind_fixed::Vector{Ti2}) where {Tv,Ti1,Ti2}

    fixed_aug_index = Int[]
    for i in ind_fixed
        append!(fixed_aug_index,append!(collect(mat.colptr[i]+1:mat.colptr[i+1]-1)))
    end
    append!(fixed_aug_index,setdiff!(Base._findin(mat.rowval,ind_fixed),mat.colptr))

    return fixed_aug_index
end
fixed_variable_treatment_vec!(vec,ind_fixed) = (vec[ind_fixed] .= 0.)
function fixed_variable_treatment_z!(zl,zu,f,jacl,ind_fixed)
    @simd for i in ind_fixed
        z = f[i]+jacl[i]
        z >=0 ? (zl[i] = z; zu[i] = 0.) : (zl[i] = 0.; zu[i] = -z)
    end
end

function dual_inf_perturbation!(px,ind_llb,ind_uub,mu,kappa_d)
    @simd for i in ind_llb
        @inbounds px[i] -= mu*kappa_d
    end
    @simd for i in ind_uub
        @inbounds px[i] += mu*kappa_d
    end
end


# ------------ Trace typing rules ------------ #

# Define an interpreter - allows multiple dispatch to handle multiple abstract interpretations.
struct TraceTypingInterpreter <: InterpretationContext end

# Fallback.
absint(ctx::TraceTypingInterpreter, args...) = Union{}
absint(ctx::TraceTypingInterpreter, trace, args...) = Missing

# Distributions.
absint(ctx::TraceTypingInterpreter, x::D, args...) where D <: Distributions.Distribution =  x
function absint(ctx::TraceTypingInterpreter, ::Type{D}, args...) where D <: Distributions.Distribution
    all(map(args) do a
            _lit(a)
        end) || return D
    return D(args...)
end

# Note - the usage of the @abstract macro below is strictly equivalent to the longform definitions for absint with ctx::T above.

# Learnables.
@abstract TraceTypingInterpreter learnable(::Symbol) = Any

# Calls to trace.
@abstract TraceTypingInterpreter trace(::Symbol, ::Type{DiscreteUniform}, args...) = Integers{1}
@abstract TraceTypingInterpreter trace(::Symbol, d::DiscreteUniform, args...) = DiscreteInterval{d.a, d.b}
@abstract TraceTypingInterpreter trace(::Symbol, ::Type{Uniform}, args...) = Reals{1}
@abstract TraceTypingInterpreter trace(::Symbol, d::Uniform, args...) = RealInterval(d.a, d.b)
@abstract TraceTypingInterpreter trace(::Symbol, ::Type{Gamma}, args...) = PositiveReals{1}
@abstract TraceTypingInterpreter trace(::Symbol, ::Gamma, args...) = PositiveReals{1}
@abstract TraceTypingInterpreter trace(::Symbol, ::Type{Normal}, args...) = Reals{1}
@abstract TraceTypingInterpreter trace(::Symbol, ::Normal, args...) = Reals{1}
@abstract TraceTypingInterpreter trace(::Symbol, ::Type{Bernoulli}, args...) = Discrete{2}
@abstract TraceTypingInterpreter trace(::Symbol, ::Bernoulli, args...) = Discrete{2}
@abstract TraceTypingInterpreter trace(::Symbol, jfn::JFunction, args...) = get_trace_type(jfn)

# Combinators.
@abstract TraceTypingInterpreter function trace(::Symbol, u::Gen.Unfold, args...)
    tt = get_trace_type(u.kernel)
    List{tt}
end
@abstract TraceTypingInterpreter function trace(::Symbol, u::Gen.Map, args...)
    List{get_trace_type(u.value.kernel)}
end
@abstract TraceTypingInterpreter function trace(::Symbol, u::Gen.Switch, args...)
    Sum{Tuple(map(u.value.branches) do br
                  get_trace_type(br)
              end)}
end


using Gamess
using Test

import Gamess: Configuration, substitutions

load_file(name) = Gamess.load(joinpath(dirname(@__FILE__), "data", "$(name).out"))

@testset "Gamess.jl" begin
    @testset "Water CIS" begin
        water = load_file("water")
        atoms = water.molecule.atoms
        @test size(atoms) == (3, 5)
        @test atoms.Atom == ["OXYGEN", "HYDROGEN", "HYDROGEN"]

        ci = water.ci
        states = ci.excited_states
        @test size(states) == (100,3)

        dipoles = ci.length_gauge_dipole
        @test size(dipoles) == (2601,10)
        @test dipoles[1, :z] ≈ 0.958273779741262
        @test dipoles[2599,:z] ≈ 0.044333
        @test dipoles[2599,:A] ≈ -5.4312e-41
        @test dipoles[2599,:B] ≈ 797650.0

        expansions = ci.state_expansions
        @test length(expansions) == 100

        e = first(expansions)
        reference = e.reference
        N = 13
        @test length(reference.α) == length(reference.β) == N
        @test reference == Configuration(vcat(trues(5), falses(N-5)), vcat(trues(5), falses(N-5)))

        @test substitutions(reference, e.configurations[1]) == ([5 => 6], [])
        @test substitutions(reference, e.configurations[2]) == ([5 => 9], [])
        @test substitutions(reference, e.configurations[3]) == ([],[5 => 6])
        @test substitutions(reference, e.configurations[4]) == ([],[5 => 9])
    end

    @testset "Neon CIS" begin
        @testset "Ne I" begin
            neon = load_file("ne-cis")
            atoms = neon.molecule.atoms
            @test size(atoms) == (1, 5)
            @test atoms.Atom == ["NE"]

            ci = neon.ci
            states = ci.excited_states
            @test size(states) == (100,3)

            dipoles = ci.length_gauge_dipole
            @test size(dipoles) == (2603,10)
            @test dipoles[4, :z] ≈ -0.250121

            expansions = ci.state_expansions
            @test length(expansions) == 100

            e = first(expansions)
            reference = e.reference
            N = 121
            @test length(reference.α) == length(reference.β) == N
            @test reference == Configuration(vcat(trues(5), falses(N-5)), vcat(trues(5), falses(N-5)))

            @test substitutions(reference, e.configurations[1]) == ([4 => 10], [])
            @test substitutions(reference, e.configurations[2]) == ([4 => 19], [])
            @test substitutions(reference, e.configurations[13]) == ([],[4 => 10])
            @test substitutions(reference, e.configurations[14]) == ([],[4 => 19])
        end

        @testset "Ne II" begin
            neon = load_file("ne+-guga")
            atoms = neon.molecule.atoms
            @test size(atoms) == (1, 5)
            @test atoms.Atom == ["NE"]

            ci = neon.ci
            states = ci.excited_states
            @test size(states) == (5,2)

            dipoles = ci.length_gauge_dipole
            @test size(dipoles) == (15,10)
            @test dipoles[11, :z] ≈ -0.521326

            dipoles = ci.velocity_gauge_dipole
            @test size(dipoles) == (10,8)
        end
    end
end


struct Font
    surface::Cairo.CairoSurface
    charDict::Dict{Char, Tuple{Int, Int, Int, Int}}
    charSpacing::Int
    lineSpacing::Int

    Font(filename::String, fontString::String, charSpacing::Int=1, lineSpacing::Int=1) = new(open(Cairo.read_from_png, filename), generateFontDict(fontString), charSpacing, lineSpacing)
end

const stringLineAlignOffsets = Dict{Symbol, Function}(
    :left => (widest, width) -> 0,
    :center => (widest, width) -> round(Int, (widest - width) / 2),
    :right => (widest, width) -> widest - width,
)

function getQuadTuple(row::Int, col::Int, width::Int=3, height::Int=5)
    return (8 * row - 8, 8 * col - 8, width, height)
end

# Use first occurance of a symbol as the target quad
function generateFontDict(fontString::String)
    fontMap = Dict{Char, Tuple{Int, Int, Int, Int}}()

    for (i, row) in enumerate(split(fontString, "\n"))
        for (j, c) in enumerate(row)
            get!(fontMap, c) do
                getQuadTuple(j, i, 3, 5)
            end
        end
    end

    return fontMap
end

# Multiline string
function lineWidth(font::Font, s::AbstractString, replaceMissing::Char=' ')
    return isempty(s) ? 0 : sum([get(font.charDict, c, font.charDict[replaceMissing])[3] for c in s]) + (length(s) - 1) * font.charSpacing
end

function lineHeight(font::Font, s::AbstractString, replaceMissing::Char=' ')
    return isempty(s) ? 0 : maximum([get(font.charDict, c, font.charDict[replaceMissing])[4] for c in s])
end

function Base.size(font::Font, s::AbstractString, replaceMissing::Char=' ')
    lines = split(s, "\n")

    width = maximum(lineWidth.(Ref(font), lines))
    height = sum(lineHeight.(Ref(font), lines)) + (length(lines) - 1) * font.lineSpacing

    return width, height
end

const stringSurfaceCache = Dict{String, Cairo.CairoSurface}()

function getStringSurface(font::Font, s; replaceMissing::Char=' ', align::Symbol=:left)
    return get!(stringSurfaceCache, s) do
        accX = 0
        accY = 0

        lineAllignFunc = get(stringLineAlignOffsets, align, :left)
        widest, height = size(font, s, replaceMissing)

        surface = CairoARGBSurface(widest, height)
        ctx = getSurfaceContext(surface)

        for line in split(s, "\n")
            width = lineWidth(font, line, replaceMissing)
            accX = lineAllignFunc(widest, width)

            for c in line
                quadX, quadY, width, height = get(font.charDict, c, font.charDict[replaceMissing])
                drawImage(ctx, font.surface, accX, accY, quadX, quadY, width, height)

                accX += width + font.charSpacing
            end

            accY += lineHeight(font, line, replaceMissing) + font.lineSpacing
        end

        return surface
    end
end

function drawString(ctx::Cairo.CairoContext, font::Font, s, x=0, y=0; scale=1.0, replaceMissing::Char=' ', align::Symbol=:left, tint=colors.canvas_font_color)
    shouldSave = x != 0 || y != 0 || scale != 1

    if shouldSave
        Cairo.save(ctx)
        Cairo.translate(ctx, x, y)
        Cairo.scale(ctx, scale, scale)
    end

    surface = getStringSurface(font, s, replaceMissing=replaceMissing, align=align)
    drawImage(ctx, surface, 0, 0, tint=tint)

    if shouldSave
        Cairo.restore(ctx)
    end
end

const pico8FontString = raw"""

 !"#$% ´()*+,-./
0123456789:;<=>?
@abcdefghijklmno
pqrstuvwxyz[\]^_
`ABCDEFGHIJKLMNO
PQRSTUVWXYZ{|}~
"""

const pico8Font = Font(abs"../assets/pico8_font.png", pico8FontString)


using Base.Meta

"""
    AbstractVertex{P}

A supertype for all vertex types with a payload of type `P`.
"""
abstract type AbstractVertex{P} end

"""
    payload(v)

Get a vertex's payload
"""
payload(v::AbstractVertex) = v.payload

zero!(v::AbstractVertex{P}) where P = v.payload = zero(P)

"""
    id(v)

Get the id of a vertex.
"""
id(v::AbstractVertex) = v.id

"""
    above(v)

Get an array of all vertices above `v`.
"""
above(v::AbstractVertex) = v.above

"""
    below(v)

Get an array of all vertices below `v`.
"""
below(v::AbstractVertex) = v.below

Base.length(v::AbstractVertex) = length(id(v))

Base.eachindex(v::AbstractVertex) = eachindex(id(v))

Base.getindex(v::AbstractVertex, idx...) = getindex(id(v), idx...)

bookend(x::AbstractString, left::AbstractString="[", right::AbstractString="]") = left * x * right

prettyname(v::AbstractVertex) = bookend(join(repr.(id(v)), ", "))

function Base.show(io::IO, v::AbstractVertex)
    ex = Meta.parse(string(typeof(v)))
    t = if ex isa Symbol
        string(ex)
    else
        string(ex.args[1])
    end
    print(io, t, "(", prettyname(v), ", ", payload(v), ")")
end

function Base.show(io::IO, ::MIME"text/dot", v::AbstractVertex)
    print(io, "[label=\"", prettyname(v), "\\n")
    try
        show(io, MIME("text/dot"), payload(v))
    catch
        print(io, payload(v))
    end
    print(io, "\"]")
end

abstract type AbstractUnnamedVertex{P} <: AbstractVertex{P} end

abstract type AbstractNamedVertex{N,P} <: AbstractVertex{P} end

name(v::AbstractNamedVertex) = v.name

function prettyname(v::AbstractNamedVertex)
    ns = String[]
    for n in name(v)
        push!(ns, if length(n) == 1
            join(repr.(n), ", ")
        else
            bookend(join(repr.(n), ", "), "{", "}")
        end)
    end
    bookend(join(ns, ", "))
end

"""
    VertexID = Vector{Int}

A type alias for integer arrays used as ids of vertices.
"""
const VertexID = Vector{Int}

"""
    UnnamedVertex{P} <: AbstractVertex{P}

A standard vertex containing a `id`, `payload`, and arrays of vertices
`above` and `below` it.
"""
mutable struct UnnamedVertex{P, V <: AbstractUnnamedVertex{P}} <: AbstractUnnamedVertex{P}
    id::VertexID
    payload::P
    above::Vector{V}
    below::Vector{V}
end

"""
    UnnamedVertex(id, p)

Construct a `Vertex` with a given `id` and payload (`p`), with no vertices
above or below.
"""
function UnnamedVertex(id::VertexID, p::P) where P
    UnnamedVertex(id, p, UnnamedVertex{P}[], UnnamedVertex{P}[])
end

"""
    UnnamedVertex{P}(id)

Construct a `Vertex` with a given `id`, a "zeroed" payload of type `P` and no
vertices above or below.
"""
UnnamedVertex{P}(id::VertexID) where P = UnnamedVertex(id, zero(P))

clone(v::UnnamedVertex) = UnnamedVertex(id(v), payload(v))

mutable struct Vertex{N, P, V <: AbstractNamedVertex{N,P}} <: AbstractNamedVertex{N,P}
    id::VertexID
    name::Vector{Vector{N}}
    payload::P
    above::Vector{V}
    below::Vector{V}
end

function Vertex(id::VertexID, name::AbstractVector{Vector{N}}, p::P) where {N,P}
    Vertex(id, name, p, Vertex{N,P}[], Vertex{N,P}[])
end

function Vertex{N,P}(id::VertexID, name::AbstractVector{Vector{N}}) where {N,P}
    Vertex(id, name, zero(P))
end

clone(v::Vertex) = Vertex(id(v), name(v), payload(v))

function Base.convert(::Type{UnnamedVertex{T}}, v::Vertex{N,T}) where {N,T}
    a = convert(Vector{UnnamedVertex{T}}, above(v))
    b = convert(Vector{UnnamedVertex{T}}, below(v))
    UnnamedVertex(id(v), payload(v), a, b)
end

"""
    isbelow(a, b)

Determine whether or not `a` is below `b`.
"""
@inline function isbelow(xs::VertexID, ys::VertexID)
    for i in eachindex(ys)
        is_valid = false
        for j in eachindex(xs)
            if xs[j] & ys[i] == xs[j]
                is_valid = true
                break
            end
        end
        if !is_valid
            return false
        end
    end
    true
end

isbelow(a::AbstractVertex, b::AbstractVertex) = isbelow(id(a), id(b))

"""
    genvertices(::Type{V}, n) where {V <: AbstractUnnamedVertex}
"""
function genvertices(::Type{V}, n::Int) where {V <: AbstractUnnamedVertex}
    if n < 1
        throw(DomainError(n, "at least one node is required"))
    end
    vs = V[]
    m = (1 << n) - 1;
    for i in 1:m
        push!(vs, V([i]))
        genvertices!(vs, i + 1, m, [i])
    end
    vs
end

function genvertices!(vs::AbstractVector{V}, i, m, c) where {V <: AbstractUnnamedVertex}
    if i < m
        genvertices!(vs, i+1, m, c[:])
    end

    if i <= m
        z = 0
        for j in 1:length(c)
            z = i & c[j]
            if (z == i || z == c[j])
                return vs
            end
        end
        push!(c, i)
        push!(vs, V(c))
        genvertices!(vs, i+1, m, c[:])
    end

    vs
end

function extractname(nodenames::AbstractVector{N}, id::Int) where N
    name = N[]
    for i in eachindex(nodenames)
        if id & (1 << (i-1)) != 0
            push!(name, nodenames[i])
        end
    end
    name
end

extractname(nodenames::AbstractVector, id::VertexID) = map(i -> extractname(nodenames, i), id)

function genvertices(::Type{V}, names::AbstractVector{N}) where {N, V <: AbstractNamedVertex{N}}
    if isempty(names)
        throw(DomainError(n, "at least one node is required"))
    end
    n = length(names)
    vs = V[]
    m = (1 << n) - 1;
    for i in 1:m
        id = [i]
        push!(vs, V(id, extractname(names, id)))
        genvertices!(vs, names, i + 1, m, [i])
    end
    vs
end

function genvertices!(vs::AbstractVector{V}, names, i, m, c) where {V <: AbstractNamedVertex}
    if i < m
        genvertices!(vs, names, i+1, m, c[:])
    end

    if i <= m
        z = 0
        for j in 1:length(c)
            z = i & c[j]
            if (z == i || z == c[j])
                return vs
            end
        end
        push!(c, i)
        push!(vs, V(c, extractname(names, c)))
        genvertices!(vs, names, i+1, m, c[:])
    end

    vs
end

"""
    toposort!(vs)

Topologically sort vertices according to the `isbelow` relation.
"""
function toposort!(vs::AbstractVector{V}) where {V <: AbstractVertex}
    n, u, v = length(vs), 1, 1
    while v < n - 1
        u = v
        for i in u:n
            is_bottom = true
            for j in u:n
                if i != j && isbelow(vs[j], vs[i])
                    is_bottom = false
                    break
                end
            end
            if is_bottom
                vs[v], vs[i] = vs[i], vs[v]
                v += 1
            end
        end
    end
    vs
end

"""
    Hasse{V}

A Hasse diagram (lattice) with vertices of type `V`
"""
mutable struct Hasse{V <: AbstractVertex}
    top::V
    bottom::V
    vertices::Vector{V}
end

function Hasse(::Type{P}, n::Int) where P
    if n > 5
        @warn """Hasse diagrams are computationally expensive to build for n > 5.
        This is gonna take a while and lot of memory""" n
    end
    Hasse(genvertices(UnnamedVertex{P}, n))
end

function Hasse(::Type{P}, names::AbstractVector{N}) where {N,P}
    if length(names) > 5
        @warn """Hasse diagrams are computationally expensive to build for n > 5.
        This is gonna take a while and lot of memory""" n=length(names)
    end
    Hasse(genvertices(Vertex{N,P}, names))
end

function Hasse(vs::AbstractVector{V}; sort::Bool=true) where {V <: AbstractVertex}
    if isempty(vs)
        throw(ArgumentError("vertex list is empty"))
    end

    if sort
        toposort!(vs)
    end

    for i in 1:length(vs)
        for j in i+1:length(vs)
            if isbelow(vs[i], vs[j])
                stop = false
                for v in above(vs[i])
                    if isbelow(v, vs[j])
                        stop = true
                        break
                    end
                end
                if stop
                    break
                else
                    push!(above(vs[i]), vs[j])
                    push!(below(vs[j]), vs[i])
                end
            end
        end
    end

    Hasse(vs[end], vs[1], vs)
end

"""
    top(h)

Get the top vertex of the Hasse diagram `h`.
"""
top(h::Hasse) = h.top

"""
    bottom(h)

Get the bottom vertex of the Hasse diagram `h`.
"""
bottom(h::Hasse) = h.bottom

"""
    vertices(h)

Get the array of vertices of the Hasse diagram `h`.
"""
vertices(h::Hasse) = h.vertices

Base.length(h::Hasse) = length(vertices(h))

Base.eachindex(h::Hasse) = eachindex(vertices(h))

Base.getindex(h::Hasse, idx...) = getindex(vertices(h), idx...)

zero!(h::Hasse) = zero!.(vertices(h))

"""
    edgelist(h)

Get the array of edges — pairs of vertex ids — of the Hasse diagram `h`.
"""
function edgelist(h::Hasse)
    edges = NTuple{2, VertexID}[]
    vs = vertices(h)
    for v in vs, w in v.above
        push!(edges, (id(v), id(w)))
    end
    edges
end

function prune(f::Function, h::Hasse)
    vs = clone.(filter(f, vertices(h)))
    if isempty(vs)
        throw(ErrorException("pruned diagram has no vertices"))
    end
    Hasse(vs)
end
prune(h::Hasse) = prune(!iszero ∘ payload, h)

Base.istextmime(::MIME"text/dot") = true

function Base.show(io::IO, ::MIME"text/dot", h::Hasse; names=nothing)
    println(io, "digraph Hasse {")
    println(io, "  rankdir=\"BT\";")
    vs = vertices(h)
    for i in eachindex(vs)
        print(io, "  ", i, " ")
        show(io, MIME("text/dot"), vs[i])
        println(io, ";")
    end
    println(io)
    for i in eachindex(vs)
        for v in above(vs[i])
            j = findfirst(w -> v == w, vs)
            println(io, "  ", i, " -> ", j, ";")
        end
    end
    println(io, "}")
end

function graphviz(filename::AbstractString, h::Hasse;
                  format=lowercase(last(splitext(filename)))[2:end])
    open(`dot -T$format -o$filename`, "w", stdout) do io
        show(io, MIME("text/dot"), h)
    end
end


## SEXPREC types as defined in ENV["R_INCLUDE_DIR"]/Rinternals.h

const NILSXP     =  0 # NULL in R
const SYMSXP     =  1 # R Symbol
const LISTSXP    =  2 # internal "pairs list", the R list type is 19
const CLOSXP     =  3 # closures
const ENVSXP     =  4 # environments
const PROMSXP    =  5 # promises: unevaluated closure arguments
const LANGSXP    =  6 # language constructs (special lists
const SPECIALSXP =  7 # special forms
const BUILTINSXP =  8 # builtin non-special forms
const CHARSXP    =  9 # "scalar" string type (internal only
const LGLSXP     = 10 # logical vectors
# 11 and 12 were factors and ordered factors in the 1990s
const INTSXP     = 13 # integer vectors
const REALSXP    = 14 # real variables
const CPLXSXP    = 15 # complex variables
const STRSXP     = 16 # string vectors
const DOTSXP     = 17 # dot-dot-dot object
const ANYSXP     = 18 # make "any" args work.
const VECSXP     = 19 # generic vectors
const EXPRSXP    = 20 # expressions vectors
const BCODESXP   = 21 # byte code
const EXTPTRSXP  = 22 # external pointer
const WEAKREFSXP = 23 # weak reference
const RAWSXP     = 24 # raw bytes
const S4SXP      = 25 # S4, non-vector