mutable struct MCVIActionBackup{S,A}
    act::A
    ba::MCVIBelief
    sa::Vector{S}

    alpha_edges::Vector{AlphaEdge}
    to_append
    n                           # MCVINode
end

mutable struct MCVIBeliefBackup
    belief::MCVIBelief
    act_backupers::Vector{MCVIActionBackup}
    last_index::Int64
    max_node
    maxv::Float64
end

MCVIBeliefBackup(belief::MCVIBelief) = MCVIBeliefBackup(belief, Vector{MCVIActionBackup}(), 1, nothing, -Inf)

function initialize_belief_backup!(bb::MCVIBeliefBackup, pomdp::POMDPs.POMDP{S}, num_state::Int64, rng::AbstractRNG) where {S}
    for a in actions(pomdp)
        if a == init_lower_action(pomdp)
            continue
        end
        pb = next(bb.belief, a, pomdp, rng)
        sa = S[rand(rng, pb) for i in 1:num_state]
        ac = MCVIActionBackup(a, pb, sa, Vector{AlphaEdge}(), nothing, nothing)
        push!(bb.act_backupers, ac)
    end
    return bb
end

function prune_alpha_edges(alpha_edges, actback::MCVIActionBackup{S,A}, pomdp::POMDPs.POMDP, num_prune_obs::Int64, rng::AbstractRNG) where {S,A}
    ba = actback.ba
    sa = actback.sa
    keep = Vector{Bool}(undef, length(alpha_edges))
    for i in 1:num_prune_obs
        # Sample observation from belief
        obs = gen(DDNNode(:o), pomdp, rand(rng, ba), rng)
        obswt = zeros(length(sa))
        for (k,s) in enumerate(sa)
            obswt[k] = obs_weight(pomdp, s, obs)
        end
        maxv = -Inf
        local maxj::Int64
        for (j, a) in enumerate(alpha_edges)
            v = dot(obswt, a.vec)
            if v > maxv
                maxv = v
                maxj = j
            end
        end
        keep[maxj]= true
    end
    pruned = Vector{AlphaEdge}()
    for (i, ae) in enumerate(alpha_edges)
        if keep[i]
            push!(pruned, ae)
        end
    end
    return pruned
end

"""
Calls the least square compute function
"""
function compute_alpha_edges(nodes::Vector{MCVINode}, actback::MCVIActionBackup{S,A}, policy::MCVIPolicy, sim::MCVISimulator, pomdp::POMDPs.POMDP, scratch::Scratch; debug=false) where {S,A}
    ba = actback.ba
    sa = actback.sa
    scratch.X = zeros(size(scratch.X,1), length(sa))
    alpha_edges = Vector{AlphaEdge}(undef, length(nodes))
    # println("len_nodes: $(length(nodes))")
    #debug && tic()
    alpha_edges = pmap((n)->compute(sa, policy, sim, pomdp, n), nodes) # FIXME simple
    # for (i,n) in enumerate(nodes)
    #     alpha_edges[i] = compute(sa, policy, sim, pomdp, n)
    # end

    # alpha_edges = pmap((n)->compute(sa, policy, sim, pomdp, n, ba, scratch), nodes)
    # for i in 1:length(nodes)
    #     n = nodes[i]
    #     alpha_edges[i] = compute(sa, policy, sim, pomdp, n, ba, scratch) # FIXME Slowwww
    # end
    #debug && toc()
    return alpha_edges
end

function add_alpha_edges!(actback::MCVIActionBackup{S,A}, edges, updater::MCVIUpdater, pomdp::POMDPs.POMDP, num_prune_obs::Int64, rng::AbstractRNG) where {S,A}
    new_alpha_edges = deepcopy(actback.alpha_edges)
    append!(new_alpha_edges, edges)
    # Prune
    new_alpha_edges = prune_alpha_edges(new_alpha_edges, actback, pomdp, num_prune_obs, rng) # Slow?

    # check
    l1 = length(actback.alpha_edges)
    l2 = length(new_alpha_edges)
    if (l1 == l2) && actback.alpha_edges[l1].id == new_alpha_edges[l2].id
        return actback.n
    end

    actback.alpha_edges = new_alpha_edges
    actback.n = create_node(updater, actback.act, actback.sa, actback.alpha_edges)
    return actback.n
end

"""
Backup action
"""
function backup(actback::MCVIActionBackup{S,A}, policy::MCVIPolicy, sim::MCVISimulator, pomdp::POMDPs.POMDP, nodes, num_prune_obs::Int64, scratch::Scratch; debug=false) where {S,A}
    # Compute alpha edges
    new_alpha_edges = compute_alpha_edges(nodes, actback, policy, sim, pomdp, scratch, debug=debug)
    # Add alpha edges
    n = add_alpha_edges!(actback, new_alpha_edges, policy.updater, pomdp, num_prune_obs, sim.rng)
    return n
end

"""
Returns the best node and its value
"""
function find_best_node(nodes::Vector{MCVINode}, policy::MCVIPolicy, sim::MCVISimulator, pomdp::POMDPs.POMDP, belief::MCVIBelief, num_eval_belief::Int64)
    @assert length(nodes) > 0
    vs = pmap((n)->evaluate(belief, policy, sim, pomdp, n, num_eval_belief), nodes)
    # vs = zeros(length(nodes))
    # for i in 1:length(nodes)
    #     n = nodes[i]
    #     vs[i] = evaluate(belief, policy, sim, pomdp, n, num_eval_belief)
    # end
    (maxv, maxi) = findmax(vs)
    max_node = nodes[maxi]
    return (max_node, maxv)
end

"""
Update belief backup's best node
"""
function update!(bb::MCVIBeliefBackup, nodes::Vector{MCVINode}, pomdp::POMDPs.POMDP, policy::MCVIPolicy, sim::MCVISimulator, eps::Float64, num_eval_belief::Int64)
    (n, v) = find_best_node(nodes, policy, sim, pomdp, bb.belief, num_eval_belief)
    if bb.maxv + eps < v
        bb.maxv = v
        bb.max_node = n
    end
end

"""
Backup belief
"""
function backup(bb::MCVIBeliefBackup, policy::MCVIPolicy, sim::MCVISimulator, pomdp::POMDPs.POMDP, num_prune_obs::Int64, num_eval_belief::Int64, scratch::Scratch; debug=false)
    # Get newer nodes
    nodes = policy.updater.nodes_queue[bb.last_index:end]
    # Update best belief
    t1 = @elapsed begin
        update!(bb, nodes, pomdp, policy, sim, -0.1, num_eval_belief) # Try ϵ less
    end
    debug && printstyled("update", color=:cyan)
    debug && println(" (nodes): $(t1)s")

    # Get new nodes from action backup
    t2 = @elapsed begin
        new_nodes = Vector{MCVINode}(undef, length(bb.act_backupers))
        for (i, actback) in enumerate(bb.act_backupers)
            new_nodes[i] = backup(actback, policy, sim, pomdp, nodes, num_prune_obs, scratch, debug=debug) # Backup action, FIXME Slowwww
        end
    end
    debug && printstyled("backup action", color=:cyan)
    debug && println(" (nodes): $(t2)s")

    bb.last_index += length(nodes)
    update!(bb, new_nodes, pomdp, policy, sim, +0.1, num_eval_belief) # Try ϵ more

    return (bb.max_node, bb.maxv)
end

function backup(belief::MCVIBelief, policy::MCVIPolicy, sim::MCVISimulator, pomdp::POMDPs.POMDP, num_state::Int64, num_prune_obs::Int64, num_eval_belief::Int64, scratch::Scratch; debug=false)
    # Belief backup struct
    bb = MCVIBeliefBackup(belief)
    initialize_belief_backup!(bb, pomdp, num_state, sim.rng)
    return backup(bb, policy, sim, pomdp, num_prune_obs, num_eval_belief, scratch, debug=debug)
end


import Random

sample_snapshot_id() = sample_snapshot_ids(1) |> first
sample_snapshot_ids(num_samples) = rand(get_snapshot_ids(), num_samples)

function sample_daughters(snapshot, organism_to_replace::Organism, genotype, num_samples)
    daughters = []

    while length(daughters) < num_samples
        append!(daughters, get_daughters(snapshot, organism_to_replace, genotype))
    end

    return daughters[1:num_samples]
end

function sample_fitness_values(snapshot, organism_to_replace::Organism, genotype, num_samples)
    [get_fitness(snapshot, organism_to_replace, genotype) for _ in 1:num_samples]
end


#################### Model Expression Operators ####################

function modelfx(literalargs::Vector{Tuple{Symbol,DataType}}, f::Function)
    modelfxsrc(literalargs, f)[1]
end

function modelfxsrc(literalargs::Vector{Tuple{Symbol,DataType}}, f::Function)
    args = Expr(:tuple, map(arg -> Expr(:(::), arg[1], arg[2]), literalargs)...)
    expr, src = modelexprsrc(f, literalargs)
    fx = Core.eval(Main, Expr(:function, args, expr))
    (fx, src)
end

function modelexprsrc(f::Function, literalargs::Vector{Tuple{Symbol,DataType}})

    m = first(methods(f).ms) # returns a mehtod datatype
    argnames = Vector{Any}(undef, m.nargs)
    argnames = Base.method_argnames(m)
    fkeys = Symbol[argnames[2:end]...]

    ftypes = DataType[m.sig.parameters[2:end]...]
    n = length(fkeys)

    literalinds = Int[]
    for (key, T) in literalargs
        i = findfirst(fkey -> fkey == key, fkeys)
        if i !== nothing && ftypes[i] == T
            push!(literalinds, i)
        end
    end
    nodeinds = setdiff(1:n, literalinds)

    all(T -> T == Any, ftypes[nodeinds]) ||
        throw(ArgumentError("model node arguments are not all of type Any"))

    modelargs = Array{Any}(undef, n)
    for i in nodeinds
        modelargs[i] = Expr(:call, MCPhylo.mod_value, Expr(:ref, :model, QuoteNode(fkeys[i])))
    end
    for i in literalinds
        modelargs[i] = fkeys[i]
    end
    expr = Expr(:block, Expr(:(=), :f, f), Expr(:call, :f, modelargs...))
    (expr, fkeys[nodeinds])
end

mod_value(d::AbstractVariate) = d.value
mod_value(d) = d

#################### Mathematical Operators ####################

# isprobvec(p::AbstractVector) = isprobvec(convert(Vector{Float64}, p))
isprobvec(s::AbstractVariate) = isprobvec(s.value)

cummean(x::AbstractArray) = mapslices(cummean, x, dims = 1)

function cummean(x::AbstractVector{T}) where {T<:Real}
    y = similar(x, Float64)
    xs = 0.0
    for i = 1:length(x)
        xs += x[i]
        y[i] = xs / i
    end
    y
end

@inline logit(x::Real) = log(x / (1.0 - x))
@inline invexpit(x::Real) = 1.0 / (x - x^2)
@inline invlogit(x::Real, λ::Real = 1.0) = 1.0 / (exp(-λ * x) + 1.0)
@inline loginvlogit(x::Real, λ::Real = 1.0) = -log(exp(-λ * x) + 1.0)

function loginvlogitder(x::T, λ::T)::T where {T<:Real}
    λ / (exp(λ * x) + 1)
end

function invlogder(x::T, λ::T)::T where {T<:Real}
    invlogit(x, λ) * invlogit(-x, λ)
end

Zygote.@adjoint function invlogit(x::Real, λ::Real)
    invlogit(x, λ), Δ -> (nothing, invlogder(Δ, λ))
end

Zygote.@adjoint function loginvlogit(x::Real, λ::Real)
    loginvlogit(x, λ), Δ -> (nothing, loginvlogitder(Δ, λ))
end

## Csorgo S and Faraway JJ. The exact and asymptotic distributions of the
## Cramer-von Mises statistic. Journal of the Royal Statistical Society,
## Series B, 58: 221-234, 1996.
function pcramer(q::Real)
    p = 0.0
    for k = 0:3
        c1 = 4.0 * k + 1.0
        c2 = c1^2 / (16.0 * q)
        p += gamma(k + 0.5) / factorial(k) * sqrt(c1) * exp(-c2) * besselk(0.25, c2)
    end
    p / (pi^1.5 * sqrt(q))
end

#################### Auxiliary Functions ####################
## pmap2 is a partial work-around for the pmap issue in julia 0.4.0 of worker
## node errors being blocked.  In single-processor mode, pmap2 calls map
## instead to avoid the error handling issue.  In multi-processor mode, pmap is
## called and will apply its error processing.

function pmap2(f::Function, lsts::AbstractArray)
    if (nprocs() > 1)
        pmap(f, lsts)
    else
        map(f, lsts)
    end
end

## Custom Macro to get embedded errror messages
macro async_showerr(ex)
    quote
        t = @async try
            eval($(esc(ex)))
        catch err
            bt = catch_backtrace()
            println()
            showerror(stderr, err, bt)
        end
    end
end

"""
  assign_mcmc_work(f::Function, lsts::AbstractArray
                   )::Tuple{Vector{Tuple{Chains, Model, ModelState}},
                      Array{Float64, 2}, Vector{AbstractString}}

--- INTERNAL ---
Starts the MCMC chain generation (on multiple workers if possible) and also
starts parallel ASDSF - if possible and requested by the user.
"""
function assign_mcmc_work(
    f::Function,
    lsts::AbstractArray,
    sp::SimulationParameters,
    conv_storage::Union{Nothing,ConvergenceStorage},
)::Tuple{
    Vector{Tuple{Chains,Model,ModelState}},
    Array{Float64,2},
    Vector{AbstractString},
    Union{Nothing,ConvergenceStorage},
}

    ASDSF::Bool = sp.asdsf
    statnames::Vector{AbstractString} = []
    # count the number of trees per step per chain
    tree_dim::Int64 = 0
    for i in lsts[1][1].nodes
        if isa(i[2], Stochastic{<:GeneralNode})
            push!(statnames, string("asdsf_", string(i[1])))
            tree_dim += 1
        end # if
    end # for
    nchains::Int64 = length(lsts)
    ntrees::Int64 = ASDSF ? floor((last(lsts[1][3]) - lsts[1][4]) / sp.freq) : 0
    results = Dict{Int64,Tuple{Chains,Model,ModelState}}()
    if ASDSF
        # set up remote channels for communication across workers for ASDSF-on-the-fly
        if nworkers() >= nchains + 1
            r_channels =
                [RemoteChannel(() -> Channel{Vector{AbstractString}}(100)) for x = 1:nchains]
        else
            r_channels =
                [RemoteChannel(() -> Channel{Vector{AbstractString}}(ntrees)) for x = 1:nchains]
        end # if / else
        ASDSF_vals::Vector{Vector{Float64}} = []
        # add necessary ASDSF related arguments to lsts for mcmc_worker function
        for (ind, lst) in enumerate(lsts)
            append!(lst, [sp.freq, r_channels[ind]])
        end # for
        # add a list to lsts array, that contains args for calculate_convergence
        push!(lsts, [sp, conv_storage, r_channels, ntrees, 1:tree_dim])
    end # if

    # set up a ProgressMeter for each chain...
    meters = [Progress(lsts[1][3][end] - lsts[1][3][1] + 1; desc="Chain $c: ", enabled=sp.verbose, offset=c-1, showspeed=true) for c in 1:nchains]
    # ... & one RemoteChannel that will be used to communicate the updates of the progress bar across parallel functions
    channel = RemoteChannel(() -> Channel{Integer}(1))
    # add the RemoteChannels to the lists that are later passed to the mcmc_worker
    for c in 1:nchains
        insert!(lsts[c], 8, (channel, c))
    end # for
    results_vec = []
    @sync begin
        # check RemoteChannel for new entries and updates the ProgressMeters
        finished_chains = 0
        @async while finished_chains < nchains
            chain = take!(channel)
            chain > 0 ? ProgressMeter.next!(meters[chain]) : finished_chains += 1
        end # while
        results_vec = pmap2(f, lsts)
    end # @sync
    ASDSF && close.(r_channels)
    if ASDSF
        asdsf_results = results_vec[end]
        append!(ASDSF_vals, asdsf_results[1])
        conv_storage = asdsf_results[2]
        stats = Array{Float64,2}(undef, length(ASDSF_vals[1]), length(ASDSF_vals))
        for i = 1:length(ASDSF_vals)
            stats[:, i] = ASDSF_vals[i]
        end # for
    else
        stats = zeros(Float64, 0, tree_dim)
    end # if/else
    println("\n")
    return [results_vec[i] for i = 1:nchains], stats, statnames, conv_storage
end # assign_mcmc_work

"""
  mcmc_or_convergence(args::AbstractArray)
    ::Union{Tuple{Chains, Model, ModelState}},
            Tuple{Vector{Vector{Float64}}, ConvergenceStorage}

--- INTERNAL ---
Used in pmap call in assign_mcmc_work to correctly dispatch arguments to either
calculate convergence statistics or mcmc chains.
"""
function mcmc_or_convergence(args::AbstractArray
                            )::Union{
                            Tuple{Chains, Model, ModelState},
                            Tuple{Vector{Vector{Float64}}, ConvergenceStorage}
                            }
    # depending on the elements of the args array either start calculate_convergence or...
    if isa(args[1], SimulationParameters)
        calculate_convergence(args...)
    # ...mcmc_worker
    else
        # if no convergence statistics are wanted, then the mcmc_worker will be called with exactly 8 elements in args...
        if length(args) == 8
            mcmc_worker!(args)
        # ...but with on-the-fly convergence statistics, we have 2 more elements (the step range & the remote channel),
        # that we need to dispatch correctly to the mcmc_worker
        else
            mcmc_worker!(args[1:end-2], args[end-1:end]...)
        end # if/else
    end # if/else
end

ind2sub(dims, ind) = Tuple(CartesianIndices(dims)[ind])

showall(v) = showall(stdout, v)

"""
    lcp(str1::T, str2::T)::T where T <: AbstractString

Get the longest common prefix.
"""
function lcp(str1::T, str2::T)::T where {T<:AbstractString}
    minl::Int64 = min(length(str1), length(str2))
    outs::T = ""
    minl == 0 && return outs
    for i = 1:minl
        str1[i] == str2[i] ? outs *= str1[i] : return outs
    end # for
    return outs
end


module QMSimFiles

using Printf
using Distributions
using HDF5

# chrosomome map
struct QMSimChromosomeMap
   nLoci::Int
   nSNP::Int
   nQTL::Int
   seqQTL::Vector{Int}
   pos::Vector{Float64}
   maxAllele::Int
   naQTL::Vector{Int}
   effQTL::Matrix{Float64}
end

# a set of chromosome maps
struct QMSimMap
   nchr::Int
   totalSNP::Int
   totalQTL::Int
   chr::Vector{QMSimChromosomeMap}
end

# genotypes in a chromosome
struct QMSimChromosomeGenome
   gp::Vector{Int8}
   gm::Vector{Int8}
end

# genotypes of an individual
struct QMSimIndividualGenome
   tbv::Float64
   chr::Vector{QMSimChromosomeGenome}
end

# genotypes
struct QMSimPopulationGenome
   map::QMSimMap
   individual::Vector{QMSimIndividualGenome}
end

import Base: ==
export QMSimChromosomeMap, QMSimMap, QMSimChromosomeGenome, QMSimIndividualGenome, QMSimPopulationGenome
export read_maps, read_genotypes, get_true_breeding_value, read_qmsim_data,
       save_qmsim_data_hdf5, read_qmsim_map_hdf5, read_qmsim_data_hdf5, get_qmsim_genotype_size,
       read_qmsim_individual_hdf5, add_qmsim_individual_hdf5, create_qmsim_map_hdf5,
       export_qmsim_individual_blupf90
export mating

include("io.jl")
include("hdf5io.jl")
include("tool.jl")
include("mating.jl")

end


"ALTRO solve"
function solve!(prob::Problem{T,Discrete}, opts::ALTROSolverOptions) where T<:AbstractFloat
    t0 = time()

    # Set terminal condition if using projected newton
    if opts.projected_newton
        opts_al = opts.opts_al
        if opts.projected_newton_tolerance >= 0
            opts_al.constraint_tolerance = opts.projected_newton_tolerance
        else
            opts_al.constraint_tolerance = 0
            opts_al.kickout_max_penalty = true
        end
    end

    # create ALTRO problem
    prob_altro, state = altro_problem(prob,opts)

    # Create ALTRO solver
    solver = ALTROSolver(prob_altro, opts)
    to = solver.stats[:timer]

    # primary solve (augmented Lagrangian)
    t_al = time()
    @info "Augmented Lagrangian solve..."
    @timeit to "AL solve" solve!(prob_altro, solver.solver_al)
    time_al = time() - t_al

    t_pn = time()
    pn_solver = ProjectedNewtonSolver(prob_altro, opts.opts_pn)
    if opts.projected_newton
        @info "Projected Newton solve..."
        # _pn_solver = solve!(prob_altro, pn_solver)
        # return pn_solver, solver.solver_pn
        # @assert pn_solver == solver.solver_pn

        solver.solver_pn.V = PrimalDual(prob_altro)
        @timeit to "Projection solve" solve!(prob_altro,solver.solver_pn)#pn_solver)
    end
    time_pn = time()- t_pn

    # process primary solve results
    @timeit to "Process results" process_results!(prob,prob_altro,state,opts)

    # Stats
    time_total = time() - t0
    stats = solver.stats
    stats[:time] = time_total
    stats[:time_al] = time_al
    stats[:time_pn] = time_pn

    return solver
end

"Processes ALTRO solve results, including: remove slack controls, add minimum time controls"
function process_results!(prob::Problem{T},prob_altro::Problem{T},
        state::NamedTuple,opts::ALTROSolverOptions{T}) where T

    # move results to original problem
    copyto!(prob.X,prob_altro.X,prob.model.n)
    copyto!(prob.U,prob_altro.U,prob.model.m)

    if !state.infeasible && !state.minimum_time && !state.projected_newton
        return nothing
    else
        # remove infeasible, perform feasible projection, resolve
        if state.infeasible
            # infeasible problem -> feasible problem
            prob_altro = infeasible_to_feasible_problem(prob,prob_altro,state,opts)

            # # secondary solve (augmented Lagrangian)
            if opts.resolve_feasible_problem
                # println("Resolving feasible")
                solve!(prob_altro,opts.opts_al)
                copyto!(prob.X,prob_altro.X,prob.model.n)
                copyto!(prob.U,prob_altro.U,prob.model.m)
            end
        end

        # update original problem (minimum time solve will return dt as controls at U[k][end])
        if state.minimum_time
            for k = 1:prob.N-1
                prob.U[k] = [prob_altro.U[k][1:prob.model.m]; prob_altro.U[k]]
            end
            # copyto!(prob.X,prob_altro.X,prob.model.n)
        end

        # if state.projected_newton
        #     #TODO
        #     nothing
        # end
    end

    return nothing
end

"Return ALTRO problem from original problem, includes: adding slack controls, adding minimum time controls"
function altro_problem(prob::Problem{T},opts::ALTROSolverOptions{T}) where T
    prob_altro = prob

    # create infeasible problem
    if !all(x->isnan(x),prob_altro.X[1])
        @info "Infeasible Solve"
        prob_altro = infeasible_problem(prob_altro,opts.R_inf)
        infeasible = true
    else
        infeasible = false
    end

    # create minimum time problem
    if prob_altro.tf == 0.0
        @info "Minimum Time Solve"
        prob_altro = minimum_time_problem(prob_altro,opts.R_minimum_time,
            opts.dt_max,opts.dt_min)
        minimum_time = true
    else
        minimum_time = false
    end

    state = (infeasible=infeasible,minimum_time=minimum_time,
        projected_newton=opts.projected_newton)

    return prob_altro, state
end


"""
    gams = parsegams(filename)

Do low-level lexing and parsing, but do not try to compile into a module.
"""
function parsegams(lexed::Vector{AbstractLex})
    gams = Dict{String,Any}()
    i = 1
    while i <= length(lexed)
        stmt = lexed[i]
        if stmt isa StatementEnd
            # do nothing
        elseif stmt isa Keyword
            kwname = stmt.text
            if kwname == "sets"
                sets = getdefault!(gams, kwname, Dict{String,Any})
                newsets = Dict{String,Vector{String}}()
                i = parse_slashed!(newsets, lexed, i)
                # Scan for i*j range declarations
                for (k, s) in newsets
                    if length(s) == 1
                        m = match(r"(\d+)\s*\*\s*(\d+)", s[1])
                        if m != nothing
                            lo, hi = parse(Int, m.captures[1]), parse(Int, m.captures[2])
                            sets[k] = lo == 1 ? Base.OneTo(hi) : (lo:hi)
                        else
                            sets[k] = s
                        end
                    else
                        sets[k] = s
                    end
                end
            elseif kwname == "alias"
                alias = lexed[i+=1]
                @assert(alias isa Parens && length(alias.args) == 3 && alias.args[2] == GText(","))
                sets = gams["sets"]
                symold, symnew = string(alias.args[1]), string(alias.args[3])
                if !haskey(sets, symold)  # GAMS allows arbitrary ordering
                    symnew, symold = symold, symnew
                end
                sets[symnew] = sets[symold]
            elseif kwname == "scalars" || kwname == "parameters"
                params = getdefault!(gams, "parameters", Dict{Any,Any})
                i = parse_slashed!(params, lexed, i; parseitems=false)
            elseif kwname == "table"
                table = getdefault!(gams, "tables", Dict{Any,Any})
                item = lexed[i+=1]
                table[item.name] = item.body
                # Some GAMS files declare it a parameter and then a table; remove it
                # from parameters
                if haskey(gams, "parameters")
                    params = gams["parameters"]
                    key = GText(getname(item.name))
                    if haskey(params, key)
                        delete!(params, key)
                    end
                end
            elseif kwname == "variables" || kwname ∈ vartypes
                attr = "free"
                if kwname ∈ vartypes
                    attr = kwname
                    kw = lexed[i+=1]
                    @assert(kw == Keyword("variables"))
                end
                sets = get(gams, "sets", nothing)
                variables = getdefault!(gams, "variables", OrderedDict{Any,VarInfo})
                while i < length(lexed) && !isa(lexed[i+1], StatementEnd) && !isa(lexed[i+1], Keyword)
                    var = lexed[i+=1]
                    var == GText(",") && continue
                    @assert(var isa GText || var isa GArray)
                    push!(variables, varinfo(var, attr, sets))
                end
            elseif kwname == "models"
                models = getdefault!(gams, kwname, Dict{String,ModelInfo})
                newmodels = Dict{String,Any}()
                i = parse_slashed!(newmodels, lexed, i)
                for (k,v) in newmodels
                    models[k] = ModelInfo(v)
                end
            elseif kwname == "options"
                options = getdefault!(gams, kwname)
                str = lexed[i+=1]
                while !isa(lexed[i+1], StatementEnd)
                    if lexed[i+1] == GText("=")
                        options[str.text] = string(lexed[i+2])
                        i+=1
                    else
                        options[str.text] = ""
                    end
                    i += 1
                end
            elseif kwname == "solve"
                newkw = lexed[i+=1]
                while !isa(newkw, Keyword) || newkw.text == "using"
                    newkw = lexed[i+=1]
                end
                @assert newkw.text == "minimizing" || newkw.text == "maximizing"
                modelname = lexed[i+=1]
                gams[newkw.text] = modelname.text
            elseif kwname == "equations"
                eqs = getdefault!(gams, kwname, OrderedDict{Any,Any})
                while i < length(lexed) && !isa(lexed[i+1], StatementEnd) && !isa(lexed[i+1], Keyword)
                    eq = lexed[i+=1]
                    @assert(eq isa GText || eq isa GArray)
                    # push!(eqs, eq=>"")  # add later so the ordering respects the definition order, not decl order
                end
            elseif kwname == "files"
                filename = lexed[i+=1]
                newkw = lexed[i+=1]
                while newkw != GText("putclose") || lexed[i+1] != filename
                    newkw = lexed[i+=1]
                end
                newkw = lexed[i+=1]
                while !isa(newkw, StatementEnd)
                    newkw = lexed[i+=1]
                end
            elseif kwname ∈ Set(["display"])
                newkw = lexed[i+=1]
                while !isa(newkw, StatementEnd)
                    newkw = lexed[i+=1]
                end
            else
                error("unhandled keyword ", kwname)
            end
        else
            @assert(stmt isa GText || stmt isa GArray)
            if length(lexed) < i+1
                # some files seem to have junk after the last semicolon
                println("skipping trailing ", stmt)
                break
            end
            if lexed[i+1] == Dots("..")
                # an equation definition
                eqs = gams["equations"]
                iend = seek_to_end(lexed, i+=2)
                eqs[stmt] = parseexprs1(lexed[i:iend])
                i = iend
            elseif lexed[i+1] == Dots(".")
                # a variable or model attribute
                @assert(lexed[i+3] == GText("="))
                attr = lexed[i+=2]
                iend = seek_to_end(lexed, i+=2)
                rhs = parseexprs1(lexed[i:iend])
                if attr isa GText && string(attr) ∈ modelattributes
                    m = gams["models"][string(stmt)]
                    m.assignments[string(attr)] = rhs
                else
                    attrstr = isa(attr, GText) ? string(attr) : attr.name
                    varkey = isa(attr, GText) ? stmt : GArray(string(stmt), attr.indices)
                    if attrstr ∈ varattributes
                        v = gams["variables"][varkey]  # gams["variables"][stmt] ?
                        push!(v.assignments, attr=>rhs)
                    else
                        error(attr, " not a recognized attribute")
                    end
                end
                i = iend
            else
                # general assignment (to, e.g., a parameter)
                @assert lexed[i+1] == GText("=")
                iend = seek_to_end(lexed, i+=2)
                assign = getdefault!(gams, "assignments", Vector{Pair{Any,Any}})
                push!(assign, stmt=>parseexprs1(lexed[i:iend]))
                i = iend
            end
        end
        i += 1
    end
    return gams
end
parsegams(arg) = parsegams(lex(arg))

function parse_slashed!(dest, lexed, i; parseitems::Bool=true)
    function parse_slashed1!(dest, lexed, i)
        name = lexed[i+=1]
        @assert(name isa GText || name isa GArray)
        if lexed[i+1] isa StatementEnd
            dest[name] = ""
            return i
        elseif lexed[i+1] isa Slashed
            val = lexed[i+=1]
        elseif lexed[i+1] isa GText && lexed[i+2] isa Slashed
            val = lexed[i+=2]
        else
            dest[name] = ""
            return i
        end
        @assert(val isa Slashed)
        if parseitems
            items = strip.(split(val.text, (',','\n')))
            # drop text for each item
            for j = 1:length(items)
                itemname, itemtext = splitws(items[j])
                items[j] = itemname
            end
            dest[convert(keytype(dest), name)] = items
        else
            dest[name] = val.text
        end
        return i
    end
    i = parse_slashed1!(dest, lexed, i)
    while i < length(lexed) && !isa(lexed[i+1], StatementEnd) && !isa(lexed[i+1], Keyword)
        i = parse_slashed1!(dest, lexed, i)
    end
    return i
end

function parseexprs(lexed, debug=false)
    # Step 1: split at commas
    args = []
    thisarg = []
    for i = 1:length(lexed)
        tok = lexed[i]
        if tok == GText(",")
            @assert(!isempty(thisarg))
            push!(args, copy(thisarg))
            empty!(thisarg)
        else
            push!(thisarg, tok)
        end
    end
    if !isempty(thisarg)
        push!(args, copy(thisarg))
        empty!(thisarg)
    end
    # Step 2: parse each arg
    n = length(args)
    parsed = Vector{AbstractLex}(undef, n)
    todelete = Int[]
    for j = 1:n
        thisarg = args[j]
        # Step 2a: recurse into calls and parse their args
        debug && sexprlf.(thisarg)
        for i = 1:length(thisarg)
            tok = thisarg[i]
            if tok isa GCall
                callargs = parseexprs(tok.args, false)
                thisarg[i] = GCall(tok.name, callargs)
            elseif tok isa Parens
                pargs = parseexprs(tok.args, false)
                thisarg[i] = Parens(pargs)
            end
        end
        # Step 2b: process ^ (highest-precedence operator)
        debug && sexprlf.(thisarg)
        process_ops!(thisarg, ("^",), debug)
        # Step 2c: process * and / (next-highest)
        debug && sexprlf.(thisarg)
        process_ops!(thisarg, ("*", "/"), debug)
        # Step 2d: process + and -
        debug && println("prior to +/-")
        debug && sexprlf.(thisarg)
        process_ops!(thisarg, ("+", "-"), debug)
        # Step 2e: process "and" and "or"
        debug && println("post to +/-")
        debug && sexprlf.(thisarg)
        process_ops!(thisarg, ("and", "or"), debug)
        # Step 2f: process other boolean-returning operations
        process_ops!(thisarg, ("<", "<=", ">", ">="), debug)
        # Step 2g: process relations
        debug && sexprlf.(thisarg)
        process_ops!(thisarg, ("=e=", "=g=", "=l="), debug)
        if length(thisarg) != 1
            println(stderr, "thisarg:")
            foreach(x->(print(stderr, "  "); sexpr(stderr, x); println(stderr)), thisarg)
        end
        @assert(length(thisarg) == 1)
        parsed[j] = stripnewcall!(thisarg[1])
    end
    return parsed
end

function parseexprs1(lexed)
    parsed = parseexprs(lexed)
    @assert(length(parsed) == 1)
    return parsed[1]
end

# A "temporary type" for marking calls that arose from operator-precedence resolution
# Future lower-priority operators will need to recurse into any NewCalls
struct NewCall <: AbstractLex
    call
end

function process_ops!(tokens::Vector, ops, debug)
    debug = false
    # i = length(tokens)
    i = 1
    # while i > 0
    while i < length(tokens)
        debug && @show i
        tok = tokens[i]
        debug && @show tok
        if tok isa NewCall
            newtok = tok.call
            for arg in newtok.args
                process_ops!(arg, ops)
            end
        end
        if !isa(tok, GText) || tok.text ∉ ops
            i += 1
            continue
        end
        if i > 1 && !isoperator(tokens[i-1])
            debug && @show tokens[i-1] isoperator(tokens[i-1])
            tokens[i-1] = NewCall(GCall(tok.text, [tokens[i-1], tokens[i+1]]))
            deleteat!(tokens, (i,i+1))
            i -= 1
        else
            # unary + or -
            @assert(tok.text ∈ ("+", "-"))
            if tok.text == "+"
                tokens[i] = tokens[i+1]
            else
                nexttok = tokens[i+1]
                if nexttok isa GNumber
                    tokens[i] = GNumber(-nexttok.val)
                else
                    tokens[i] = NewCall(GCall("-", [nexttok]))
                end
            end
            deleteat!(tokens, (i+1,))
        end
        i += 1
    end
    return tokens
end
process_ops!(tokens, ops) = nothing

isoperator(op) = op isa GText && op.text ∈ operatortokens

function sexpr(io::IO, l::NewCall)
    print(io, "NewCall(")
    sexpr(io, l.call)
    print(io, ")")
end

function stripnewcall!(expr::NewCall)
    call = expr.call
    for i = 1:length(call.args)
        call.args[i] = stripnewcall!(call.args[i])
    end
    return call
end
stripnewcall!(expr) = expr

function deparen!(expr)
    if expr isa GCall
        for i = 1:length(expr.args)
            expr.args[i] = deparen!(expr.args[i])
        end
    elseif expr isa Parens && length(expr.args) == 1
        expr = expr.args[1]
        expr = deparen!(expr)
    end
    return expr
end

function varinfo(name::GText, attr, sets)
    return name=>VarInfo(attr, ())
end
function varinfo(name::GArray, attr, sets)
    return name=>VarInfo(attr, getaxes(name.indices, sets))
end

function seek_to_end(lexed, i)
    iend = i
    while iend < length(lexed) && !isa(lexed[iend+1], Union{StatementEnd, Keyword})
        iend += 1
    end
    return iend
end

function getwithkey(dict, key::AbstractString)
    for (k, v) in dict
        if getname(k) == key
            return v
        end
    end
    error(key, " not found")
end

function getdefault!(gams, tag, ::Type{T}) where T
    if !haskey(gams, tag)
        gams[tag] = T<:AbstractDict ? T() : (T<:Vector ? T(undef, 0) : error("type $T not handled"))
    end
    return gams[tag]
end
getdefault!(gams, tag) = getdefault!(gams, tag, Dict{String,String})


@testset "ModelFrame (legacy API)" begin

    @testset "#133" begin
        df = (x = rand(12),
              y = categorical(repeat(1:3, inner=4)),
              z = categorical(repeat(1:2, outer=6)));
        f = @formula(x ~ y * z);
        mf = ModelFrame(f, df)
        mm = ModelMatrix(mf)
        @test mm.assign == [1, 2, 2, 3, 4, 4]

        # `q` isn't in df  - should throw ArgumentError
        f2 = @formula(x ~ y * q)
        @test_throws ArgumentError ModelFrame(f2, df)

    end

end


export BatchedSampler

mutable struct BatchedGradSampler{BN, P, S, Sc, Sv, Pv, Gv, Gvv, Gva, LC, Lv, Lgv, Pa, Pc, Pd} <: AbstractIterativeSampler
    bnet::BN
    Ĉ::P

    sampler::S
    sampler_cache::Sc
    samples::Sv

    ψvals::Pv
    ∇vals::Gv
    ∇vals_vec::Gvv
    ∇vec_avg::Gva

    accum::LC

    local_vals::Lv
    ∇local_vals::Lgv

    precond_alg::Pa
    precond_cache::Pc

    observables_sampler

    parallel_cache::Pd
end

function BatchedGradSampler(net,
                        sampl,
                        prob,
                        algo=Gradient();
                        batch_sz=2^4,
                        local_batch_sz=batch_sz,
                        obs_sampler=sampl,
                        par_type=automatic_parallel_type())
    if net isa CachedNet
        throw("Only takes standard network.")
    end

    par_cache      = parallel_execution_cache(par_type)

    bnet           = cached(net, batch_sz)
    v              = state(prob, bnet)
    sampler_cache  = cache(sampl, basis(prob), bnet, par_cache)

    ch_len         = chain_length(sampl, sampler_cache)

    samples        = state(prob, bnet, ch_len)
    ψvals          = similar(trainable_first(bnet), out_type(bnet), 1, batch_sz, ch_len)
    ∇vals, ∇vec    = grad_cache(bnet, batch_sz, ch_len)
    ∇vec_avg       = tuple([similar(∇vec_i, size(∇vec_i, 1)) for ∇vec_i=∇vec]...)

    local_acc      = AccumulatorObsGrad(net, basis(prob), v, local_batch_sz)
    Llocal_vals    = collect(similar(ψvals, size(ψvals)[2:end]...)) #ensure it's on cpu #tocheck
    ∇Llocal_vals   = tuple([similar(ψvals, size(∇vec_i, 1), size(ψvals)[2:end]...) for ∇vec_i=∇vec]...)

    precond        = algorithm_cache(algo, prob, net, par_cache)

    if prob isa KLocalLiouvillian
        obs        = BatchedObsDMSampler(bnet, obs_sampler, basis(prob), batch_sz=batch_sz)
    else
        obs        = BatchedObsKetSampler(samples, ψvals, local_acc)
    end

    nq = BatchedGradSampler(bnet, prob,
            sampl, sampler_cache, samples,
            ψvals, ∇vals, ∇vec, ∇vec_avg,
            local_acc, Llocal_vals, ∇Llocal_vals,
            algo, precond,
            obs,
            par_cache)

    return nq
end

function sample!(is::BatchedGradSampler; sample=true)
    ch_len         = chain_length(is.sampler, is.sampler_cache)
    batch_sz       = size(is.local_vals, 1)

    # Sample the new configurations
    sample && _sample_state!(is)

    # Compute logψ and ∇logψ
    logψ_and_∇logψ!(is.∇vals, is.ψvals, is.bnet, is.samples)

    # Compute LdagL
    for i=1:ch_len
        for j = 1:batch_sz
            σv = state_uview(is.samples, j, i)
            init!(is.accum, σv, is.ψvals[1,j,i], vec_data(is.∇vals[i])[1][:,j] )
            accumulate_connections!(is.accum, is.Ĉ, σv)
            L_loc, ∇L_loc = NeuralQuantum.finalize!(is.accum)
            is.local_vals[j, i]   = L_loc
            uview(is.∇local_vals[1], :, j, i) .= vec_data(∇L_loc)[1]
        end
    end

    Ĉ2 = abs2.(is.local_vals)

    L_stat = stat_analysis(Ĉ2, is.parallel_cache)

    # Center the gradient so that it has zero-average
    _center_gradient!(is)

    # Compute the gradient
    # Flatten the batches and the iterations
    ∇vr = reshape(is.∇vals_vec[1], :, ch_len*batch_sz)
    ∇Ĉr = reshape(is.∇local_vals[1], :, ch_len*batch_sz)
    Ĉr  = reshape(is.local_vals, 1, :)
    #Ĉ2r = reshape(Ĉ2, 1, :)

    # Compute the gradient
    ∇C  = Ĉr*∇Ĉr'
    ∇C ./= (ch_len*batch_sz)
    ∇C .-= L_stat.mean .* is.∇vec_avg[1]'
    ∇C .= conj.(∇C)
    workers_mean!(∇C, is.parallel_cache)

    setup_algorithm!(is.precond_cache, reshape(∇C,:), ∇vr, is.parallel_cache)

    return L_stat, is.precond_cache
end

function compute_observables(is::BatchedGradSampler)
    return compute_observables(is.observables_sampler)
end

function add_observable!(is::BatchedGradSampler, name::String, obs)
    return add_observable!(is.observables_sampler, name, obs)
end

Base.show(io::IO, is::BatchedGradSampler) = print(io,
    "BatchedGradSampler for :"*
    "\n\tnet\t\t: $(is.bnet)"*
    "\n\tproblem\t: $(is.Ĉ)"*
    "\n\tsampler\t: $(is.sampler)")


struct SimulationInputs
    dynamic_injectors::Vector{DynamicWrapper{<:PSY.DynamicInjection}}
    static_injectors::Vector
    static_loads::Vector
    dynamic_branches::Vector{BranchWrapper}
    injection_n_states::Int
    branches_n_states::Int
    variable_count::Int
    inner_vars_count::Int
    bus_count::Int
    ode_range::UnitRange{Int}
    ybus_rectangular::SparseArrays.SparseMatrixCSC{Float64, Int}
    dyn_lines::Bool
    total_shunts::SparseArrays.SparseMatrixCSC{Float64, Int}
    lookup::Dict{Int, Int}
    DAE_vector::Vector{Bool}
    mass_matrix::LinearAlgebra.Diagonal{Float64}
    global_vars_update_pointers::Dict{Int, Int}
    global_state_map::MAPPING_DICT

    function SimulationInputs(
        sys::PSY.System,
        ::Type{T},
    ) where {T <: Union{ConstantFrequency, ReferenceBus}}
        n_buses = get_n_buses(sys)
        Ybus, lookup = _get_ybus(sys)

        TimerOutputs.@timeit BUILD_TIMER "Wrap Branches" begin
            wrapped_branches = _wrap_dynamic_branches(sys, lookup)
            has_dyn_lines = !isempty(wrapped_branches)
            branch_state_counts = 2 * length(wrapped_branches)
            injection_start = 2 * n_buses + branch_state_counts + 1
        end

        TimerOutputs.@timeit BUILD_TIMER "Wrap Dynamic Injectors" begin
            wrapped_injectors = _wrap_dynamic_injector_data(sys, lookup, injection_start)
            var_count = wrapped_injectors[end].ix_range[end]
        end

        TimerOutputs.@timeit BUILD_TIMER "Calculate MM, DAE_vector, Total Shunts" begin
            mass_matrix = _make_mass_matrix(wrapped_injectors, var_count, n_buses)
            DAE_vector = _make_DAE_vector(mass_matrix, var_count, n_buses)
            total_shunts = _make_total_shunts(wrapped_branches, n_buses)
        end

        TimerOutputs.@timeit BUILD_TIMER "Wrap Static Injectors" begin
            wrapped_loads = _wrap_loads(sys, lookup)
            wrapped_static_injectors = _wrap_static_injectors(sys, lookup)
        end

        _adjust_states!(
            DAE_vector,
            mass_matrix,
            total_shunts,
            n_buses,
            PSY.get_frequency(sys),
        )

        global_vars =
            _make_global_variable_index(wrapped_injectors, wrapped_static_injectors, T)

        inner_vars_count = 0
        for i in length(wrapped_injectors):-1:1
            if length(wrapped_injectors[i].inner_vars_index) > 0
                inner_vars_count = wrapped_injectors[i].inner_vars_index[end]
                break
            end
        end
        IS.@assert_op inner_vars_count > 0
        new(
            wrapped_injectors,
            wrapped_static_injectors,
            wrapped_loads,
            wrapped_branches,
            var_count - 2 * n_buses - branch_state_counts,
            branch_state_counts,
            var_count,
            inner_vars_count,
            n_buses,
            injection_start:var_count,
            Ybus,
            has_dyn_lines,
            total_shunts,
            lookup,
            DAE_vector,
            mass_matrix,
            global_vars,
            MAPPING_DICT(),
        )
    end
end

get_dynamic_injectors(inputs::SimulationInputs) = inputs.dynamic_injectors
get_dynamic_branches(inputs::SimulationInputs) = inputs.dynamic_branches
get_static_injectors(inputs::SimulationInputs) = inputs.static_injectors
get_static_loads(inputs::SimulationInputs) = inputs.static_loads
get_ybus(inputs::SimulationInputs) = inputs.ybus_rectangular
get_total_shunts(inputs::SimulationInputs) = inputs.total_shunts
get_lookup(inputs::SimulationInputs) = inputs.lookup
get_DAE_vector(inputs::SimulationInputs) = inputs.DAE_vector
get_mass_matrix(inputs::SimulationInputs) = inputs.mass_matrix
has_dyn_lines(inputs::SimulationInputs) = inputs.dyn_lines
get_global_vars_update_pointers(inputs::SimulationInputs) =
    inputs.global_vars_update_pointers

get_injection_n_states(inputs::SimulationInputs) = inputs.injection_n_states
get_branches_n_states(inputs::SimulationInputs) = inputs.branches_n_states
get_variable_count(inputs::SimulationInputs) = inputs.variable_count
get_inner_vars_count(inputs::SimulationInputs) = inputs.inner_vars_count
get_ode_ouput_range(inputs::SimulationInputs) = inputs.ode_range
get_bus_count(inputs::SimulationInputs) = inputs.bus_count
get_bus_range(inputs::SimulationInputs) = 1:(2 * inputs.bus_count)

# Utility function not to be used for performance sensitive operations
function get_voltage_buses_ix(inputs::SimulationInputs)
    n_buses = get_bus_count(inputs::SimulationInputs)
    return findall(get_DAE_vector(inputs)[1:n_buses])
end

# Utility function not to be used for performance sensitive operations
function get_current_buses_ix(inputs::SimulationInputs)
    n_buses = get_bus_count(inputs::SimulationInputs)
    return findall(.!get_DAE_vector(inputs)[1:n_buses])
end

"""
SimulationInputs build function for MassMatrixModels
"""
function SimulationInputs(::Type{MassMatrixModel}, sys::PSY.System, frequency_reference)
    return SimulationInputs(sys, frequency_reference)
end

"""
SimulationInputs build function for ResidualModels
"""
function SimulationInputs(::Type{ResidualModel}, sys::PSY.System, frequency_reference)
    return SimulationInputs(sys, frequency_reference)
end

function _wrap_dynamic_injector_data(sys::PSY.System, lookup, injection_start::Int)
    injector_data = get_injectors_with_dynamics(sys)
    isempty(injector_data) && error("System doesn't contain any DynamicInjection devices")
    # TODO: Needs a better container that isn't parametrized on an abstract type
    wrapped_injector = Vector(undef, length(injector_data))
    injection_count = 1
    inner_vars_count = 1
    sys_base_power = PSY.get_base_power(sys)
    sys_base_freq = PSY.get_frequency(sys)
    @assert !isempty(injector_data)
    for (ix, device) in enumerate(injector_data)
        @debug "Wrapping $(PSY.get_name(device))"
        dynamic_device = PSY.get_dynamic_injector(device)
        n_states = PSY.get_n_states(dynamic_device)
        ix_range = range(injection_start, length = n_states)
        ode_range = range(injection_count, length = n_states)
        bus_n = PSY.get_number(PSY.get_bus(device))
        bus_ix = lookup[bus_n]
        inner_vars_range =
            range(inner_vars_count, length = get_inner_vars_count(dynamic_device))
        @debug "ix_range=$ix_range ode_range=$ode_range inner_vars_range= $inner_vars_range"
        wrapped_injector[ix] = DynamicWrapper(
            device,
            bus_ix,
            ix_range,
            ode_range,
            inner_vars_range,
            sys_base_power,
            sys_base_freq,
        )
        injection_count += n_states
        injection_start += n_states
        inner_vars_count =
            length(inner_vars_range) > 0 ? inner_vars_range[end] : inner_vars_count
    end
    return wrapped_injector
end

function _wrap_dynamic_branches(sys::PSY.System, lookup::Dict{Int, Int})
    branches_start = 2 * get_n_buses(sys) + 1
    sys_base_power = PSY.get_base_power(sys)
    sys_base_freq = PSY.get_frequency(sys)
    dynamic_branches = get_dynamic_branches(sys)
    wrapped_branches = Vector{BranchWrapper}(undef, length(dynamic_branches))
    if !isempty(wrapped_branches)
        branches_count = 1
        for (ix, br) in enumerate(dynamic_branches)
            @debug "Wrapping Branch $(PSY.get_name(br))"
            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 = lookup[from_bus_number]
            bus_ix_to = lookup[to_bus_number]
            ix_range = range(branches_start, length = n_states)
            ode_range = range(branches_count, length = n_states)
            @debug "ix_range=$ix_range ode_range=$ode_range"
            wrapped_branches[ix] = BranchWrapper(
                br,
                bus_ix_from,
                bus_ix_to,
                ix_range,
                ode_range,
                sys_base_power,
                sys_base_freq,
            )
            branches_count += n_states
            branches_start += n_states
        end
    else
        @debug("System doesn't contain Dynamic Branches")
    end
    return wrapped_branches
end

function _wrap_static_injectors(sys::PSY.System, lookup::Dict{Int, Int})
    static_injection_data = get_injection_without_dynamics(sys)
    container = Vector{StaticWrapper}(undef, length(static_injection_data))
    for (ix, ld) in enumerate(static_injection_data)
        bus_n = PSY.get_number(PSY.get_bus(ld))
        bus_ix = lookup[bus_n]
        container[ix] = StaticWrapper(ld, bus_ix)
    end
    return container
end

function _wrap_loads(sys::PSY.System, lookup::Dict{Int, Int})
    # This needs to change if we implement dynamic load models
    static_loads =
        PSY.get_components(PSY.ElectricLoad, sys, x -> !isa(x, PSY.FixedAdmittance))
    container = Vector(undef, length(static_loads))
    for (ix, ld) in enumerate(static_loads)
        bus_n = PSY.get_number(PSY.get_bus(ld))
        bus_ix = lookup[bus_n]
        container[ix] = StaticWrapper(ld, bus_ix)
    end
    return container
end

function _get_ybus(sys::PSY.System)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    dyn_lines = PSY.get_components(PSY.DynamicBranch, sys)
    if !isempty(PSY.get_components(PSY.ACBranch, sys))
        Ybus_ = PSY.Ybus(sys)
        ybus = Ybus_[:, :]
        lookup = Ybus_.lookup[1]
        ybus_rectangular = transform_ybus_to_rectangular(ybus)
        for br in dyn_lines
            ybus_update!(ybus_rectangular, br, lookup, -1.0)
        end
    else
        ybus_rectangular =
            SparseArrays.SparseMatrixCSC{Float64, Int}(zeros(2 * n_buses, 2 * n_buses))
        lookup = Dict{Int.Int}()
    end
    return ybus_rectangular, lookup
end

function _static_injection_inputs!(inputs::SimulationInputs, ::Vector{Int}, sys::PSY.System)
    for s in PSY.get_components(PSY.StaticInjection, sys)
        # Temporary condition needs to be removed for other static injections
        PSY.get_dynamic_injector(s) !== nothing && continue
        index_static_injection(inputs, s)
        _attach_control_refs!(s)
    end
    return
end

function _make_mass_matrix(wrapped_injectors, var_count::Int, n_buses::Int)
    mass_matrix = LinearAlgebra.Diagonal(ones(var_count))
    mass_matrix[1:(2 * n_buses), 1:(2 * n_buses)] .= 0.0
    for d in wrapped_injectors
        device_mass_matrix_entries!(mass_matrix, d)
    end
    return mass_matrix
end

function _init_DAE_vector(var_count::Int, n_buses::Int)
    DAE_vector = trues(var_count)
    DAE_vector[1:(n_buses * 2)] .= false
    return DAE_vector
end

function _make_DAE_vector(mass_matrix::AbstractArray, var_count::Int, n_buses::Int)
    DAE_vector = _init_DAE_vector(var_count, n_buses)
    bus_range = 1:(2 * n_buses)
    ode_range = (2 * n_buses + 1):var_count
    mass_buses = @view mass_matrix[bus_range, bus_range]
    DAE_ode = @view DAE_vector[ode_range]
    mass_ode = @view mass_matrix[ode_range, ode_range]
    for i in eachindex(DAE_vector[bus_range])
        IS.@assert_op DAE_vector[bus_range][i] == (mass_buses[i, i] > 0.0)
    end
    for i in eachindex(DAE_ode)
        DAE_ode[i] = (mass_ode[i, i] > 0.0)
    end
    return DAE_vector
end

function _make_total_shunts(wrapped_branches, n_buses::Int)
    shunts = SparseArrays.SparseMatrixCSC{Float64, Int}(zeros(2 * n_buses, 2 * n_buses))
    for br in wrapped_branches
        bus_ix_from = get_bus_ix_from(br)
        bus_ix_to = get_bus_ix_to(br)
        b_from = PSY.get_b(br).from
        b_to = PSY.get_b(br).to
        shunts[bus_ix_from, bus_ix_from + n_buses] += b_from
        shunts[bus_ix_to, bus_ix_to + n_buses] += b_to
        shunts[bus_ix_from + n_buses, bus_ix_from] -= b_from
        shunts[bus_ix_to + n_buses, bus_ix_to] -= b_to
    end
    return shunts
end

function _adjust_states!(
    DAE_vector::BitVector,
    mass_matrix::LinearAlgebra.Diagonal{Float64},
    total_shunts::SparseArrays.SparseMatrixCSC{Float64, Int},
    n_buses::Int,
    sys_f::Float64,
)
    all(iszero.(total_shunts)) && return
    line_constant = 1 / (2.0 * π * sys_f)
    # Takes the upper quadrant of the rectangular shunts matrix
    shunts = LinearAlgebra.diag(total_shunts[1:n_buses, (n_buses + 1):end])
    for (ix, val) in enumerate(shunts)
        if val > 0
            @debug "Found shunt with value $val in bus index $ix"
            mass_matrix[ix, ix] =
                mass_matrix[ix + n_buses, ix + n_buses] = val * line_constant
            DAE_vector[ix] = DAE_vector[ix + n_buses] = true
        end
    end
    return
end

function _make_global_variable_index(
    wrapped_injectors::Vector,
    static_injection_data::Vector,
    frequency_reference::Type{T},
) where {T <: Union{ConstantFrequency, ReferenceBus}}
    global_vars_dict = get_vars_ix()
    global_vars_dict[GLOBAL_VAR_SYS_FREQ_INDEX] = get_frequency_reference!(
        frequency_reference,
        wrapped_injectors,
        static_injection_data,
    )
    return global_vars_dict
end

function get_dynamic_wrapper(inputs::SimulationInputs, name::AbstractString)
    for w in get_dynamic_injectors(inputs)
        if PSY.get_name(w) == name
            return w
        end
    end
    println("DynamicWrapper with name $(name) was not found.")
    return nothing
end

function get_setpoints(inputs::SimulationInputs)
    dic = Dict{String, Dict{String, Float64}}()
    for w in get_dynamic_injectors(inputs)
        dic_w = Dict{String, Float64}()
        dic_w["P_ref"] = get_P_ref(w)
        dic_w["Q_ref"] = get_Q_ref(w)
        dic_w["ω_ref"] = get_ω_ref(w)
        dic_w["V_ref"] = get_V_ref(w)
        dic[PSY.get_name(w)] = dic_w
    end
    return dic
end


#====================== Auxilliary precomputation =============================#
"""
    EvalInfo

EvalInfo is an auxilliary data structure for storing useful computations for the evaluation of the board.
"""
struct EvalInfo
    wrammedpawns::Bitboard
    brammedpawns::Bitboard
    wmobility::Bitboard
    bmobility::Bitboard
    stage::Int
end

"""
    EvalAttackInfo

EvalAttackInfo is an auxilliary data structure for storing useful computations for the evaluation of the board.
"""
mutable struct EvalAttackInfo
    wpawnattacks::Bitboard
    wknightattacks::Bitboard
    wbishopattacks::Bitboard
    wrookattacks::Bitboard
    wqueenattacks::Bitboard
    wkingattacks::Bitboard
    bpawnattacks::Bitboard
    bknightattacks::Bitboard
    bbishopattacks::Bitboard
    brookattacks::Bitboard
    bqueenattacks::Bitboard
    bkingattacks::Bitboard
    passed::Bitboard
end

function initEvalInfo(board::Board)
    # extract pawn positions
    wpawns = white(board) & pawns(board)
    bpawns = black(board) & pawns(board)

    # check for rammed pawns
    wrammedpawns = pawnAdvance(bpawns, wpawns, BLACK)
    brammedpawns = pawnAdvance(wpawns, bpawns, WHITE)

    # generate list of all positions attacked by a pawn
    wpawnattacks = pawnCapturesWhite(wpawns, FULL)
    bpawnattacks = pawnCapturesBlack(bpawns, FULL)

    # generate list of all positions attacked by a knight
    wknightattacks = bknightattacks = EMPTY
    wbishopattacks = bbishopattacks = EMPTY
    wrookattacks = brookattacks = EMPTY
    wqueenattacks = bqueenattacks = EMPTY
    wkingattacks = bkingattacks = EMPTY

    # mobility regions
    wmobility = ~((white(board) & kings(board)) | bpawnattacks)
    bmobility = ~((black(board) & kings(board)) | wpawnattacks)

    passed = EMPTY

    gamestage = stage(board)

    ei = EvalInfo(wrammedpawns, brammedpawns, wmobility, bmobility, gamestage)
    ea = EvalAttackInfo(wpawnattacks, wknightattacks, wbishopattacks, wrookattacks, wqueenattacks, wkingattacks, bpawnattacks, bknightattacks, bbishopattacks, brookattacks, bqueenattacks, bkingattacks, passed)
    ei, ea
end

#============================== Game stage functions ==========================#

function stage(board::Board)
    stage = 6156 - 1024count(queens(board)) - 512count(rooks(board)) - 256count(knights(board) | bishops(board))
    stage = fld(stage, 24)
end

#=============================== Endgame scoring scaling ======================#

function scale_factor(board::Board, eval::Int)
    if isone(board[WHITEBISHOP]) && isone(board[BLACKBISHOP]) && isone(bishops(board) & LIGHT)
        if isempty(knights(board) | rooklike(board))
            return SCALE_OCB_BISHOPS
        end
        if isempty(rooklike(board)) && isone(board[WHITEKNIGHT]) && isone(board[BLACKKNIGHT])
            return SCALE_OCB_ONE_KNIGHT
        end
        if isempty(knights(board) | queens(board)) && isone(board[WHITEROOK]) && isone(board[BLACKROOK])
            return SCALE_OCB_ONE_ROOK
        end
    end
    if (eval > 0) && (count(white(board)) == 2) && ismany(board[WHITEKNIGHT] | board[WHITEBISHOP])
        return SCALE_DRAW
    elseif (eval < 0) && (count(black(board)) == 2) && ismany(board[BLACKKNIGHT] | board[BLACKBISHOP])
        return SCALE_DRAW
    end
    return SCALE_NORMAL
end

#=============================== Main evaluation ==============================#

"""
    evaluate(board)

Naive evaluation function to get the code development going.
"""
function evaluate(board::Board, pktable::PawnKingTable)
    ei, ea = initEvalInfo(board)

    score = 0
    score += board.psqteval

    if (pkt_entry = getPKentry(pktable, board.pkhash)) !== PKT_BLANK
        score += pkt_entry.score
        ea.passed = pkt_entry.passed
    else
        kpscore = 0
        kpscore += evaluate_pawns(board, ea)
        kpscore += evaluate_kingpawns(board, ea)
        storePKentry!(pktable, board.pkhash, ea.passed, kpscore)
        score += kpscore
    end

    v = fld(scoreEG(score) + scoreMG(score), 2)
    v = round(Int, v * (100.0 - Float64(board.halfmovecount)) / 100.0)
    if abs(v) > LAZY_THRESH
        if board.turn == WHITE
            return v
        else
            return -v
        end
    end

    score += evaluate_knights(board, ei, ea)
    score += evaluate_bishops(board, ei, ea)
    score += evaluate_rooks(board, ei, ea)
    score += evaluate_queens(board, ei, ea)
    score += evaluate_kings(board, ea)
    score += evaluate_pins(board)
    score += evaluate_space(board, ea)
    score += evaluate_threats(board, ei, ea)

    score += evaluate_passed(board, ea)

    score += evaluate_initiative(board, ea.passed, score)

    scale_f = scale_factor(board, scoreEG(score))

    eval = (256 - ei.stage) * scoreMG(score) + ei.stage * scoreEG(score) * fld(scale_f, SCALE_NORMAL)
    eval = fld(eval, 256)

    if board.turn == WHITE
        eval += TEMPO_BONUS#
        eval = round(Int, eval * (100.0 - Float64(board.halfmovecount)) / 100.0)
        return eval
    else
        eval -= TEMPO_BONUS
        eval = round(Int, eval * (100.0 - Float64(board.halfmovecount)) / 100.0)
        return -eval
    end
end

function evaluate_pawns(board::Board, ea::EvalAttackInfo)

    w_pawns = white(board) & pawns(board)
    b_pawns = black(board) & pawns(board)

    w_king = square(white(board) & kings(board))
    b_king = square(black(board) & kings(board))

    score = 0

    @inbounds for pawn in w_pawns
        file = fileof(pawn)
        rank = rankof(pawn)
        blocked    = b_pawns & (Bitboard(pawn) << 8)
        stoppers   = b_pawns & PASSED_PAWN_MASKS[1][pawn]
        opposed    = stoppers & FILE[file]
        lever      = b_pawns & PAWN_CAPTURES_WHITE[pawn]
        leverpush  = b_pawns & PAWN_CAPTURES_WHITE[pawn + 8]
        doubled    = !isempty(w_pawns & (Bitboard(pawn) >> 8))
        neighbours = w_pawns & NEIGHBOUR_FILE_MASKS[file]
        phalanx    = neighbours & RANK[rank]
        support    = neighbours & RANK[rank - 1]

        # Passed pawns
        if isempty(stoppers)
            ea.passed |= Bitboard(pawn)
        end

        # Backward pawns
        backward = isempty(neighbours & PASSED_PAWN_MASKS[2][pawn - 8] & FILE[file]) && !isempty(leverpush | blocked)

        # file mapping
        file_psqt = FILE_TO_QSIDE_MAP[file]

        # Connected pawns
        if !isempty(support) || !isempty(phalanx)
            score += CONNECTED_PAWN_PSQT[rank][file_psqt]

        # Isolated pawns
        elseif isempty(neighbours)
            score -= ISOLATED_PAWN_PENALTY + WEAK_UNOPPOSED * isempty(opposed)

        # Backward pawns
        elseif backward
            score -= BACKWARD_PAWN + WEAK_UNOPPOSED * isempty(opposed)
        end

        # Doubled pawns
        if isempty(support)
            score -= DOUBLE_PAWN_PENALTY * doubled + WEAK_LEVER * ismany(lever)
        end
    end

    @inbounds for pawn in b_pawns
        file = fileof(pawn)
        rank = rankof(pawn)
        blocked    = w_pawns & (Bitboard(pawn) >> 8)
        stoppers   = w_pawns & PASSED_PAWN_MASKS[2][pawn]
        opposed    = stoppers & FILE[file]
        lever      = w_pawns & PAWN_CAPTURES_BLACK[pawn]
        leverpush  = w_pawns & PAWN_CAPTURES_BLACK[pawn - 8]
        doubled    = !isempty(b_pawns & (Bitboard(pawn) << 8))
        neighbours = b_pawns & NEIGHBOUR_FILE_MASKS[file]
        phalanx    = neighbours & RANK[rank]
        support    = neighbours & RANK[rank + 1]

        # Passed pawns
        if isempty(stoppers)
            ea.passed |= Bitboard(pawn)
        end

        # Backward pawns
        backward = isempty(neighbours & PASSED_PAWN_MASKS[1][pawn + 8] & FILE[file]) && !isempty(leverpush | blocked)

        # file mapping
        file_psqt = FILE_TO_QSIDE_MAP[file]

        # Connected pawns
        if !isempty(support) || !isempty(phalanx)
            score -= CONNECTED_PAWN_PSQT[9 - rank][file_psqt]

        # Isolated pawns
        elseif isempty(neighbours)
            score += ISOLATED_PAWN_PENALTY + WEAK_UNOPPOSED * isempty(opposed)

        # Backward pawns
        elseif backward
            score += BACKWARD_PAWN + WEAK_UNOPPOSED * isempty(opposed)
        end

        # Doubled pawns
        if isempty(support)
            score += DOUBLE_PAWN_PENALTY * doubled + WEAK_LEVER * ismany(lever)
        end
    end

    if isempty(pawns(board) & KINGFLANK[fileof(w_king)])
        score -= PAWNLESS_FLANK
    end
    if isempty(pawns(board) & KINGFLANK[fileof(b_king)])
        score += PAWNLESS_FLANK
    end

    return score
end

function evaluate_knights(board::Board, ei::EvalInfo, ea::EvalAttackInfo)
    w_knights = white(board) & knights(board)
    b_knights = black(board) & knights(board)

    w_pawns = pawns(board) & white(board)
    b_pawns = pawns(board) & black(board)

    w_king_sqr = square(kings(board) & white(board))
    b_king_sqr = square(kings(board) & black(board))

    score = 0

    w_outposts = (RANK_4 | RANK_5 | RANK_6) & ea.wpawnattacks & ~ea.bpawnattacks
    @inbounds for knight in w_knights
        attacks = knightMoves(knight)
        ea.wknightattacks |= attacks
        score += KNIGHT_MOBILITY[count(attacks & ei.wmobility) + 1]
        score -= MINOR_KING_PROTECTION * DISTANCE_BETWEEN[w_king_sqr, knight]
    end

    # Score knight outposts
    score += count(w_outposts & w_knights) * KNIGHT_OUTPOST_BONUS

    # Score reachable knight outposts
    score += count(w_outposts & ea.wknightattacks & ~white(board)) * KNIGHT_POTENTIAL_OUTPOST_BONUS

    # Score knights behind pawns
    score += count((w_knights << 8) & w_pawns) * PAWN_SHIELD_BONUS

    b_outposts = (RANK_3 | RANK_4 | RANK_5) & ~ea.wpawnattacks & ea.bpawnattacks
    @inbounds for knight in b_knights
        attacks = knightMoves(knight)
        ea.bknightattacks |= attacks
        score -= KNIGHT_MOBILITY[count(attacks & ei.bmobility) + 1]
        score += MINOR_KING_PROTECTION * DISTANCE_BETWEEN[b_king_sqr, knight]
    end

    # Score knight outposts
    score -= count(b_outposts & b_knights) * KNIGHT_OUTPOST_BONUS

    # Score reachable knight outposts
    score -= count(b_outposts & ea.bknightattacks & ~black(board)) * KNIGHT_POTENTIAL_OUTPOST_BONUS

    # Score knights behind pawns
    score -= count((b_knights >> 8) & b_pawns) * PAWN_SHIELD_BONUS

    # bonus for knights in rammed positions
    num2_rammed = count(ei.wrammedpawns)^2
    score += div(count(w_knights) * num2_rammed, 4) * KNIGHT_RAMMED_BONUS
    score -= div(count(b_knights) * num2_rammed, 4) * KNIGHT_RAMMED_BONUS

    score
end

function evaluate_bishops(board::Board, ei::EvalInfo, ea::EvalAttackInfo)
    w_bishops = white(board) & bishops(board)
    b_bishops = black(board) & bishops(board)

    w_pawns = pawns(board) & white(board)
    b_pawns = pawns(board) & black(board)

    w_king_sqr = square(kings(board) & white(board))
    b_king_sqr = square(kings(board) & black(board))

    score = 0

    attacks = EMPTY
    occ = occupied(board)

    w_outposts = (RANK_4 | RANK_5 | RANK_6) & ea.wpawnattacks & ~ea.bpawnattacks
    @inbounds for bishop in w_bishops
        attacks = bishopMoves(bishop, occ)
        ea.wbishopattacks |= attacks
        score += BISHOP_MOBILITY[count(attacks & ei.wmobility) + 1]
        if ismany(attacks & CENTRAL_SQUARES)
            score += BISHOP_CENTRAL_CONTROL
        end
        score -= MINOR_KING_PROTECTION * DISTANCE_BETWEEN[w_king_sqr, bishop]
    end

    # Outpost bonus
    score += count(w_outposts & w_bishops) * BISHOP_OUTPOST_BONUS

    # Add a bonus for being behind a pawn.
    score += count((w_bishops << 8) & w_pawns) * PAWN_SHIELD_BONUS

    b_outposts = (RANK_3 | RANK_4 | RANK_5) & ~ea.wpawnattacks & ea.bpawnattacks
    @inbounds for bishop in b_bishops
        attacks = bishopMoves(bishop, occ)
        ea.bbishopattacks |= attacks
        score -= BISHOP_MOBILITY[count(attacks & ei.bmobility) + 1]
        if ismany(attacks & CENTRAL_SQUARES)
            score -= BISHOP_CENTRAL_CONTROL
        end
        score += MINOR_KING_PROTECTION * DISTANCE_BETWEEN[b_king_sqr, bishop]
    end

    # Outpost bonus
    score -= count(b_outposts & b_bishops) * BISHOP_OUTPOST_BONUS

    # Add a bonus for being behind a pawn
    score -= count((b_bishops >> 8) & b_pawns) * PAWN_SHIELD_BONUS

    # bishop pair
    if count(w_bishops) >= 2
        score += BISHOP_PAIR_BONUS
    end
    if count(b_bishops) >= 2
        score -= BISHOP_PAIR_BONUS
    end

    # penalty for bishops on colour of own pawns
    if !isempty(w_bishops & LIGHT)
        score -= BISHOP_COLOR_PENALTY * count(w_pawns & LIGHT)
        score -= BISHOP_RAMMED_COLOR_PENALTY * count(ei.wrammedpawns & LIGHT)
    end
    if !isempty(w_bishops & DARK)
        score -= BISHOP_COLOR_PENALTY * count(w_pawns & DARK)
        score -= BISHOP_RAMMED_COLOR_PENALTY * count(ei.wrammedpawns & DARK)
    end
    if !isempty(b_bishops & LIGHT)
        score += BISHOP_COLOR_PENALTY * count(b_pawns & LIGHT)
        score += BISHOP_RAMMED_COLOR_PENALTY * count(ei.brammedpawns & LIGHT)
    end
    if !isempty(b_bishops & DARK)
        score += BISHOP_COLOR_PENALTY * count(b_pawns & DARK)
        score += BISHOP_RAMMED_COLOR_PENALTY * count(ei.brammedpawns & DARK)
    end

    score
end

function evaluate_rooks(board::Board, ei::EvalInfo, ea::EvalAttackInfo)
    w_rooks = (white(board) & rooks(board))
    b_rooks = (black(board) & rooks(board))

    score = 0

    occ = occupied(board)
    @inbounds for rook in w_rooks
        rfile = file(rook)
        if isempty(rfile & pawns(board))
            score += ROOK_OPEN_FILE_BONUS
        elseif isempty(rfile & white(board) & pawns(board))
            score += ROOK_SEMIOPEN_FILE_BONUS
        end
        if !isempty(rfile & queens(board))
            score += ROOK_ON_QUEEN_FILE
        end
        attacks = rookMoves(rook, occ)
        ea.wrookattacks |= attacks
        mob_cnt = count(attacks & ei.wmobility) + 1
        score += ROOK_MOBILITY[mob_cnt]
    end

    @inbounds for rook in b_rooks
        rfile = file(rook)
        if isempty(rfile & pawns(board))
            score -= ROOK_OPEN_FILE_BONUS
        elseif isempty(rfile & black(board) & pawns(board))
            score -= ROOK_SEMIOPEN_FILE_BONUS
        end
        if !isempty(rfile & queens(board))
            score -= ROOK_ON_QUEEN_FILE
        end
        attacks = rookMoves(rook, occ)
        ea.brookattacks |= attacks
        mob_cnt = count(attacks & ei.bmobility) + 1
        score -= ROOK_MOBILITY[mob_cnt]
    end

    score
end

function evaluate_queens(board::Board, ei::EvalInfo, ea::EvalAttackInfo)
    w_queens = board[WHITEQUEEN]
    b_queens = board[BLACKQUEEN]

    score = 0

    occ = occupied(board)

    @inbounds for queen in w_queens
        attacks = queenMoves(queen, occ)
        ea.wqueenattacks |= attacks
        score += QUEEN_MOBILITY[count(attacks & ei.wmobility) + 1]
    end

    @inbounds for queen in b_queens
        attacks = queenMoves(queen, occ)
        ea.bqueenattacks |= attacks
        score -= QUEEN_MOBILITY[count(attacks & ei.bmobility) + 1]
    end

    score
end

function evaluate_kings(board::Board, ea::EvalAttackInfo)
    w_king = white(board) & kings(board)
    b_king = black(board) & kings(board)
    w_king_sqr = square(w_king)
    b_king_sqr = square(b_king)
    ea.wkingattacks |= kingMoves(w_king_sqr)
    ea.bkingattacks |= kingMoves(b_king_sqr)

    score = 0

    if cancastlekingside(board, WHITE)
        score += CASTLE_OPTION_BONUS
    end
    if cancastlequeenside(board, WHITE)
        score += CASTLE_OPTION_BONUS
    end
    if cancastlekingside(board, BLACK)
        score -= CASTLE_OPTION_BONUS
    end
    if cancastlequeenside(board, BLACK)
        score -= CASTLE_OPTION_BONUS
    end

    score
end

function evaluate_kingpawns(board::Board, ea::EvalAttackInfo)
    w_king = white(board) & kings(board)
    b_king = black(board) & kings(board)
    w_king_sqr = square(w_king)
    b_king_sqr = square(b_king)

    score = 0

    wkingfile = fileof(w_king_sqr)
    lower = max(1, wkingfile - 1)
    upper = min(8, wkingfile + 1)
    _wpawns = pawns(board) & white(board) & PASSED_PAWN_MASKS[1][w_king_sqr]
    _bpawns = pawns(board) & black(board) & PASSED_PAWN_MASKS[1][w_king_sqr]
    wksqr = rankof(w_king_sqr)
    @inbounds for file in lower:upper
        wpawns = _wpawns & FILE[file]
        if isempty(wpawns)
            w_dist = 8
        else
            w_dist = abs(wksqr - rankof(trailing_zeros(wpawns.val)+1))
        end

        bpawns = _bpawns & FILE[file]
        if isempty(bpawns)
            b_dist = 8
        else
            b_dist = abs(wksqr - rankof(trailing_zeros(bpawns.val)+1))
        end

        if file === wkingfile
            score += KING_SHELTER_ONFILE[file][w_dist]
        else
            score += KING_SHELTER_OFFFILE[file][w_dist]
        end

        blocked = (b_dist !== 8) && (w_dist === b_dist - 1)
        fmap = FILE_TO_QSIDE_MAP[file]
        if blocked
            score += KING_STORM_BLOCKED[fmap][b_dist]
        else
            score += KING_STORM_UNBLOCKED[fmap][b_dist]
        end
    end

    bkingfile = fileof(b_king_sqr)
    lower = max(1, bkingfile - 1)
    upper = min(8, bkingfile + 1)
    _bpawns = pawns(board) & black(board) & PASSED_PAWN_MASKS[2][b_king_sqr]
    _wpawns = pawns(board) & white(board) & PASSED_PAWN_MASKS[2][b_king_sqr]
    bksqr = rankof(b_king_sqr)
    @inbounds for file in lower:upper
        bpawns = _bpawns & FILE[file]
        if isempty(bpawns)
            b_dist = 8
        else
            b_dist = abs(bksqr - rankof(64-leading_zeros(bpawns.val)))
        end

        wpawns = _wpawns & FILE[file]
        if isempty(wpawns)
            w_dist = 8
        else
            w_dist = abs(bksqr - rankof(64-leading_zeros(wpawns.val)))
        end

        if file === bkingfile
            score -= KING_SHELTER_ONFILE[file][b_dist]
        else
            score -= KING_SHELTER_OFFFILE[file][b_dist]
        end

        blocked = (w_dist !== 8) && (b_dist === w_dist - 1)
        fmap = FILE_TO_QSIDE_MAP[file]
        if blocked
            score -= KING_STORM_BLOCKED[fmap][w_dist]
        else
            score -= KING_STORM_UNBLOCKED[fmap][w_dist]
        end
    end

    score
end

function evaluate_pins(board::Board)
    eval = 0

    # switch turn and find all pins
    board.turn = !board.turn
    opp_pinned = findpins(board)
    board.turn = !board.turn

    if board.turn == WHITE
        eval -= count(board.pinned)
        eval += count(opp_pinned)
    else
        eval += count(board.pinned)
        eval -= count(opp_pinned)
    end

    eval *= PIN_BONUS

    # specific additional pin bonus
    eval -= count(pinned(board) & queens(board)) * 30
    eval += count(opp_pinned & queens(board)) * 30
    eval -= count(pinned(board) & rooks(board)) * 10
    eval += count(opp_pinned & rooks(board)) * 10

    makescore(eval, eval)
end

function evaluate_space(board::Board, ea::EvalAttackInfo)
    eval = 0

    w_filter = (RANK_2 | RANK_3 | RANK_4) & CENTERFILES
    b_filter = (RANK_5 | RANK_6 | RANK_7) & CENTERFILES
    w_attacks = ea.wpawnattacks | ea.wknightattacks | ea.wbishopattacks | ea.wrookattacks | ea.wqueenattacks | ea.wkingattacks
    b_attacks = ea.bpawnattacks | ea.bknightattacks | ea.bbishopattacks | ea.brookattacks | ea.bqueenattacks | ea.bkingattacks
    eval += count(w_filter & ~b_attacks)
    eval -= count(b_filter & ~w_attacks)
    eval *= SPACE_BONUS

    makescore(eval, eval)
end

function evaluate_threats(board::Board, ei::EvalInfo, ea::EvalAttackInfo)
    w_attacks = ea.wpawnattacks
    w_double_attacks = w_attacks & ea.wknightattacks
    w_attacks |= ea.wknightattacks
    w_double_attacks |= w_attacks & ea.wbishopattacks
    w_attacks |= ea.wbishopattacks
    w_double_attacks |= w_attacks & ea.wrookattacks
    w_attacks |= ea.wrookattacks
    w_double_attacks |= w_attacks & ea.wqueenattacks
    w_attacks |= ea.wqueenattacks
    w_double_attacks |= w_attacks & ea.wkingattacks
    w_attacks |= ea.wkingattacks

    b_attacks = ea.bpawnattacks
    b_double_attacks = b_attacks & ea.bknightattacks
    b_attacks |= ea.bknightattacks
    b_double_attacks |= b_attacks & ea.bbishopattacks
    b_attacks |= ea.bbishopattacks
    b_double_attacks |= b_attacks & ea.brookattacks
    b_attacks |= ea.brookattacks
    b_double_attacks |= b_attacks & ea.bqueenattacks
    b_attacks |= ea.bqueenattacks
    b_double_attacks |= b_attacks & ea.bkingattacks
    b_attacks |= ea.bkingattacks

    occ = occupied(board)

    score = 0

    #========================= Evaluation w.r.t. white ========================#
    # strongly protected by the enemy.
    strongly_protected = ea.bpawnattacks | (b_double_attacks & ~w_double_attacks)
    # well defended by the enemy
    defended = (black(board) & ~pawns(board)) & strongly_protected
    # not well defended by the enemy
    weak = black(board) & ~strongly_protected & w_attacks

    # Case where our opponent is hanging pieces
    case = ~b_attacks | ((black(board) & ~pawns(board)) & w_double_attacks)
    # Bonus if opponent is hanging pieces
    score += HANGING_BONUS * count(weak & case)

    if !isempty(weak & ea.wkingattacks)
        score += THREAT_BY_KING
    end

    # Case where our opponent is defended or weak, and attacked by a bishop or knight.
    case = (defended | weak) & (ea.wknightattacks | ea.wbishopattacks)
    case_pawns = count(case & pawns(board))
    case_knights = count(case & knights(board))
    case_bishops = count(case & bishops(board))
    case_rooks = count(case & rooks(board))
    case_queens = count(case & queens(board))
    @inbounds score += THREAT_BY_MINOR[1] * case_pawns
    @inbounds score += THREAT_BY_MINOR[2] * case_knights
    @inbounds score += THREAT_BY_MINOR[3] * case_bishops
    @inbounds score += THREAT_BY_MINOR[4] * case_rooks
    @inbounds score += THREAT_BY_MINOR[5] * case_queens

    # Case where our opponent is weak and attacked by our rook
    case = weak & ea.wrookattacks
    case_pawns = count(case & pawns(board))
    case_knights = count(case & knights(board))
    case_bishops = count(case & bishops(board))
    case_rooks = count(case & rooks(board))
    case_queens = count(case & queens(board))
    @inbounds score += THREAT_BY_ROOK[1] * case_pawns
    @inbounds score += THREAT_BY_ROOK[2] * case_knights
    @inbounds score += THREAT_BY_ROOK[3] * case_bishops
    @inbounds score += THREAT_BY_ROOK[4] * case_rooks
    @inbounds score += THREAT_BY_ROOK[5] * case_queens

    # Case where we get a bonus for restricting our opponent
    case = b_attacks & ~strongly_protected & w_attacks
    score += count(case) * RESTRICTION_BONUS

    # Find squares which are safe
    safe = ~b_attacks | w_attacks

    # Case where we have safe pawn attacks
    case = pawns(board) & white(board) & safe
    case = pawnCapturesWhite(case, black(board) & ~pawns(board))
    score += THREAT_BY_PAWN * count(case)

    # Case threat by pawn push
    case = pawnAdvance(white(board) & pawns(board), empty(board), WHITE)
    case |= pawnDoubleAdvance(white(board) & pawns(board), empty(board), WHITE)
    case &= ~ea.bpawnattacks & safe
    case = pawnCapturesWhite(case, black(board) & ~pawns(board))
    score += THREAT_BY_PUSH * count(case)

    # Evaluate if there are pieces that can reach a checking square, while being safe.
    # King Danger evaluations.
    # As WHITE, we are seeing how in danger the BLACK king is.
    b_king_sqr = square(black(board) & kings(board))
    b_king_box = ea.bkingattacks | Bitboard(b_king_sqr)
    safe = ~white(board) & (~b_attacks | (weak & w_double_attacks))
    king_bishop_moves = bishopMoves(b_king_sqr, occ)
    king_rook_moves = rookMoves(b_king_sqr, occ)
    king_knight_moves = knightMoves(b_king_sqr)
    knightcheck_sqrs = king_knight_moves & safe & ea.wknightattacks
    rookcheck_sqrs   = king_rook_moves & safe & ea.wrookattacks
    # Remove rook checks from queen checks, because they are more valuable.
    # Don't count if queens trade.
    queencheck_sqrs  = (king_rook_moves | king_bishop_moves) & safe & ea.wqueenattacks & ~rookcheck_sqrs & ~ea.bqueenattacks
    # Remove queen checks from bishop checks, because they are more valuable.
    bishopcheck_sqrs = king_bishop_moves & safe & ea.wbishopattacks & ~queencheck_sqrs

    king_danger = 0
    unsafechecks = EMPTY
    if !isempty(knightcheck_sqrs)
        king_danger += KNIGHT_SAFE_CHECK
    else
        unsafechecks |= king_knight_moves & ea.wknightattacks
    end
    if !isempty(bishopcheck_sqrs)
        king_danger += BISHOP_SAFE_CHECK
    else
        unsafechecks |= king_bishop_moves & ea.wbishopattacks
    end
    if !isempty(rookcheck_sqrs)
        king_danger += ROOK_SAFE_CHECK
    else
        unsafechecks |= king_rook_moves & ea.wrookattacks
    end
    if !isempty(queencheck_sqrs)
        king_danger += QUEEN_SAFE_CHECK
    end
    if isempty(board[WHITEQUEEN])
        king_danger -= KD_QUEEN
    end
    if isempty(board[WHITEROOK])
        king_danger -= KD_ROOK
    end
    if isempty(board[WHITEBISHOP])
        king_danger -= KD_BISHOP
    end
    if isempty(board[WHITEKNIGHT])
        king_danger -= KD_KNIGHT
    end
    king_danger += 9 * count(w_attacks & ea.bkingattacks)
    if king_danger > 0
        score += makescore(king_danger, fld(king_danger, 2))
    end

    # Score the number of attacks on the king's flank
    camp = FULL ⊻ RANK_1 ⊻ RANK_2 ⊻ RANK_3
    black_king_flank = KINGFLANK[fileof(b_king_sqr)]
    king_flank_attacks = w_attacks & black_king_flank & camp
    king_flank_double_attacks = king_flank_attacks & w_double_attacks
    score += (count(king_flank_attacks) + count(king_flank_double_attacks)) * KING_FLANK_ATTACK
    king_flank_defence = b_attacks & black_king_flank & camp
    score -= count(king_flank_defence) * KING_FLANK_DEFEND

    # King ring attack bonus
    score += count(b_king_box & weak)^2 * KING_BOX_WEAK
    score += count(unsafechecks)^2 * UNSAFE_CHECK

    # Threats on black queen
    black_queens = black(board) & queens(board)
    if isone(black_queens)
        safe = ei.wmobility & ~strongly_protected
        case = ea.wknightattacks & knightMoves(square(black_queens)) & safe
        score += count(case) * KNIGHT_ON_QUEEN
    end

    #========================= Evaluation w.r.t. black ========================#

    # strongly protected by the enemy.
    strongly_protected = ea.wpawnattacks | (w_double_attacks & ~b_double_attacks)
    # well defended by the enemy
    defended = (white(board) & ~pawns(board)) & strongly_protected
    # not well defended by the enemy
    weak = white(board) & ~strongly_protected & b_attacks

    # Case where our opponent is hanging pieces
    case = ~w_attacks | ((white(board) & ~pawns(board)) & b_double_attacks)
    # Bonus if opponent is hanging pieces
    score -= HANGING_BONUS * count(weak & case)

    if !isempty(weak & ea.bkingattacks)
        score -= THREAT_BY_KING
    end

    # Case where our opponent is defended or weak, and attacked by a bishop or knight.
    case = (defended | weak) & (ea.bknightattacks | ea.bbishopattacks)
    case_pawns = count(case & pawns(board))
    case_knights = count(case & knights(board))
    case_bishops = count(case & bishops(board))
    case_rooks = count(case & rooks(board))
    case_queens = count(case & queens(board))
    @inbounds score -= THREAT_BY_MINOR[1] * case_pawns
    @inbounds score -= THREAT_BY_MINOR[2] * case_knights
    @inbounds score -= THREAT_BY_MINOR[3] * case_bishops
    @inbounds score -= THREAT_BY_MINOR[4] * case_rooks
    @inbounds score -= THREAT_BY_MINOR[5] * case_queens

    # Case where our opponent is weak and attacked by our rook
    case = weak & ea.brookattacks
    case_pawns = count(case & pawns(board))
    case_knights = count(case & knights(board))
    case_bishops = count(case & bishops(board))
    case_rooks = count(case & rooks(board))
    case_queens = count(case & queens(board))
    @inbounds score -= THREAT_BY_ROOK[1] * case_pawns
    @inbounds score -= THREAT_BY_ROOK[2] * case_knights
    @inbounds score -= THREAT_BY_ROOK[3] * case_bishops
    @inbounds score -= THREAT_BY_ROOK[4] * case_rooks
    @inbounds score -= THREAT_BY_ROOK[5] * case_queens

    # Case where we get a bonus for restricting our opponent
    case = b_attacks & ~strongly_protected & w_attacks
    score -= count(case) * RESTRICTION_BONUS

    # Case threat by pawn
    safe = ~w_attacks | b_attacks
    case = pawns(board) & black(board) & safe
    case = pawnCapturesBlack(case, white(board) & ~pawns(board))
    score -= THREAT_BY_PAWN * count(case)

    # Case threat by pawn push
    case = pawnAdvance(black(board) & pawns(board), empty(board), BLACK)
    case |= pawnDoubleAdvance(black(board) & pawns(board), empty(board), BLACK)
    case &= ~ea.wpawnattacks & safe
    case = pawnCapturesBlack(case, white(board) & ~pawns(board))
    score -= THREAT_BY_PUSH * count(case)

    # Evaluate if there are pieces that can reach a checking square, while being safe.
    # King Danger evaluations.
    # As BLACK, we are seeing how in danger the WHITE king is.
    w_king_sqr = square(white(board) & kings(board))
    w_king_box = ea.wkingattacks | Bitboard(w_king_sqr)
    safe = ~black(board) & (~w_attacks | (weak & b_double_attacks))
    king_bishop_moves = bishopMoves(w_king_sqr, occ)
    king_rook_moves = rookMoves(w_king_sqr, occ)
    king_knight_moves = knightMoves(w_king_sqr)
    knightcheck_sqrs = king_knight_moves & safe & ea.bknightattacks
    rookcheck_sqrs   = king_rook_moves & safe & ea.brookattacks
    # Remove rook checks from queen checks, because they are more valuable.
    # Don't count if queens trade.
    queencheck_sqrs  = (king_rook_moves | king_bishop_moves) & safe & ea.bqueenattacks & ~rookcheck_sqrs & ~ea.wqueenattacks
    # Remove queen checks from bishop checks, because they are more valuable.
    bishopcheck_sqrs = king_bishop_moves & safe & ea.bbishopattacks & ~queencheck_sqrs

    king_danger = 0
    unsafechecks = EMPTY
    if !isempty(knightcheck_sqrs)
        king_danger += KNIGHT_SAFE_CHECK
    else
        unsafechecks |= king_knight_moves & ea.bknightattacks
    end
    if !isempty(bishopcheck_sqrs)
        king_danger += BISHOP_SAFE_CHECK
    else
        unsafechecks |= king_bishop_moves & ea.bbishopattacks
    end
    if !isempty(rookcheck_sqrs)
        king_danger += ROOK_SAFE_CHECK
    else
        unsafechecks |= king_rook_moves & ea.brookattacks
    end
    if !isempty(queencheck_sqrs)
        king_danger += QUEEN_SAFE_CHECK
    end
    if isempty(board[BLACKQUEEN])
        king_danger -= KD_QUEEN
    end
    if isempty(board[BLACKROOK])
        king_danger -= KD_ROOK
    end
    if isempty(board[BLACKBISHOP])
        king_danger -= KD_BISHOP
    end
    if isempty(board[BLACKKNIGHT])
        king_danger -= KD_KNIGHT
    end
    king_danger += 9 * count(b_attacks & ea.wkingattacks)
    if king_danger > 0
        score -= makescore(king_danger, fld(king_danger, 2))
    end

    # Count king flank attacks
    camp = (FULL ⊻ RANK_8 ⊻ RANK_7 ⊻ RANK_6)
    white_king_flank = KINGFLANK[fileof(w_king_sqr)]
    king_flank_attacks = b_attacks & white_king_flank & camp
    king_flank_double_attacks = king_flank_attacks & b_double_attacks
    score -= (count(king_flank_attacks) + count(king_flank_double_attacks)) * KING_FLANK_ATTACK
    king_flank_defence = w_attacks & white_king_flank & camp
    score += count(king_flank_defence) * KING_FLANK_DEFEND

    # King ring attack bonus
    score -= count(w_king_box & weak)^2 * KING_BOX_WEAK
    score -= count(unsafechecks)^2 * UNSAFE_CHECK

    # Attacks on white queen
    white_queens = white(board) & queens(board)
    if isone(white_queens)
        safe = ei.bmobility & ~strongly_protected
        case = ea.bknightattacks & knightMoves(square(white_queens)) & safe
        score -= count(case) * KNIGHT_ON_QUEEN
    end

    score
end

function evaluate_initiative(board::Board, passed::Bitboard, score::Int)
    eg = scoreEG(score)
    mg = scoreMG(score)
    black_king = square(kings(board) & black(board))
    white_king = square(kings(board) & white(board))
    black_king_rank = rankof(black_king)
    white_king_rank = rankof(white_king)
    black_king_file = fileof(black_king)
    white_king_file = fileof(white_king)

    outflank = abs(white_king_file - black_king_file) - abs(white_king_rank - black_king_rank)

    infiltrating = (white_king_rank > 4) || (black_king_rank < 5)

    pawns_both_flanks = !isempty(pawns(board) & KINGSIDE) || !isempty(pawns(board) & QUEENSIDE)

    slimchance = isempty(passed) && (outflank < 0) && !pawns_both_flanks

    complexity = 8 * count(passed) + 10 * count(pawns(board)) + 8 * outflank +
        (infiltrating ? 11 : 0) + (pawns_both_flanks ? 19 : 0) + (has_non_pawn_material(board) ? 0 : 44) -
        (slimchance ? 39 : 0) - 90

    c_mg = (Int(mg > 0) - Int(mg < 0)) * clamp(complexity+45, -abs(mg), 0)
    c_eg = (Int(eg > 0) - Int(eg < 0)) * max(complexity, -abs(eg))

    makescore(c_mg, c_eg)
end

function evaluate_passed(board::Board, ea::EvalAttackInfo)
    black_king = square(kings(board) & black(board))
    white_king = square(kings(board) & white(board))
    white_pawns = white(board) & ea.passed
    black_pawns = black(board) & ea.passed

    w_attacks = ea.wknightattacks | ea.wpawnattacks | ea.wbishopattacks | ea.wrookattacks | ea.wqueenattacks | ea.wkingattacks
    b_attacks = ea.bknightattacks | ea.bpawnattacks | ea.bbishopattacks | ea.brookattacks | ea.bqueenattacks | ea.bkingattacks

    block_sqr = 0

    score = 0
    @inbounds for pawn in white_pawns
        bonus = 0
        file = fileof(pawn)
        rank = rankof(pawn)
        bonus += PASS_PAWN_THREAT[rank]
        if rank > 3
            factor = muladd(5, rank, -13)
            block_sqr = pawn + 8
            bonus += makescore(0, (fld(min(DISTANCE_BETWEEN[block_sqr, black_king],5)*19, 4) - min(DISTANCE_BETWEEN[block_sqr, white_king],5)*2)*factor)
            if rank !== 7
                block_sqr_2 = block_sqr + 8
                bonus -= makescore(0, DISTANCE_BETWEEN[block_sqr_2, white_king]*factor)
            end
            if board[block_sqr] === BLANK
                sqrs_to_queen = PASSED_PAWN_MASKS[1][pawn] & FILE[file]
                unsafe_sqrs = PASSED_PAWN_MASKS[1][pawn]
                case = (PASSED_PAWN_MASKS[2][pawn] & FILE[file]) & rooklike(board)
                if isempty(case & black(board))
                    unsafe_sqrs &= b_attacks
                end

                if isempty(unsafe_sqrs)
                    k = 32
                elseif isempty(unsafe_sqrs & sqrs_to_queen)
                    k = 18
                elseif isempty(unsafe_sqrs & Bitboard(block_sqr))
                    k = 8
                else
                    k = 0
                end

                if !isempty(white(board) & case) || !isempty(w_attacks & Bitboard(block_sqr))
                    k += 5
                end

                bonus += makescore(k * factor, k * factor)
            end
        end
        if board[block_sqr] === WHITEPAWN
            bonus = makescore(fld(scoreMG(bonus), 2), fld(scoreEG(bonus), 2))
        end
        score += bonus
    end

    @inbounds for pawn in black_pawns
        bonus = 0
        file = fileof(pawn)
        rank = 9 - rankof(pawn)
        bonus += PASS_PAWN_THREAT[rank]
        if rank > 3
            factor = muladd(5, rank, -13)
            block_sqr = pawn - 8
            bonus += makescore(0, (fld(min(DISTANCE_BETWEEN[block_sqr, white_king],5)*14, 3) - min(DISTANCE_BETWEEN[block_sqr, black_king],5)*2)*factor)
            if rank !== 7
                block_sqr_2 = block_sqr - 8
                bonus -= makescore(0, DISTANCE_BETWEEN[block_sqr_2, black_king]*factor)
            end
            if board[block_sqr] === BLANK
                sqrs_to_queen = PASSED_PAWN_MASKS[2][pawn] & FILE[file]
                unsafe_sqrs = PASSED_PAWN_MASKS[2][pawn]
                case = (PASSED_PAWN_MASKS[1][pawn] & FILE[file]) & rooklike(board)
                if isempty(case & white(board))
                    unsafe_sqrs &= w_attacks
                end

                if isempty(unsafe_sqrs)
                    k = 32
                elseif isempty(unsafe_sqrs & sqrs_to_queen)
                    k = 18
                elseif isempty(unsafe_sqrs & Bitboard(block_sqr))
                    k = 8
                else
                    k = 0
                end

                if !isempty(black(board) & case) || !isempty(b_attacks & Bitboard(block_sqr))
                    k += 5
                end

                bonus += makescore(k * factor, k * factor)
            end
        end
        if board[block_sqr] === BLACKPAWN
            bonus = makescore(fld(scoreMG(bonus), 2), fld(scoreEG(bonus), 2))
        end
        score -= bonus
    end

    score
end


module SKCrossValidators

using PyCall

# standard included modules
using DataFrames
using Random
using AutoMLPipeline.AbsTypes
using AutoMLPipeline.Utils

import AutoMLPipeline.CrossValidators.crossvalidate
export crossvalidate

function __init__()
    global SKM = pyimport("sklearn.metrics")
    global metric_dict = Dict(
          "roc_auc_score" => SKM.roc_auc_score,
          "accuracy_score" => SKM.accuracy_score,
          "auc" => SKM.auc,
          "average_precision_score" => SKM.average_precision_score,
          "balanced_accuracy_score" => SKM.balanced_accuracy_score,
          "brier_score_loss" => SKM.brier_score_loss,
          "classification_report" => SKM.classification_report,
          "cohen_kappa_score" => SKM.cohen_kappa_score,
          "confusion_matrix" => SKM.confusion_matrix,
          "f1_score" => SKM.f1_score,
          "fbeta_score" => SKM.fbeta_score,
          "hamming_loss" => SKM.hamming_loss,
          "hinge_loss" => SKM.hinge_loss,
          "log_loss" => SKM.log_loss,
          "matthews_corrcoef" => SKM.matthews_corrcoef,
          "multilabel_confusion_matrix" => SKM.multilabel_confusion_matrix,
          "precision_recall_curve" => SKM.precision_recall_curve,
          "precision_recall_fscore_support" => SKM.precision_recall_fscore_support,
          "precision_score" => SKM.precision_score,
          "recall_score" => SKM.recall_score,
          "roc_auc_score" => SKM.roc_auc_score,
          "roc_curve" => SKM.roc_curve,
          "jaccard_score" => SKM.jaccard_score,
          "zero_one_loss" => SKM.zero_one_loss
         )
end

function crossvalidate(pl::Machine,X::DataFrame,Y::Vector,
                       sfunc::String,nfolds=10)
    @assert sfunc in keys(metric_dict)
    pfunc = metric_dict[sfunc]
    metric(X,Y) = pfunc(X,Y)
    crossvalidate(pl,X,Y,metric,nfolds)
end

function crossvalidate(pl::Machine,X::DataFrame,Y::Vector,
                       sfunc::String,averagetype::String,nfolds=10)
    @assert sfunc in keys(metric_dict)
    pfunc = metric_dict[sfunc]
    metric(X,Y) = pfunc(X,Y,average=averagetype)
    crossvalidate(pl,X,Y,metric,nfolds)
end

end


function setticks(y::Real)
    if y == 10
        return 0:3:10
    elseif 10 < y <= 20
        return 0:5:y
    elseif 20 < y <= 50
        return 0:10:y
    else
        s = div(y, 4)
        d, r = divrem(s, 10)
        return r < 5 ? (0:10d:y) : (0:((d + 1) * 10):y)
    end
end

function calculateld!(
    gwas::DataFrame,
    ref::SnpData;
    snp::Union{AbstractString, Tuple{AbstractString, Int}} = "index"
)

    gwas.ind = findsnps(gwas, ref; matchalleles = false)
    dropmissing!(gwas, "ind")
    n = size(gwas, 1)
    gwas.LD = fill(0.0, n)
    if snp == "index"
        i = argmax(gwas.P)
        snp = gwas.SNP[i]
        geno = convert(Matrix{Float64}, ref.snparray[:, gwas.ind])
        for j in 1:n
            gwas.LD[j] = cor(geno[:, i], geno[:, j])^2
        end
        gwas.index = fill(snp, n)
        return
    else
        if snp isa AbstractString
            (chr, bp) = getsnpinfo(snp, ref)
        else
            (chr, bp) = snp
        end
        i = findfirst((gwas.CHR .== chr) .& (gwas.BP .== bp))
        if isnothing(i)
            snp = ""
            i = findfirst((ref.snp_info.chromosome .== chr) .& (ref.snp_info.position .== bp))
            geno = convert(Matrix{Float64}, ref.snparray[:, [i; gwas.ind]])
            for j in 1:n
                gwas.LD[j] = cor(geno[:, 1], geno[:, j + 1])^2
            end
        else
            snp = gwas.SNP[i]
            geno = convert(Matrix{Float64}, ref.snparray[:, gwas.ind])
            for j in 1:n
                gwas.LD[j] = cor(geno[:, i], geno[:, j])^2
            end
        end
        gwas.index = fill(snp, n)
        return
    end
end

"""
    plotlocus!(ax::Axis, chromosome::AbstractString, range1::Real, range2::Real, gwas::DataFrame; kwargs)
    plotlocus!(ax::Axis, chromosome::AbstractString, bp::Real, gwas::DataFrame; kwargs)
    plotlocus!(ax::Axis, gene::AbstractString, gwas::DataFrame, gencode::DataFrame; kwargs)

Plot `gwas` results within a given `chromosome` and genomic range between `range1`
and `range2`.

Alternatively, plot within a given `chromosome` and a certain `window` around a
genomic coordinate `bp` or plot within a certain `window` around `gene`.

# Arguments
- `ld::Union{Nothing, SnpData, Tuple{SnpData, Union{AbstractString, Tuple{AbstractString, Int}}}} = nothing`:
    the reference panel for which LD is calculated.
- `ymax::Real`: the maximum value for y axis.
- `window::Real = 1e6`: the window around `bp` or `gene`.
- `showsnplabel::Bool = true`: draw SNP ID text label.
"""
function plotlocus!(
    ax::Axis,
    chromosome::AbstractString,
    range1::Real,
    range2::Real,
    gwas::DataFrame;
    ld::Union{Nothing, SnpData, Tuple{SnpData, Union{AbstractString, Tuple{AbstractString, Int}}}} = nothing,
    ymax::Real = 0,
    showsnplabel::Bool = true
)

    df = filter(x -> (x.CHR == chromosome) && (x.BP >= range1) && (x.BP <= range2), gwas)
    if nrow(df) == 0
        ymax == 0 ? ymax = 10 : nothing
        ax.spinewidth = 0.75
        ax.ytickwidth = 0.75
        ax.ylabelsize = 6
        ax.yticklabelsize = 6
        ax.yticksize = 3
        ax.yticks = setticks(ymax)
        xlims!(ax, range1, range2)
        ylims!(ax, 0, ymax)
        hidespines!(ax, :t, :r)
        hidexdecorations!(ax)
        hideydecorations!(ax, ticks = false, label = false, ticklabels = false)
        return
    end
    df.P = -log.(10, df.P)
    if ymax == 0
        ymax = maximum(df.P) / 4 * 5
        ymax <= 10 ? ymax = 10 : nothing
    end
    if !isnothing(ld)
        typeof(ld) == SnpData ? calculateld!(df, ld) : calculateld!(df, ld[1]; snp = ld[2])
        scatter!(ax, df.BP, df.P, color = df.LD, colorrange = (0, 1),
            colormap = (:gray60, :red2), markersize = 1.5)
        if typeof(ld) == SnpData
            ind = argmax(df.P)
            bp = df.BP[ind]
            p = df.P[ind]
            scatter!(ax, [bp], [p], color = :purple1, markersize = 4.0, marker = '◆')
            if showsnplabel
                text!(ax, "$(df.index[1])", position = (bp, p), textsize = 6, align = (:center, :bottom))
            end
        elseif length(df.index[1]) > 0
            ind = findfirst(df.SNP .== df.index[1])
            bp, p = df.BP[ind], df.P[ind]
            scatter!(ax, [bp], [p], color = :purple1, markersize = 4.0, marker = '◆')
            if showsnplabel
                text!(ax, "$(df.index[1])", position = (bp, p), textsize = 6, align = (:center, :bottom))
            end
        end
    else
        scatter!(ax, df.BP, df.P, color = :gray60, markersize = 1.5)
    end
    ax.spinewidth = 0.75
    ax.ytickwidth = 0.75
    ax.ylabelsize = 6
    ax.yticklabelsize = 6
    ax.yticksize = 3
    ax.yticks = setticks(ymax)
    xlims!(ax, range1, range2)
    ylims!(ax, 0, ymax)
    hidespines!(ax, :t, :r)
    hidexdecorations!(ax)
    hideydecorations!(ax, ticks = false, label = false, ticklabels = false)
end

plotlocus!(ax::Axis, chromosome::AbstractString, bp::Real, gwas::DataFrame; window::Real = 1e6, kwargs...) =
    plotlocus!(ax, chromosome, bp - window, bp + window, gwas; kwargs...)

function plotlocus!(ax::Axis, gene::AbstractString, gwas::DataFrame, gencode::DataFrame; window::Real = 1e6, kwargs...)
    chr, start, stop = findgene(gene, gencode)
    plotlocus!(ax, chr, start - window, stop + window, gwas; kwargs...)
end