"""
    AbstractΔSizeStrategy

Abstract base type for strategies for how to change the size.

Only used as a transition until JuMP approach has been fully verified.
"""
abstract type AbstractΔSizeStrategy end

"""
    OnlyFor <: AbstractΔSizeStrategy

Change size only for the provided vertex.
"""
struct OnlyFor <: AbstractΔSizeStrategy end

"""
    AbstractJuMPΔSizeStrategy <: AbstractΔSizeStrategy

Abstract type for strategies to change or align the sizes of vertices using JuMP.
"""
abstract type AbstractJuMPΔSizeStrategy <: AbstractΔSizeStrategy end

"""
    ΔSizeFailError <: AbstractJuMPΔSizeStrategy
    ΔSizeFailError(msg::String)

Throws an `ErrorException` with message `msg`.
"""
struct ΔSizeFailError <: AbstractJuMPΔSizeStrategy
    msg::String
end

"""
    ΔSizeFailNoOp <: AbstractJuMPΔSizeStrategy
    ΔSizeFailNoOp()

Does not perform any action.
"""
struct ΔSizeFailNoOp <: AbstractJuMPΔSizeStrategy end

"""
    LogΔSizeExec <: AbstractJuMPΔSizeStrategy
    LogΔSizeExec(msg::String)
    LogΔSizeExec(level::Logging.LogLevel, msg::String)
    LogΔSizeExec(level::Logging.LogLevel, msg::String, andthen::AbstractJuMPΔSizeStrategy)

Logs `msg` at log level `level`, then executes `AbstractJuMPΔSizeStrategy andthen`.
"""
struct LogΔSizeExec <: AbstractJuMPΔSizeStrategy
    level::LogLevel
    msg::String
    andthen::AbstractJuMPΔSizeStrategy
end
LogΔSizeExec(msg::String) = LogΔSizeExec(Logging.Info, msg)
LogΔSizeExec(level::LogLevel, msg::String) = LogΔSizeExec(level, msg, ΔSizeFailNoOp())

"""
    DefaultJuMPΔSizeStrategy <: AbstractJuMPΔSizeStrategy

Default strategy intended to be used when adding some extra constraints or objectives to a model on top of the default.
"""
struct DefaultJuMPΔSizeStrategy <: AbstractJuMPΔSizeStrategy end

struct Exact end
struct Relaxed end

"""
    ΔNout{T} <: AbstractJuMPΔSizeStrategy
    ΔNout{T}(vertex::AbstractVertex, Δ::Integer)
    ΔNoutExact(vertex::AbstractVertex, Δ::Integer, fallback::AbstractJuMPΔSizeStrategy)
    ΔNoutRelaxed(vertex::AbstractVertex, Δ::Integer, fallback::AbstractJuMPΔSizeStrategy)

Strategy for changing nout of `vertex` by `Δ`, i.e new size is `nout(vertex) + Δ`.

If `T == Exact`, size change will be added as a constraint to the model which means that the operation will fail if it is not possible to change `nout(vertex)` by exactly `Δ`. If the operation fails, it will be retried with the `fallback` strategy (default `ΔNoutRelaxed`).

If `T == Relaxed`, size change will be added as an objective to the model which means that `nout(vertex)` might not change by exactly `Δ`. In addition, a constraint that `nout(vertex)` must change is also added.

If the operation fails, it will be retried with the `fallback` strategy (default `ΔNout{Relaxed}` if `T==Exact` and `ΔSizeFailError` if `T==Relaxed`).
"""
struct ΔNout{T} <: AbstractJuMPΔSizeStrategy
    vertex::AbstractVertex
    Δ::Integer
    fallback::AbstractJuMPΔSizeStrategy
end
ΔNoutExact(v::AbstractVertex, Δ::Integer) = ΔNout{Exact}(v, Δ, LogΔSizeExec(Logging.Warn, "Could not change nout of $v by $(Δ)! Relaxing constraints...", ΔNoutRelaxed(v, Δ)))
ΔNoutRelaxed(v::AbstractVertex, Δ::Integer) = ΔNout{Relaxed}(v, Δ, ΔSizeFailError("Could not change nout of $v by $(Δ)!!"))
fallback(s::ΔNout) = s.fallback

"""
    ΔNin{T} <: AbstractJuMPΔSizeStrategy
    ΔNin{T}(vertex::AbstractVertex, Δs::Vector{Maybe{Int}}, fallback::AbstractJuMPΔSizeStrategy)
    ΔNinExact(vertex::AbstractVertex, Δs::Vector{Maybe{Int}})
    ΔNinRelaxed(vertex::AbstractVertex, Δs::Vector{Maybe{Int}})

Strategy for changing nin of `vertex` by `Δs`, i.e new size is `nin(vertex) .+ Δs`. Note that `Δs` must have the same number of elements as `nin(vertex)`.

Use `missing` to indicate "no change required" as 0 will be interpreted as "must not change".

If `T == Exact`, size change will be added as a constraint to the model which means that the operation will fail if it is not possible to change `nin(vertex)` by exactly `Δs`.

If `T == Relaxed`, size change will be added as an objective to the model which means that `nin(vertex)` might not change by exactly `Δs`. In addition, a constraint that `nin(vertex)` must change is also added.

If the operation fails, it will be retried with the `fallback` strategy (default `ΔNin{Relaxed}` if `T==Exact` and `ΔSizeFailError` if `T==Relaxed`).
"""
struct ΔNin{T} <: AbstractJuMPΔSizeStrategy
    vertices::Vector{<:AbstractVertex}
    Δs::Vector{Int}
    fallback::AbstractJuMPΔSizeStrategy
    function ΔNin{T}(v, Δs, fallback) where T
        @assert size(Δs) == size(inputs(v)) "Must supply same number of Δs as v has inputs! Got $Δs for $v."
        inds = .!ismissing.(Δs)
        new(inputs(v)[inds], Δs[inds], fallback)
    end
end
ΔNinExact(v::AbstractVertex, Δs::Vector{<:Maybe{Int}}) = ΔNin{Exact}(v, Δs, LogΔSizeExec(Logging.Warn, "Could not change nin of $v by $(join(Δs, ", "))! Relaxing constraints...", ΔNinRelaxed(v, Δs)))
ΔNinExact(v::AbstractVertex, Δ::Integer) = ΔNin{Exact}(v, [Δ])
ΔNinRelaxed(v::AbstractVertex, Δs::Vector{<:Maybe{Int}}) = ΔNin{Relaxed}(v, Δs, ΔSizeFailError("Could not change nin of $vertex by $(join(Δs, ", "))!!"))
ΔNinRelaxed(v::AbstractVertex, Δ::Integer) = ΔNin{Relaxed}(v, [Δ])
fallback(s::ΔNin) = s.fallback

"""
    AlignNinToNout <: AbstractJuMPΔSizeStrategy
    AlignNinToNout(vstrat=DefaultJuMPΔSizeStrategy())

Adds variables and constraints for `nin(vi) == nout.(inputs(vi))`.

If it fails, the operation will be retried with the `fallback` strategy (default `ΔSizeFailError`).
"""
struct AlignNinToNout <: AbstractJuMPΔSizeStrategy
    nindict::Dict{AbstractVertex, Vector{JuMP.VariableRef}}
    vstrat::AbstractJuMPΔSizeStrategy
    fallback::AbstractJuMPΔSizeStrategy
end
AlignNinToNout(;vstrat=DefaultJuMPΔSizeStrategy(), fallback=ΔSizeFailError("Failed to align Nin to Nout!!")) = AlignNinToNout(vstrat, fallback)
AlignNinToNout(vstrat, fallback) = AlignNinToNout(Dict{AbstractVertex, JuMP.VariableRef}(), vstrat, fallback)
fallback(s::AlignNinToNout) = s.fallback

"""
    AlignNinToNoutVertices <: AbstractJuMPΔSizeStrategy
    AlignNinToNoutVertices(vin, vout, inds::Integer...;vstrat=AlignNinToNout(), fallback=ΔSizeFailError())
    AlignNinToNoutVertices(vin, vout, inds::AbstractArray{<:Integer},vstrat=AlignNinToNout(), fallback=ΔSizeFailError())

Same as [`AlignNinToNout`](@ref) with an additional constraint that `nin(s.vin)[s.ininds] == nout(s.vout)` where `s` is a `AlignNinToNoutVertices`.

Useful in the context of removing vertices and/or edges.

If it fails, the operation will be retried with the `fallback` strategy (default `ΔSizeFailError`).
"""
struct AlignNinToNoutVertices <: AbstractJuMPΔSizeStrategy
    vin::AbstractVertex
    vout::AbstractVertex
    ininds::AbstractArray{<:Integer}
    vstrat::AlignNinToNout
    fallback::AbstractJuMPΔSizeStrategy
end
AlignNinToNoutVertices(vin, vout, inds::Integer...;vstrat=AlignNinToNout(), fallback=failToAlign(vin, vout)) = AlignNinToNoutVertices(vin, vout, collect(inds), vstrat, fallback)
AlignNinToNoutVertices(vin, vout, inds::AbstractArray{<:Integer}, vstrat::AlignNinToNout, fallback::AbstractJuMPΔSizeStrategy=failToAlign(vin,vout)) = AlignNinToNoutVertices(vin, vout, inds, vstrat, fallback)
AlignNinToNoutVertices(vin, vout, inds::AbstractArray{<:Integer}, innerstrat::AbstractJuMPΔSizeStrategy, fallback=failToAlign(vin, vout)) = AlignNinToNoutVertices(vin, vout, inds, AlignNinToNout(vstrat=innerstrat), fallback)
fallback(s::AlignNinToNoutVertices) = s.fallback

failToAlign(vin, vout) = ΔSizeFailError("Could not align nout of $vin to nin of $(vout)!!")

#TODO: Ugh, this is too many abstraction layers for too little benefit. Refactor so
# all MutationVertex has state?
nin(v::InputSizeVertex) = []
nin(v::MutationVertex) = nin(v, op(v))
nin(v::AbstractVertex, op::MutationState) = nin(op)
nin(v::AbstractVertex, op::MutationOp) = nin(trait(v), v)
nin(t::DecoratingTrait, v::AbstractVertex) = nin(base(t), v)

# SizeTransparent does not need mutation state to keep track of sizes
nin(::SizeTransparent, v::AbstractVertex) = nout.(inputs(v))

nout(v::InputSizeVertex) = v.size
nout(v::MutationVertex) = nout(v, op(v))
nout(v::AbstractVertex, op::MutationState) = nout(op)
nout(v::AbstractVertex, op::MutationOp) = nout(trait(v), v)
nout(t::DecoratingTrait, v::AbstractVertex) = nout(base(t), v)

# SizeTransparent does not need mutation state to keep track of sizes
nout(t::SizeInvariant, v::AbstractVertex) = nin(v)[1]
nout(t::SizeStack, v::AbstractVertex) = sum(nin(v))

nout_org(v::AbstractVertex) = nout_org(trait(v), v)
nout_org(t::DecoratingTrait, v) = nout_org(base(t), v)
nout_org(::MutationSizeTrait, v::MutationVertex) = nout_org(op(v))
nout_org(::Immutable, v) = nout(v)

nin_org(v::InputSizeVertex) = []
nin_org(v::AbstractVertex) = nin_org(trait(v), v)
nin_org(t::DecoratingTrait, v) = nin_org(base(t), v)
nin_org(::MutationSizeTrait, v::MutationVertex) = nin_org(op(v))
nin_org(::Immutable, v) = nin(v)

"""
    minΔnoutfactor(v::AbstractVertex)

Returns the smallest `k` so that allowed changes to `nout` of `v` as well as `nin` of its outputs are `k * n` where `n` is an integer.
Returns `missing` if it is not possible to change `nout`.
"""
minΔnoutfactor(v::AbstractVertex) = missing
minΔnoutfactor(v::MutationVertex) = minΔnoutfactor(trait(v), v)
minΔnoutfactor(t::DecoratingTrait, v::AbstractVertex) = minΔnoutfactor(base(t), v)
minΔnoutfactor(::MutationTrait, v::AbstractVertex) = lcmsafe(vcat(minΔninfactor_only_for.(outputs(v))..., minΔnoutfactor_only_for(v)))
minΔnoutfactor(t::SizeTransparent, v::AbstractVertex) = minΔninfactor(t, v)

"""
    minΔnoutfactor(v::AbstractVertex, [s=VisitState{Int}()])

Returns the smallest `k` so that allowed changes to `nin` of `v` as well as `nout` of its inputs are `k * n` where `n` is an integer.
Returns `missing` if it is not possible to change `nin`.
"""
minΔninfactor(v::AbstractVertex) = lcmsafe(vcat(minΔnoutfactor_only_for.(inputs(v))..., minΔninfactor_only_for(v)))
minΔninfactor(v::MutationVertex) = minΔninfactor(trait(v), v)
minΔninfactor(t::DecoratingTrait, v::AbstractVertex) = minΔninfactor(base(t), v)
minΔninfactor(::MutationTrait, v::AbstractVertex) = lcmsafe(vcat(minΔnoutfactor_only_for.(inputs(v))..., minΔninfactor_only_for(v)))
minΔninfactor(::SizeTransparent, v::AbstractVertex) = lcmsafe([minΔnoutfactor_only_for(v), minΔninfactor_only_for(v)])

"""
    minΔnoutfactor_only_for(v::AbstractVertex, [s=VisitState{Int}()])

Returns the smallest `k` so that allowed changes to `nout` of `v` are `k * n` where `n` is an integer.
Returns `missing` if it is not possible to change `nout`.
"""
minΔnoutfactor_only_for(v::AbstractVertex, s=[]) = outvisit(v, s) ? 1 : minΔnoutfactor_only_for(base(v),s)
minΔnoutfactor_only_for(v::MutationVertex, s=[]) = outvisit(v, s) ? 1 : minΔnoutfactor_only_for(trait(v), v, s)
minΔnoutfactor_only_for(v::InputVertex, s=[]) = missing
minΔnoutfactor_only_for(v::CompVertex, s=[]) = minΔnoutfactor(v.computation)
minΔnoutfactor(f::Function) = 1 # TODO: Move to test as this does not make alot of sense
minΔnoutfactor_only_for(t::DecoratingTrait, v::AbstractVertex, s) = minΔnoutfactor_only_for(base(t), v, s)
minΔnoutfactor_only_for(::Immutable, v::AbstractVertex, s) = missing
minΔnoutfactor_only_for(::SizeAbsorb, v::AbstractVertex, s) = minΔnoutfactor_only_for(base(v),s)

minΔnoutfactor_only_for(::SizeInvariant, v::AbstractVertex, s) = lcmsafe(vcat(map(vi -> minΔnoutfactor_only_for(vi,s), inputs(v)), map(vo -> minΔninfactor_only_for(vo,s), outputs(v))))
function minΔnoutfactor_only_for(::SizeStack, v::AbstractVertex, s)
    absorbing = findterminating(v, inputs, outputs)

    # This is not strictly the minimum as using only one of the factors would work as well
    # However, this would create a bias as the same factor would be used all the time
    # Life is really hard sometimes :(

    # Count thingy is for duplicate outputs. Must be counted twice as it is impossible
    # to only change one of them, right?
    factors = [count(x->x==va,absorbing) * minΔnoutfactor_only_for(va,s) for va in unique(absorbing)]
    return lcmsafe(factors)
end

"""
    minΔninfactor_only_for(v::AbstractVertex)

Returns the smallest `k` so that allowed changes to `nin` of `v` are `k * n` where `n` is an integer.
Returns `missing` if it is not possible to change `nin`.
"""
minΔninfactor_only_for(v::AbstractVertex, s=[]) = invisit(v, s) ? 1 : minΔninfactor_only_for(base(v),s)
minΔninfactor_only_for(v::MutationVertex, s=[]) = invisit(v, s) ? 1 : minΔninfactor_only_for(trait(v), v, s)
minΔninfactor_only_for(v::InputVertex,s=[]) = missing
minΔninfactor_only_for(v::CompVertex,s=[]) = minΔninfactor(v.computation)
minΔninfactor(f::Function) = 1 # TODO: Move to test as this does not make alot of sense
minΔninfactor_only_for(t::DecoratingTrait, v::AbstractVertex, s) = minΔninfactor_only_for(base(t), v, s)
minΔninfactor_only_for(::Immutable, v::AbstractVertex, s) = missing
minΔninfactor_only_for(::SizeAbsorb, v::AbstractVertex, s) = minΔninfactor_only_for(base(v), s)
minΔninfactor_only_for(t::SizeInvariant, v::AbstractVertex, s) = minΔnoutfactor_only_for(t, v, s)
function minΔninfactor_only_for(::SizeStack, v::AbstractVertex, s)
    absorbing = findterminating(v, outputs, inputs)
    # This is not strictly the minimum as using only one of the factors would work as well
    # However, this would create a bias as the same factor would be used all the time
    # Life is really hard sometimes :(

    # Count thingy is for duplicate inputs. Must be counted twice as it is impossible
    # to only change one of them, right?
    factors = [count(x->x==va,absorbing) * minΔnoutfactor_only_for(va, s) for va in unique(absorbing)]
    return lcmsafe(factors)
end

#lcm which also checks for missing and arrays of undefined type
function lcmsafe(x)
    isempty(x) && return 1
    return any(ismissing.(x)) ? missing : lcm(Integer.(x))
end

invisit(v, visited) = hasvisited(v, Input(), visited)
outvisit(v, visited) = hasvisited(v, Output(), visited)

function hasvisited(v, d::Direction, visited)
    (v, d) in visited && return true
    push!(visited, (v, d))
    return false
end

"""
    findterminating(v::AbstractVertex, direction::Function, other::Function= v -> [], visited = [])

Return an array of all vertices which terminate size changes (i.e does not propagate them) seen through the given direction (typically inputs or outputs). A vertex will be present once for each unique path through which its seen.

The `other` direction may be specified and will be traversed if a SizeInvariant vertex is encountered.

Will return the given vertex if it is terminating.

# Examples
```julia-repl

julia> v1 = inputvertex("v1", 3);

julia> v2 = inputvertex("v2", 3);

julia> v3 = conc(v1,v2,v1,dims=1);

julia> name.(findterminating(v1, outputs, inputs))
1-element Array{String,1}:
 "v1"

julia> name.(findterminating(v3, outputs, inputs))
0-element Array{Any,1}

julia> name.(findterminating(v3, inputs, outputs))
3-element Array{String,1}:
 "v1"
 "v2"
 "v1"

 julia> v5 = v3 + inputvertex("v4", 9);

 julia>  # Note, + creates a SizeInvariant vertex and this causes its inputs to be seen through the output direction

 julia> name.(findterminating(v3, outputs, inputs))
 1-element Array{String,1}:
  "v4"
```
"""
function findterminating(v::AbstractVertex, direction::Function, other::Function=v->[], visited = Set{AbstractVertex}())
    v in visited && return []
    push!(visited, v)
    res = findterminating(trait(v), v, direction, other, visited)
    delete!(visited, v)
    return res
 end
findterminating(t::DecoratingTrait, v, d::Function, o::Function, visited) = findterminating(base(t), v, d, o, visited)
findterminating(::SizeAbsorb, v, d::Function, o::Function, visited) = [v]
findterminating(::Immutable, v, d::Function, o::Function, visited) = [v]

findterminating(::SizeStack, v, d::Function, o::Function, visited) = collectterminating(v, d, o, visited)
findterminating(::SizeInvariant, v, d::Function, o::Function, visited) = vcat(collectterminating(v, d, o, visited), collectterminating(v, o, d, visited))
collectterminating(v, d::Function, o::Function, visited) = mapfoldl(vf -> findterminating(vf, d, o, visited), vcat, d(v), init=[])

# Just another name for Δsize with the corresponding direction
Δnout(v::AbstractVertex, Δ::Integer) = Δsize(Output(), v, Δ)
Δnin(v::AbstractVertex, Δs::Maybe{<:Integer}...) = Δsize(Input(), v, Δs...)

"""
    Δsize(d::Direction, v::AbstractVertex, Δ...)

Change size of `v` by `Δ` in direction `d`.
"""
Δsize(d::Input, v, Δs::Maybe{T}...) where T <:Integer = Δsize(ΔNinExact(v, collect(Maybe{T}, Δs)), all_in_Δsize_graph(v, d))
Δsize(d::Output, v, Δ::T) where T <: Integer = Δsize(ΔNoutExact(v, Δ), all_in_Δsize_graph(v, d))

"""
    Δsize(s::AbstractΔSizeStrategy, vertices::AbstractArray{<:AbstractVertex})

Calculate new sizes for (potentially) all provided `vertices` using the strategy `s` and apply all changes.
"""
function Δsize(s::AbstractΔSizeStrategy, vertices::AbstractVector{<:AbstractVertex})
    execute, nins, nouts = newsizes(s, vertices)
    if execute
        Δsize(nins, nouts, vertices)
    end
end

"""
    Δsize(nouts::AbstractVector{<:Integer}, vertices::AbstractVector{<:AbstractVertex})

Set output size of `vertices[i]` to `nouts[i]` for all `i` in `1:length(vertices)`.
Set input size of all keys `vi` in `nins` to `nins[vi]`.
"""
function Δsize(nins::Dict, nouts::AbstractVector{<:Integer}, vertices::AbstractVector{<:AbstractVertex})
    Δnouts = nouts .- nout.(vertices)

    for (i, vi) in enumerate(vertices)
        ninΔs = get(() -> nin(vi), nins, vi) .- nin(vi)
        Δnin(OnlyFor(), vi, ninΔs...)
        Δnout(OnlyFor(), vi, Δnouts[i])
    end

    for (i, vi) in enumerate(vertices)
        ninΔs = get(() -> nin(vi), nins, vi) .- nin(vi)
        after_Δnin(vi, ninΔs...)
        after_Δnout(vi, Δnouts[i])
    end
end

function Δnin(s::OnlyFor, v, Δs::Maybe{<:Integer}...)
    any(skipmissing(Δs) .!= 0) || return
    Δnin(s, trait(v), v, Δs)
end
Δnin(s::AbstractΔSizeStrategy, t::DecoratingTrait, v, Δs) = Δnin(s, base(t), v, Δs)
function Δnin(s::AbstractΔSizeStrategy, t::SizeChangeLogger, v, Δs)

    @logmsg t.level "Change nin of $(infostr(t, v)) by $(join(compressed_string.(Δs), ", "))"
    Δnin(s, base(t), v, Δs)
end
Δnin(::OnlyFor, ::MutationSizeTrait, v, Δs) = Δnin(op(v), Δs...)

function Δnout(s::OnlyFor, v, Δ::Integer)
    Δ == 0 && return
    Δnout(s, trait(v), v, Δ)
end
Δnout(s::AbstractΔSizeStrategy, t::DecoratingTrait, v, Δ) = Δnout(s, base(t), v, Δ)
function Δnout(s::AbstractΔSizeStrategy, t::SizeChangeLogger, v, Δ)
    @logmsg t.level "Change nout of $(infostr(t, v)) by $(compressed_string(Δ))"
    Δnout(s, base(t), v, Δ)
end
Δnout(::OnlyFor, ::MutationSizeTrait, v, Δ) = Δnout(op(v), Δ)

after_Δnin(v, Δs...) = after_Δnin(trait(v), v, Δs)
after_Δnin(t::DecoratingTrait, v, Δs) = after_Δnin(base(t), v, Δs)
after_Δnin(t::SizeChangeValidation, v, Δs) = validate_Δnin(v, Δs, () -> after_Δnin(base(t), v, Δs))
function after_Δnin(t, v, Δs) end

after_Δnout(v, Δ) = after_Δnout(trait(v), v, Δ)
after_Δnout(t::DecoratingTrait, v, Δ) = after_Δnout(base(t), v, Δ)
after_Δnout(t::SizeChangeValidation, v, Δ) = validate_Δnout(v, Δ, () -> after_Δnout(base(t), v, Δ))
function after_Δnout(t, v, Δ) end

sizeΔ(Δ::Integer) = Δ
sizeΔ(Δ::AbstractArray) = length(Δ)
function validate_Δnin(v::AbstractVertex, Δ, Δfun, validvisit = true)

    # Yeah, this is checking more than one thing. Cba to have three different structs and methods for validation
    length(Δ) == length(inputs(v)) || throw(ArgumentError("Length of Δ must be equal to number of inputs for $(v)! length(Δ) = $(length(Δ)), length(inputs(v)) = $(length(inputs(v)))"))

    if validvisit
        # TODO base(v) makes this a bit weaker than I would have wanted. Right now it is only because testcases use smaller factors to trigger SizeStack to do unusual stuff
        Δninfactor = minΔninfactor_only_for(base(v))
        any(Δi -> sizeΔ(Δi) % Δninfactor != 0, skipmissing(Δ)) && throw(ArgumentError("Nin change of $Δ to $v is not an integer multiple of $(Δninfactor)!"))
    end

    Δfun()

    if validvisit
        nout.(inputs(v)) == nin(v) || throw(ArgumentError("Nin change of $Δ to $v did not result in expected size! Expected: $(nout.(inputs(v))), actual: $(nin(v))"))
    end
end

function validate_Δnout(v::AbstractVertex, Δ, Δfun, validvisit=true)

    if validvisit
        # TODO base(v) makes this a bit weaker than I would have wanted. Right now it is only because testcases use smaller factors to trigger SizeStack to do unusual stuff
        Δnoutfactor = minΔnoutfactor_only_for(base(v))
        sizeΔ(Δ) % Δnoutfactor != 0 && throw(ArgumentError("Nout change of $Δ to $v is not an integer multiple of $(Δnoutfactor)!"))
    end

    Δfun()

    if validvisit
        nin_of_outputs = unique(mapreduce(vi -> nin(vi)[inputs(vi) .== v], vcat, outputs(v), init=nout(v)))

        nin_of_outputs == [nout(v)] || throw(ArgumentError("Nout change of $Δ to $v resulted in size mismatch! Nin of outputs: $nin_of_outputs, nout of this: $([nout(v)])"))
    end
end

newsizes(s::ΔSizeFailError, vertices::AbstractVector{<:AbstractVertex}) = error(s.msg)
newsizes(s::ΔSizeFailNoOp, vertices::AbstractVector{<:AbstractVertex}) = false, Dict(vertices .=> nin.(vertices)), nout.(vertices)
function newsizes(s::LogΔSizeExec, vertices::AbstractVector{<:AbstractVertex})
    @logmsg s.level s.msg
    return newsizes(s.andthen, vertices)
end

"""
    newsizes(s::AbstractΔSizeStrategy, vertices::AbstractArray{<:AbstractVertex})

Return a vector of new outputs sizes for and a `Dict` of new input sizes for all provided `vertices` using the strategy `s`.

Result vector is index aligned with `vertices`.
Result `Dict` has a vector of input sizes for each element of `vertices` which has an input (i.e everything except input vertices).
"""
function newsizes(s::AbstractJuMPΔSizeStrategy, vertices::AbstractVector{<:AbstractVertex})

    model = sizemodel(s, vertices)

    noutvars = @variable(model, noutvars[i=1:length(vertices)], Int)

    noutdict = Dict(zip(vertices, noutvars))
    for v in vertices
        vertexconstraints!(v, s, (model=model, noutdict=noutdict))
    end

    sizeobjective!(s, model, noutvars, vertices)

    JuMP.optimize!(model)

    if accept(s, model)
        return true, ninsAndNouts(s, vertices, noutvars)...
    end
    return newsizes(fallback(s), vertices)
end

"""
    sizemodel(s::AbstractJuMPΔSizeStrategy, vertices)

Return a `JuMP.Model` for executing strategy `s` on `vertices`.
"""
sizemodel(s::AbstractJuMPΔSizeStrategy, vertices) = JuMP.Model(JuMP.optimizer_with_attributes(Cbc.Optimizer, "loglevel"=>0))

# Just a shortcut for broadcasting on dicts
getall(d::Dict, ks, deffun=() -> missing) = get.(deffun, [d], ks)

# First we dispatch on trait in order to filter out immutable vertices
vertexconstraints!(v::AbstractVertex, s, data) = vertexconstraints!(trait(v), v, s, data)
vertexconstraints!(t::DecoratingTrait, v, s, data) = vertexconstraints!(base(t), v, s,data)
function vertexconstraints!(::Immutable, v, s, data)
    @constraint(data.model, data.noutdict[v] == nout(v))
    @constraint(data.model, getall(data.noutdict, inputs(v)) .== nin(v))
    vertexconstraints!(s, v, data)
end
vertexconstraints!(::MutationTrait, v, s, data) = vertexconstraints!(s, v, data)

# This must be applied to immutable vertices as well
function vertexconstraints!(v::AbstractVertex, s::AlignNinToNout, data)
    vertexconstraints!(v, s.vstrat, data)
    # Code below secretly assumes vo is in data.noutdict (ninarr will be left with undef entries otherwise).
    for vo in filter(vo -> vo in keys(data.noutdict), outputs(v))
        ninvar = @variable(data.model, integer=true)
        @constraint(data.model, data.noutdict[v] == ninvar)

        ninarr = get!(() -> Vector{JuMP.VariableRef}(undef, length(inputs(vo))), s.nindict, vo)
        ninarr[inputs(vo) .== v] .= ninvar
    end
end

function vertexconstraints!(v::AbstractVertex, s::AlignNinToNoutVertices, data)
    # Any s.ininds which are mapped to vertices which are not yet added in s.vstrat.nindict[s.vout] will cause undefined reference below
    # There is a check to wait for them to be added, but if they are not in data.noutdict they will never be added, so we need to check for that too. Not 100% this can even happen, so I'm just waiting for this in itself to trigger a bug. Sigh...
    neededinds = filter(i -> i != nothing, indexin(keys(data.noutdict), inputs(s.vout)))

    condition() = s.vout in keys(s.vstrat.nindict) && all(i -> isassigned(s.vstrat.nindict[s.vout], i), neededinds)

    # Just to make sure we only do this once without having to store hasadded as a field in AlignNinToNoutVertices:
    #  - Check the condition before vertexconstraints! -> Only when condition changes from false to true do we add as we assume this is the first time the condition becomes fulfilled
    hasadded = condition()
    vertexconstraints!(v, s.vstrat, data)

    if !hasadded && condition()
        @constraint(data.model, data.noutdict[s.vin] .== s.vstrat.nindict[s.vout][s.ininds])
    end
end

"""
    vertexconstraints!(s::AbstractJuMPΔSizeStrategy, v, data)

Add constraints for `AbstractVertex v` using strategy `s`.

Extra info like the model and variables is provided in `data`.
"""
function vertexconstraints!(s::AbstractJuMPΔSizeStrategy, v, data)
    ninconstraint!(s, v, data)
    compconstraint!(s, v, (data..., vertex=v))
    sizeconstraint!(s, v, data)
end

"""
    sizeconstraint!(s::AbstractJuMPΔSizeStrategy, v, data)

Add size constraints for `AbstractVertex v` using strategy `s`.

Extra info like the model and variables is provided in `data`.
"""
sizeconstraint!(s::AbstractJuMPΔSizeStrategy, v, data) = @constraint(data.model, data.noutdict[v] >= 1)

function sizeconstraint!(s::ΔNout{Exact}, v, data)
    if v == s.vertex
        @constraint(data.model, data.noutdict[v] == nout(v) + s.Δ)
    else
        sizeconstraint!(DefaultJuMPΔSizeStrategy(), v, data)
    end
end

function sizeconstraint!(s::ΔNin{Exact}, v, data)
    if v in s.vertices
        Δ = s.Δs[s.vertices .== v][1]
        @constraint(data.model, data.noutdict[v] == nout(v) + Δ)
    else
        sizeconstraint!(DefaultJuMPΔSizeStrategy(), v, data)
    end
end

"""
    ninconstraint!(s, v, data)

Add input size constraints for `AbstractVertex v` using strategy `s`.

Extra info like the model and variables is provided in `data`.
"""
ninconstraint!(s, v, data) = ninconstraint!(s, trait(v), v, data)
ninconstraint!(s, t::DecoratingTrait, v, data) = ninconstraint!(s, base(t), v, data)
function ninconstraint!(s, ::MutationTrait, v, data) end
ninconstraint!(s, ::SizeStack, v, data) = @constraint(data.model, sum(getall(data.noutdict, inputs(v))) == data.noutdict[v])
ninconstraint!(s, ::SizeInvariant, v, data) = @constraint(data.model, getall(data.noutdict, unique(inputs(v))) .== data.noutdict[v])

"""
    ninconstraint!(s, v, data)

Add constraints on the computation (e.g. neural network layer) for `AbstractVertex v` using strategy `s`.

Extra info like the model and variables is provided in `data`.
"""
compconstraint!(s, v::AbstractVertex, data) = compconstraint!(s, base(v), data)
compconstraint!(s, v::CompVertex, data) = compconstraint!(s, v.computation, data)
function compconstraint!(s, v::InputVertex, data) end
function compconstraint!(s, f, data) end

"""
    sizeobjective!(s::AbstractJuMPΔSizeStrategy, model, noutvars, sizetargets)

Add the objective for `noutvars` using strategy `s`.
"""
sizeobjective!(s::AbstractJuMPΔSizeStrategy, model, noutvars, vertices) = @objective(model, Min, objective!(s, model, noutvars, vertices))

function objective!(s, model, noutvars, vertices)
    sizetargets = nout.(vertices)
    # L1 norm prevents change in vertices which does not need to change.
    # Max norm tries to spread out the change so no single vertex takes most of the change.
    return norm!(SumNorm(0.1 => L1NormLinear(), 0.8 => MaxNormLinear()), model, @expression(model, objective[i=1:length(noutvars)], noutvars[i] - sizetargets[i]), sizetargets)
end

objective!(s::ΔNout{Relaxed}, model, noutvars, vertices) = noutrelax!(model, [s.vertex], [s.Δ], noutvars, vertices)
objective!(s::ΔNin{Relaxed}, model, noutvars, vertices) = noutrelax!(model, s.vertices, s.Δs, noutvars, vertices)

function noutrelax!(model, vs, Δs, noutvars, vertices)
    inds = mapreduce(v -> vertices .== v, (i1,i2) -> i1 .| i2, vs)
    def_obj = objective!(DefaultJuMPΔSizeStrategy(), model, noutvars[.!inds], vertices[.!inds])
    sizetarget = nout.(vs) + Δs
    Δnout_obj = norm!(L1NormLinear(), model, @expression(model, noutvars[inds] .- sizetarget))
    # Force it to change as s.Δ might be too small
    # Trick from http://lpsolve.sourceforge.net/5.1/absolute.htm
    Δnout_const = @expression(model, noutvars[inds] - nout.(vs))
    B = @variable(model, [1:length(vs)], binary=true)
    M = 1e5
    ϵ = 1e-2 # abs(Δnout_const) must be larger than this
    @constraint(model, Δnout_const .+ M .* B .>= ϵ)
    @constraint(model, Δnout_const .+ M .* B .<= M .- ϵ)

    return @expression(model, def_obj + 1e6*sum(Δnout_obj))
end

"""
    accept(::AbstractJuMPΔSizeStrategy, model::JuMP.Model)

Return true of the solution for `model` is accepted using strategy `s`.
"""
accept(::AbstractJuMPΔSizeStrategy, model::JuMP.Model) = JuMP.termination_status(model) != MOI.INFEASIBLE && JuMP.primal_status(model) == MOI.FEASIBLE_POINT

function ninsAndNouts(::AbstractJuMPΔSizeStrategy, vs, noutvars)
    nouts = round.(Int, JuMP.value.(noutvars))
    mapnout(i::Integer) = nouts[i]
    mapnout(i::Nothing) = missing

    nins = Dict(vs .=> map(vi -> mapnout.(indexin(inputs(vi), vs)), vs))
    return nins, nouts
end

function ninsAndNouts(s::AlignNinToNout, vs, noutvars)
    nouts = round.(Int, JuMP.value.(noutvars))
    nins = Dict(key => round.(Int, JuMP.value.(value)) for (key, value) in s.nindict)
    return nins,nouts
end

ninsAndNouts(s::AlignNinToNoutVertices, vs, noutvars) = ninsAndNouts(s.vstrat, vs, noutvars)

# TODO: Remove since only used for debugging. If only it wasn't so bloody cumbersome to just list the constraints in a JuMP model....
nconstraints(model) = mapreduce(tt -> JuMP.num_constraints.(model,tt...), +,  filter(tt -> tt != (JuMP.VariableRef, MOI.Integer), JuMP.list_of_constraint_types(model)), init=0)

# TODO: Remove since only used for debugging. If only it wasn't so bloody cumbersome to just list the constraints in a JuMP model....
function list_constraints(model)
    for tt in filter(tt -> tt != (JuMP.VariableRef, MOI.Integer), JuMP.list_of_constraint_types(model))
        display(JuMP.all_constraints(model, tt...))
    end
end


using Oceananigans.BoundaryConditions: default_auxiliary_bc
using Oceananigans.Fields: FunctionField, data_summary
using Oceananigans.AbstractOperations: AbstractOperation
using Oceananigans.Operators: assumed_field_location
using Oceananigans.OutputWriters: output_indices

import Oceananigans.Fields: set!, compute!, compute_at!, validate_field_data, validate_boundary_conditions
import Oceananigans.Fields: validate_indices, FieldBoundaryBuffers
import Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions
import Base: fill!
import Oceananigans.Simulations: hasnan

# Field and FunctionField (both fields with "grids attached")
const MultiRegionField{LX, LY, LZ, O} = Field{LX, LY, LZ, O, <:MultiRegionGrid} where {LX, LY, LZ, O}
const MultiRegionComputedField{LX, LY, LZ, O} = Field{LX, LY, LZ, <:AbstractOperation, <:MultiRegionGrid} where {LX, LY, LZ}
const MultiRegionFunctionField{LX, LY, LZ, C, P, F} = FunctionField{LX, LY, LZ, C, P, F, <:MultiRegionGrid} where {LX, LY, LZ, C, P, F}

const GriddedMultiRegionField = Union{MultiRegionField, MultiRegionFunctionField}
const GriddedMultiRegionFieldTuple{N, T} = NTuple{N, T} where {N, T<:GriddedMultiRegionField}
const GriddedMultiRegionFieldNamedTuple{S, N} = NamedTuple{S, N} where {S, N<:GriddedMultiRegionFieldTuple}

# Utils
Base.size(f::GriddedMultiRegionField) = size(getregion(f.grid, 1))

@inline isregional(f::GriddedMultiRegionField) = true
@inline devices(f::GriddedMultiRegionField)    = devices(f.grid)
@inline sync_all_devices!(f::GriddedMultiRegionField)  = sync_all_devices!(devices(f.grid))

@inline switch_device!(f::GriddedMultiRegionField, d) = switch_device!(f.grid, d)
@inline getdevice(f::GriddedMultiRegionField, d)      = getdevice(f.grid, d)

@inline getregion(f::MultiRegionFunctionField{LX, LY, LZ}, r) where {LX, LY, LZ} =
    FunctionField{LX, LY, LZ}(_getregion(f.func, r),
                              _getregion(f.grid, r),
                              clock = _getregion(f.clock, r),
                              parameters = _getregion(f.parameters, r))

@inline getregion(f::MultiRegionField{LX, LY, LZ}, r) where {LX, LY, LZ} =
    Field{LX, LY, LZ}(_getregion(f.grid, r),
                      _getregion(f.data, r),
                      _getregion(f.boundary_conditions, r),
                      _getregion(f.indices, r),
                      _getregion(f.operand, r),
                      _getregion(f.status, r),
                      _getregion(f.boundary_buffers, r))

@inline _getregion(f::MultiRegionFunctionField{LX, LY, LZ}, r) where {LX, LY, LZ} =
    FunctionField{LX, LY, LZ}(getregion(f.func, r),
                              getregion(f.grid, r),
                              clock = getregion(f.clock, r),
                              parameters = getregion(f.parameters, r))

@inline _getregion(f::MultiRegionField{LX, LY, LZ}, r) where {LX, LY, LZ} =
    Field{LX, LY, LZ}(getregion(f.grid, r),
                      getregion(f.data, r),
                      getregion(f.boundary_conditions, r),
                      getregion(f.indices, r),
                      getregion(f.operand, r),
                      getregion(f.status, r),
                      getregion(f.boundary_buffers, r))

@inline reconstruct_global_field(f::AbstractField) = f

function reconstruct_global_field(mrf::MultiRegionField)
    global_grid  = on_architecture(CPU(), reconstruct_global_grid(mrf.grid))
    global_field = Field(location(mrf), global_grid)

    data = construct_regionally(interior, mrf)
    data = construct_regionally(Array, data)
    compact_data!(global_field, global_grid, data, mrf.grid.partition)

    fill_halo_regions!(global_field)
    return global_field
end

## Functions applied regionally
set!(mrf::MultiRegionField, v)  = apply_regionally!(set!,  mrf, v)
fill!(mrf::MultiRegionField, v) = apply_regionally!(fill!, mrf, v)

compute_at!(mrf::GriddedMultiRegionField, time)  = apply_regionally!(compute_at!, mrf, time)
compute_at!(mrf::MultiRegionComputedField, time) = apply_regionally!(compute_at!, mrf, time)

@inline hasnan(field::MultiRegionField) = (&)(construct_regionally(hasnan, field).regions...)

validate_indices(indices, loc, mrg::MultiRegionGrid, args...) =
    construct_regionally(validate_indices, indices, loc, mrg.region_grids, args...)

FieldBoundaryBuffers(grid::MultiRegionGrid, args...; kwargs...) =
    construct_regionally(FieldBoundaryBuffers, grid, args...; kwargs...)

FieldBoundaryConditions(mrg::MultiRegionGrid, loc, args...; kwargs...) =
  construct_regionally(inject_regional_bcs, mrg, Iterate(1:length(mrg)), Reference(mrg.partition), Reference(loc), args...; kwargs...)

function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions,
                                              mrg::MultiRegionGrid,
                                              field_name::Symbol,
                                              prognostic_field_name=nothing)

  reg_bcs = regularize_field_boundary_conditions(bcs, mrg.region_grids[1], field_name, prognostic_field_name)
  loc = assumed_field_location(field_name)

  return FieldBoundaryConditions(mrg, loc; west = reg_bcs.west,
                                           east = reg_bcs.east,
                                           south = reg_bcs.south,
                                           north = reg_bcs.north,
                                           bottom = reg_bcs.bottom,
                                           top = reg_bcs.top,
                                           immersed = reg_bcs.immersed)
end

function inject_regional_bcs(grid, region, partition, loc, args...;
                              west = default_auxiliary_bc(topology(grid, 1)(), loc[1]()),
                              east = default_auxiliary_bc(topology(grid, 1)(), loc[1]()),
                             south = default_auxiliary_bc(topology(grid, 2)(), loc[2]()),
                             north = default_auxiliary_bc(topology(grid, 2)(), loc[2]()),
                            bottom = default_auxiliary_bc(topology(grid, 3)(), loc[3]()),
                               top = default_auxiliary_bc(topology(grid, 3)(), loc[3]()),
                          immersed = NoFluxBoundaryCondition())

  west  = inject_west_boundary(region, partition, west)
  east  = inject_east_boundary(region, partition, east)
  south = inject_south_boundary(region, partition, south)
  north = inject_north_boundary(region, partition, north)
  return FieldBoundaryConditions(west, east, south, north, bottom, top, immersed)
end

function Base.show(io::IO, field::MultiRegionField)

  bcs = field.boundary_conditions

  prefix =
      string("$(summary(field))\n",
             "├── grid: ", summary(field.grid), '\n',
             "├── boundary conditions: ", summary(bcs), '\n')
  middle = isnothing(field.operand) ? "" :
      string("├── operand: ", summary(field.operand), '\n',
             "├── status: ", summary(field.status), '\n')

  suffix = string("└── data: ", summary(field.data), '\n',
                  "    └── ", data_summary(field))

  print(io, prefix, middle, suffix)
end


module DoseCalculators

export dose_calculator_app

using Gtk

const _apps = GtkWindow[]
const _apps_should_persist = [true]

"""
    dose_calculator_app(func::Function, title = "Dose Calculator", rlabel = "Results")

Create a `Gtk` widget with entries for `weight`, `height`, `age`, and dose interval.
Arguments:
`func` is a function which takes named arguments age = years, weight =  kg, height = cm and returns total 24-hour dosage in mg.
`title` is the title for the app
`rlabel` is the label to be given the results
"""
function dose_calculator_app(func::Function, title = "Dose Calculator", rlabel = "Results")
    wentry, aentry, hentry, qentry = GtkEntry(), GtkEntry(), GtkEntry(), GtkEntry()
    weightunits = [GtkRadioButton("kg")]
    push!(weightunits, GtkRadioButton(weightunits[1], "lb"))
    ageunits = [GtkRadioButton("years")]
    push!(ageunits, GtkRadioButton(ageunits[1], "months"))
    heightunits = [GtkRadioButton("cm")]
    push!(heightunits, GtkRadioButton(heightunits[1], "in"))
    resultbutton = GtkButton("Calculate")
    resultlabel = GtkLabel("0")
    win = GtkWindow(title, 500, 100) |> (GtkFrame() |> (vbox = GtkBox(:v)))
    push!(_apps, win)
    wbox = GtkButtonBox(:h)
    push!(wbox, GtkLabel("Patient Weight"), wentry, weightunits...)
    abox = GtkButtonBox(:h)
    push!(abox, GtkLabel("Age"), aentry, ageunits...)
    hbox = GtkButtonBox(:h)
    push!(hbox, GtkLabel("Height"), hentry, heightunits...)
    resultbox = GtkButtonBox(:h)
    qbox = GtkButtonBox(:h)
    push!(qbox, GtkLabel("Give medication every: "), qentry, GtkLabel(" hours."))
    set_gtk_property!(qentry, :text, "12")
    push!(resultbox, GtkLabel(rlabel), resultlabel, resultbutton)
    push!(vbox, wbox, abox, hbox, qbox, resultbox)

    function calculate(w)
        wt = something(tryparse(Float64, get_gtk_property(wentry, :text, String)), 0) /
            (get_gtk_property(weightunits[1], :active, Bool) ? 1 : 0.393701)
        ht = something(tryparse(Float64, get_gtk_property(hentry, :text, String)), 0) *
            (get_gtk_property(heightunits[1], :active, Bool) ? 1 : 2.20462)
        ag = something(tryparse(Float64, get_gtk_property(aentry, :text, String)), 0) /
            (get_gtk_property(ageunits[1], :active, Bool) ? 1 : 12)
        result = func(weight = wt, height = ht, age = ag)
        dosagefraction = parse(Float64, get_gtk_property(qentry, :text, String)) / 24
        GAccessor.text(resultlabel, "                      ")
        GAccessor.text(resultlabel, string(result * dosagefraction))
    end

    signal_connect(calculate, resultbutton, :clicked)

    if _apps_should_persist[1]
        cond = Condition()
        endit(w) = notify(cond)
        signal_connect(endit, win, :destroy)
        showall(win)
        wait(cond)
    else
        sleep(5)
    end
end

end # module


using Revise
using VCFTools
using MendelImpute
using GeneticVariation
using Random
using JLSO
using LinearAlgebra

function filter_and_mask(data, maf)
    # filter chromosome data for unique snps
    #println("filtering for unique snps")
    #full_record_index = .!find_duplicate_marker(data)
    #@time VCFTools.filter(data, full_record_index, 1:nsamples(data),
    #    des = "full.uniqueSNPs.vcf.gz")

    # summarize data
    println("summarizing data")
    total_snps, samples, _, _, _, maf_by_record, _ = gtstats("full.uniqueSNPs.vcf.gz")

    # keep snps with at least 5 copies of minor alleles
    snps_tokeep = findall(x -> x ≥ 5 / 2samples, maf_by_record)

    # generate target panel with all snps
    println("generating complete target panel")
    n = 1000
    sample_idx = falses(samples)
    sample_idx[1:n] .= true
    shuffle!(sample_idx)
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", snps_tokeep,
        sample_idx, des = "target.full.vcf.gz", allow_multiallelic=false)

    #also generate reference panel without target samples
    println("generating reference panel without target samples")
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", snps_tokeep,
        .!sample_idx, des = "ref.excludeTarget.vcf.gz", allow_multiallelic=false)

    # generate target file with 1000 samples and typed snps with certain maf
    println("generating target file with typed snps only")
    my_maf = findall(x -> x > maf, maf_by_record)
    p = length(my_maf)
    record_idx = falses(total_snps)
    record_idx[my_maf] .= true
    @time VCFTools.filter("full.uniqueSNPs.vcf.gz", record_idx, sample_idx,
        des = "target.typedOnly.maf$maf.vcf.gz", allow_multiallelic=false)

    # unphase and mask 1% entries in target file
    println("unphasing and masking entries in target file with typed snps only")
    masks = falses(p, n)
    missingprop = 0.005
    for j in 1:n, i in 1:p
        rand() < missingprop && (masks[i, j] = true)
    end
    @time mask_gt("target.typedOnly.maf$maf.vcf.gz", masks,
        des="target.typedOnly.maf$maf.masked.vcf.gz", unphase=true)

    # compress reference file to jlso format
    max_d   = [1000]
    reffile = "ref.excludeTarget.vcf.gz"
    tgtfile = "target.typedOnly.maf$maf.masked.vcf.gz"
    H, H_sampleID, H_chr, H_pos, H_ids, H_ref, H_alt = convert_ht(Bool, reffile, trans=true, save_snp_info=true, msg="importing reference data...")
    X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt = VCFTools.convert_gt(UInt8, tgtfile, trans=true, save_snp_info=true, msg = "Importing genotype file...")
    for d in max_d
        outfile = "ref.excludeTarget.maxd$d.jlso"
        @time compress_haplotypes(H, X, outfile, X_pos, H_sampleID, H_chr, H_pos, H_ids, H_ref, H_alt, d, 0, 0.0)
    end
end
Random.seed!(2020)
data = "/mnt/AppRun/msprime/full3.vcf"
maf = 0.05
@time filter_and_mask(data, maf)


export hmm_viterbi

import ..HMMModel

using mlpack._Internal.io

import mlpack_jll
const hmm_viterbiLibrary = mlpack_jll.libmlpack_julia_hmm_viterbi

# Call the C binding of the mlpack hmm_viterbi binding.
function hmm_viterbi_mlpackMain()
  success = ccall((:hmm_viterbi, hmm_viterbiLibrary), Bool, ())
  if !success
    # Throw an exception---false means there was a C++ exception.
    throw(ErrorException("mlpack binding error; see output"))
  end
end

" Internal module to hold utility functions. "
module hmm_viterbi_internal
  import ..hmm_viterbiLibrary

import ...HMMModel

# Get the value of a model pointer parameter of type HMMModel.
function IOGetParamHMMModel(paramName::String)::HMMModel
  HMMModel(ccall((:IO_GetParamHMMModelPtr, hmm_viterbiLibrary), Ptr{Nothing}, (Cstring,), paramName))
end

# Set the value of a model pointer parameter of type HMMModel.
function IOSetParamHMMModel(paramName::String, model::HMMModel)
  ccall((:IO_SetParamHMMModelPtr, hmm_viterbiLibrary), Nothing, (Cstring, Ptr{Nothing}), paramName, model.ptr)
end

# Serialize a model to the given stream.
function serializeHMMModel(stream::IO, model::HMMModel)
  buf_len = UInt[0]
  buf_ptr = ccall((:SerializeHMMModelPtr, hmm_viterbiLibrary), Ptr{UInt8}, (Ptr{Nothing}, Ptr{UInt}), model.ptr, Base.pointer(buf_len))
  buf = Base.unsafe_wrap(Vector{UInt8}, buf_ptr, buf_len[1]; own=true)
  write(stream, buf)
end
# Deserialize a model from the given stream.
function deserializeHMMModel(stream::IO)::HMMModel
  buffer = read(stream)
  HMMModel(ccall((:DeserializeHMMModelPtr, hmm_viterbiLibrary), Ptr{Nothing}, (Ptr{UInt8}, UInt), Base.pointer(buffer), length(buffer)))
end
end # module

"""
    hmm_viterbi(input, input_model; [verbose])

This utility takes an already-trained HMM, specified as `input_model`, and
evaluates the most probable hidden state sequence of a given sequence of
observations (specified as '`input`, using the Viterbi algorithm.  The computed
state sequence may be saved using the `output` output parameter.

For example, to predict the state sequence of the observations `obs` using the
HMM `hmm`, storing the predicted state sequence to `states`, the following
command could be used:

```julia
julia> using CSV
julia> obs = CSV.read("obs.csv")
julia> states = hmm_viterbi(obs, hmm)
```

# Arguments

 - `input::Array{Float64, 2}`: Matrix containing observations,
 - `input_model::HMMModel`: Trained HMM to use.
 - `verbose::Bool`: Display informational messages and the full list of
      parameters and timers at the end of execution.  Default value `false`.

# Return values

 - `output::Array{Int, 2}`: File to save predicted state sequence to.

"""
function hmm_viterbi(input,
                     input_model::HMMModel;
                     verbose::Union{Bool, Missing} = missing,
                     points_are_rows::Bool = true)
  # Force the symbols to load.
  ccall((:loadSymbols, hmm_viterbiLibrary), Nothing, ());

  IORestoreSettings("Hidden Markov Model (HMM) Viterbi State Prediction")

  # Process each input argument before calling mlpackMain().
  IOSetParamMat("input", input, points_are_rows)
  hmm_viterbi_internal.IOSetParamHMMModel("input_model", convert(HMMModel, input_model))
  if verbose !== nothing && verbose === true
    IOEnableVerbose()
  else
    IODisableVerbose()
  end

  IOSetPassed("output")
  # Call the program.
  hmm_viterbi_mlpackMain()

  return IOGetParamUMat("output", points_are_rows)
end


module LSDYNA

using Printf

export cross_section_extruder, cfile_cross_section_extruder

function cross_section_extruder(xcoords, ycoords, element_info, L, dL, path, scl_filename,lsdyna_filename)

    cd(path)

    io = open(scl_filename, "w")

    println(io,"define:")
    println(io,"void main(void)")
    println(io,"{")

    num_elem = size(element_info)[1]

    for i=1:num_elem
        @printf(io, "ExecuteCommand(\"line param  0 %f %f 0 %f %f\");\n", xcoords[i],ycoords[i],xcoords[i+1],ycoords[i+1])
    end

    println(io, "ExecuteCommand(\"elgenerate shelltype 6\");")
    println(io, "ExecuteCommand(\"genselect target occobject\");")
    println(io, "ExecuteCommand(\"occfilter add Edge\");")
    println(io, "ExecuteCommand(\"genselect whole\");")

    @printf(io, "ExecuteCommand(\"elgenerate shell curvedrag 1 1 %s %s 0 0 0 1 0 0 0",string(L), string(floor(L/dL)))

    for i=1:num_elem
        @printf(io," %s",string(i))
    end

    @printf(io, "\");\n")

    println(io, "ExecuteCommand(\"elgenerate accept\");")

    println(io, "ExecuteCommand(\"genselect whole\");")
    println(io, "ExecuteCommand(\"dupnode merge 0.001000\");")

    @printf(io, "ExecuteCommand(\"geomag del ")

    for i=1:num_elem
        @printf(io," %s",string(i, "e"))
    end

    @printf(io, "\");\n")

    println(io, string("ExecuteCommand(\"save keyword \\","\"",path,"/",lsdyna_filename,"\\","\"\");"))

    println(io,"}")
    println(io,"main();")

    close(io)

end

function cfile_cross_section_extruder(xcoords, ycoords, element_info, L, dL, cfile_filename,lsdyna_filename)

    # cd(path)

    io = open(cfile_filename, "w")

    num_elem = size(element_info)[1]

    for i=1:num_elem
        @printf(io, "line param  0 %f %f 0 %f %f\n", xcoords[i],ycoords[i],xcoords[i+1],ycoords[i+1])
    end

    println(io, "elgenerate shelltype 6")
    println(io, "genselect target occobject")
    println(io, "occfilter add Edge")
    println(io, "genselect whole")

    @printf(io, "elgenerate shell curvedrag 1 1 %s %s 0 0 0 1 0 0 0",string(L), string(floor(L/dL)))

    for i=1:num_elem
        @printf(io," %s",string(i))

        if i==num_elem
            @printf(io," %s\n",string(i))
        end
    end

    println(io, "elgenerate accept")

    # println(io, "genselect whole")
    # println(io, "dupnode merge 0.4000")

    @printf(io, "geomag del ")

    for i=1:num_elem
        @printf(io," %s",string(i, "e"))

        if i==num_elem
            @printf(io," %s\n",string(i, "e"))
        end
    end

    println(io, "dupnode open 1")
    @printf(io, "dupnode showdup %f\n", dL/100)
    println(io, "dupnode select 1")
    println(io, "genselect target node")
    println(io, "genselect node add part 1/0")
    @printf(io, "dupnode merge %f\n", dL/100)  #this dL/100 is a little tricky, might merge nodes not intended to be merged for a very fine mesh
    println(io, "genselect clear")
    println(io, "dupnode accept")

    # kfile_path="/Volumes/GoogleDrive/Shared drives/RunToSolve/Projects/Geometrica/cladding_study_2020/model_generation/all/"

    # lsdyna_filename=kfile_path*"h_purlin_"*string(i)*"_"*string(i)*".k"

    println(io, string("save keyword \"", lsdyna_filename, "\""))

    close(io)

end

end #module


"""
Return snaps stored in `SeisForwExpt` after using `mod!`.

# Arguments

* `iss::Int64=1` : supersource index
"""
function Base.getindex(pa::Param, s::Symbol, iss::Int=1)
	@assert iss≤pa.c.acqgeom[1].nss
	@assert s in [:snaps]
	if(s==:snaps)
		nzd,nxd=length.(pa.c.model.mgrid)
		snaps_all=zeros(nzd,nxd,length(pa.c.itsnaps))
		ip=findall(in.(iss,pa.sschunks))[1]
		issp=findall(x->x==iss,pa.sschunks[ip])[]
		@sync begin
			p=procs(pa.p)[ip]
			@sync remotecall_wait(p) do
				pap=localpart(pa.p)
				for i in eachindex(snaps_all)
					snaps_all[i]=pap.ss[issp].snaps[i]
				end
			end
		end
		return snaps_all
	end
end


function gc_to_comp_sites(sites::SiteParam,groups::GroupParam)
    #given some groups and some sites calculated over those groups
    #calculates "flattened" sites
    comps = groups.components
    sitenames = deepcopy(sites.sites)
    gcnames = groups.flattenedgroups
    for i in 1:length(sitenames)
        gci = gcnames[i]
        sitei = sitenames[i]
        for a in 1:length(sitei)
            sitei[a] = gci * '{' * sitei[a] * '}'
        end
    end

    flattened_comp_sitenames = collect(Iterators.flatten(sitenames))
    flattened_comp_nsites = collect(Iterators.flatten(sites.n_sites))
    siteidx = deepcopy(sites.i_sites)
    for i in 1:length(siteidx)
        sitei = siteidx[i]
        for a in 1:length(sitei)
            sitei[a] = i
        end
    end
    flattened_gcidx = collect(Iterators.flatten(siteidx))
    flattened_comp_isites = collect(Iterators.flatten(sites.i_sites))
    idxdict = Dict((flattened_gcidx[i],flattened_comp_isites[i]) => i for i in 1:length(flattened_gcidx) )

    n_flattenedsites = Vector{Vector{Int}}(undef,length(comps))
    flat_groups = groups.n_flattenedgroups
    for i in 1:length(n_flattenedsites)
        n_flattenedsites[i] = flat_groups[i][flattened_gcidx] .* flattened_comp_nsites
    end
    k = length(flattened_comp_sitenames)
    pairs = [(comps[i],[flattened_comp_sitenames[j] =>n_flattenedsites[i][j] for j in 1:k]) for i in 1:length(comps)]
    n_flattenedsites

    for pair in pairs
        pairname,pairvec = pair
        filter!(x->!iszero(last(x)),pairvec)
    end

    return SiteParam(pairs),idxdict
end

#returns a compressed assoc matrix corresponding to the indices
#of the old gc values, arranged by component instead.
#that is,
function gc_to_comp_assoc_idx(param::AssocParam,sites::SiteParam,idxdict)
    pvals = param.values
    vals,outer,inner = pvals.values, pvals.outer_indices,pvals.inner_indices
    ngc = length(vals)
    nsites = length(sites.flattenedsites)
    comps = sites.components
    ncomps = length(comps)
    site1 = Vector{Int64}(undef,ngc)
    site2 = Vector{Int64}(undef,ngc)
    assoc_idxdict = Dict{Tuple{Int,Int},Int}()

    for ii in 1:ngc
        i,j = outer[ii]
        a,b = inner[ii]
        s1 =  idxdict[(i,a)]
        s2 = idxdict[(j,b)]
        site1[ii] = s1
        site2[ii] = s2
        assoc_idxdict[(s1,s2)] = ii
        assoc_idxdict[(s2,s1)] = ii
    end

    #return site1,site2
    x = Matrix{Matrix{Int}}(undef,ncomps,ncomps)
    for i = 1:ncomps
        xi = zeros(Int,nsites,nsites)
        for (a,b) in Iterators.product(sites.i_sites[i],sites.i_sites[i])
            abval = get(assoc_idxdict,(a,b),0)

            xi[a,b] = abval
            xi[b,a] = abval

        end
        x[i,i] = xi
        for j = 1:i-1
            xi = zeros(Int,nsites,nsites)
            for (a,b) in Iterators.product(sites.i_sites[i],sites.i_sites[j])
                abval = get(assoc_idxdict,(a,b),0)

                xi[a,b] = abval
                xi[b,a] = abval
            end
            x[i,j] = xi
            x[j,i] = xi
        end
    end
    val =  Compressed4DMatrix(x)
    new_inneridx = val.inner_indices
    for (idx,(i,j),(a,b)) ∈ indices(val)
        _a = sites.i_flattenedsites[i][a]
        _b = sites.i_flattenedsites[j][b]
        new_inneridx[idx] = (_a,_b)
    end

    n = findall(!=((0,0)),new_inneridx)
    values = val.values[n]
    outer_indices = val.outer_indices[n]
    inner_indices = new_inneridx[n]
    outer_size = val.outer_size
    inner_size = val.inner_size
    newvals = Compressed4DMatrix(values,outer_indices,inner_indices,outer_size,inner_size)
end


VERSION < v"0.1"

"""
Data filters:    `filter_array`, `filter_array_2`, `filter_array_time`, `getsmoothdata`, `getsmoothdata_nospace`
Data structures:  `ape_budget`, `cat_ape_budget`, `cutborders!`, `surf_quantities`, `cyclone_comp_timemean`, `Composite_Cyclone`, `Composite_Cyclone_v2`, `Composite_Cyclone_v3`
Methods:    `cyclonecompositer`, `shifter`, `smoothfilter`, `cyclonecompositer_v2`, `cyclonecompositer_v3`, `timemean_nofalseframe`, `removefalseframes`, `getapebudget`, `buoyancybudget`

"""
module AvailablePotentialEnergyFramework

using DataStructures: OrderedDict
using ImageFiltering: imfilter, imfilter!, centered, kernelfactors, mapwindow, mapwindow!, Kernel, Inner
using Images: findlocalminima, findlocalmaxima
using ImageSegmentation: SegmentedImage, segment_labels, region_adjacency_graph, seeded_region_growing, labels_map
using Interpolations: LinearInterpolation, interpolate, Gridded, Linear
using NCDatasets: Dataset, variable, defVar
using OffsetArrays: OffsetArray
using SparseArrays: SparseMatrixCSC
using Statistics: mean, median!, mean!
using Unitful: @u_str, unit, ustrip, Quantity
using LoopVectorization: @turbo, @tturbo
using Tullio: @tullio

include("apehelperfunctions.jl")
include("compositehelperfunctions.jl")
include("apebudgets.jl")
include("physicalconstants.jl")
include("physicsfunctions.jl")
include("useful_diagnostics.jl")
#include("ape_computation_from_julia_output.jl") #testing purposes
include("arrayfiltering.jl")
include("datamanagement.jl")
export
# Filters
    filter_array!,
    filter_array_2!,
    filter_array_time,
    filter_array,
    getsmoothdata!,
    getsmoothdata_nospace,
    #Data structures
    cat_ape_budget,
    cutborders!,
    #Methods
#    findlocalmaxima,
    cyclonecompositer,
    shifter,
    shifter!,
    smooth_and_mask,
    smooth_and_mask!,
    getapebudget,
    buoyancybudget,
    add_allcyclones!,
    averageallindistance,
    detect_cyclones,
    detect_cyclones!,
    get_diabatic_as_residual_buoyancy,
    run_distributed_test,
    ##Constants
    R,
    Dryair,
    Watervapor,
    Liquidwater,
    epsilon,
    g,
    #Physics functions
    compute_N2,
    compute_mse,
    get_mse_budget,
    get_vorticity,
    get_okubo_weiss,
    get_divergence,
    get_saturation_vapor_pressure,
    get_partial_vapor_pressure,
    get_mixing_ratio,
    get_virtual_temperature,
    get_lifted_condensation_level,
    get_specific_entropy,
    get_potential_temperature,
    get_virtual_temperature,
    mixing_ratio_to_specific_humidity,
    specific_humidity_to_mixing_ratio,
    get_buoyancy_of_lifted_parcel,
    surface_sensible_heat_flux_to_buoyancy,
    surface_latent_heat_flux_to_buoyancy,
    get_buoyancy,
    radiative_heating_rate_to_buoyancy,
    get_density_ideal_gas,
    #math functions
    integrate_vertically,
    average_precipitation_per_pw_bin,
    average_precipitation_per_pw_bin_dayang,
    #Files and datamanagement
    smooth_vars_and_write_to_netcdf!,
    create_APE_netcdf,
    set_netcdf_var!
end


@testset "runtestExtractInQuoteInside" begin
    println("Testing runtestExtractInQuoteInside")
    include("../src/bfgaExtractInQuoteInside.jl")
    using .bfgaExtractInQuoteInside

    include("../src/Bf.jl")
    using .BfInterpreter

    @time model = test_serial()
    #@time model = test_parallel()
    prems = model.refSet[length(model.refSet)]
    println("Generation : $(model.params.currentGeneration)")
    println("& fitness : $(model.specific_fitness.fitness(prems, model.instructionsSet))")
    #bfcode :: Array{Float64,1}  = prems.dna
    println("code :  $(prems.program)") #$(join(genesToBfInstructions(bfcode), "") ) ")

    #bft = bfType(prems.dna)
    input1 = "One Two \"One Two, repeat\" after me"
    input2 = "The \"RMA had\" been infected"
    n1 , n2 = length(input1), length(input2)

    output1, _ = execute(prems.program, input1, model.instructionsSet)
    output2, _ = execute(prems.program, input2, model.instructionsSet)

    @test length(output1) >= 15
    @test collect("One Two, repeat") == output1[1:15]
    @test length(output2) >= 7
    @test collect("RMA had") == output2[1:7]

    # do the test_parallel after ...
    #test_parallel()
end


module Matter

include("./geometry/vector.jl")
include("./geometry/vertex.jl")
export Vector, add!, sub!, mult!, div!, magnitude, normalize!, dot, det, cross, rotate!
export Vertex, Contact, vertices, center, mean, area, inertia, translate!, rotate!

end # module