import Base: ==
using Serialization
import Serialization: serialize, deserialize

###### Array Domains ######

struct ArrayDomain{N}
    indexes::NTuple{N, Any}
end

include("../lib/domain-blocks.jl")

ArrayDomain(xs...) = ArrayDomain(xs)
ArrayDomain(xs::Array) = ArrayDomain((xs...,))

indexes(a::ArrayDomain) = a.indexes
chunks(a::ArrayDomain{N}) where {N} = DomainBlocks(
    ntuple(i->first(indexes(a)[i]), Val(N)), map(x->[length(x)], indexes(a)))

(==)(a::ArrayDomain, b::ArrayDomain) = indexes(a) == indexes(b)
Base.getindex(arr::AbstractArray, d::ArrayDomain) = arr[indexes(d)...]

function intersect(a::ArrayDomain, b::ArrayDomain)
    if a === b
        return a
    end
    ArrayDomain(map((x, y) -> _intersect(x, y), indexes(a), indexes(b)))
end

function project(a::ArrayDomain, b::ArrayDomain)
    map(indexes(a), indexes(b)) do p, q
        q .- (first(p) - 1)
    end |> ArrayDomain
end

function getindex(a::ArrayDomain, b::ArrayDomain)
    ArrayDomain(map(getindex, indexes(a), indexes(b)))
end

"""
    alignfirst(a)

Make a subdomain a standalone domain. For example,

    alignfirst(ArrayDomain(11:25, 21:100))
    # => ArrayDomain((1:15), (1:80))
"""
alignfirst(a::ArrayDomain) =
    ArrayDomain(map(r->1:length(r), indexes(a)))

function size(a::ArrayDomain, dim)
    idxs = indexes(a)
    length(idxs) < dim ? 1 : length(idxs[dim])
end
size(a::ArrayDomain) = map(length, indexes(a))
length(a::ArrayDomain) = prod(size(a))
ndims(a::ArrayDomain) = length(size(a))
isempty(a::ArrayDomain) = length(a) == 0

"The domain of an array is a ArrayDomain"
domain(x::AbstractArray) = ArrayDomain([1:l for l in size(x)])

abstract type ArrayOp{T, N} <: AbstractArray{T, N} end
Base.IndexStyle(::Type{<:ArrayOp}) = IndexCartesian()

compute(ctx, x::ArrayOp) =
    compute(ctx, cached_stage(ctx, x)::DArray)

collect(ctx::Context, x::ArrayOp) =
    collect(ctx, compute(ctx, x))

collect(x::ArrayOp) = collect(Context(), x)

function Base.show(io::IO, ::MIME"text/plain", x::ArrayOp)
    write(io, string(typeof(x)))
    write(io, string(size(x)))
end

function Base.show(io::IO, x::ArrayOp)
    m = MIME"text/plain"()
    show(io, m, x)
end

"""
`DArray{T,N,F}(domain, subdomains, chunks, concat)`

An N-dimensional distributed array of element type T.

- `domain`: the whole ArrayDomain of the array
- `subdomains`: a `DomainBlocks` of the same dimensions as the array
- `chunks`: an array of chunks of dimension N
- `concat`: a function of type `F`. `concat(x, y; dims=d)` takes two chunks `x` and `y`
            and concatenates them along dimension `d`. `cat` is used by default.
"""
mutable struct DArray{T,N,F} <: ArrayOp{T, N}
    domain::ArrayDomain{N}
    subdomains::AbstractArray{ArrayDomain{N}, N}
    chunks::AbstractArray{Union{Chunk,Thunk}, N}
    concat::F
    freed::Threads.Atomic{UInt8}
    function DArray{T,N,F}(domain, subdomains, chunks, concat) where {T, N,F}
        A = new(domain, subdomains, chunks, concat, Threads.Atomic{UInt8}(0))
        refcount_chunks(A.chunks)
        finalizer(free!, A)
        A
    end
end

function serialize(io::AbstractSerializer, A::DArray)
    @async refcount_chunks(A)
    invoke(serialize, Tuple{AbstractSerializer,Any}, io, A)
end

function deserialize(io::AbstractSerializer, dt::Type{DArray{T,N,F}}) where {T,N,F}
    nf = fieldcount(dt)
    A = ccall(:jl_new_struct_uninit, Any, (Any,), dt)
    Serialization.deserialize_cycle(io, A)
    for i in 1:nf
        tag = Int32(read(io.io, UInt8)::UInt8)
        if tag != Serialization.UNDEFREF_TAG
            ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), A, i-1, Serialization.handle_deserialize(io, tag))
        end
    end
    finalizer(free!, A)
    A
end

refcount_chunks(A::DArray) = refcount_chunks(A.chunks)
function refcount_chunks(chunks)
    for c in chunks
        if c isa Chunk{<:Any, DRef}
            # increment refcount on the master node
            addrefcount(c.handle, 1)
        elseif c isa Thunk
            refcount_chunks(c.inputs)
        end
    end
end

function free_chunks(chunks)
    @sync for c in chunks
        if c isa Chunk{<:Any, DRef}
            # increment refcount on the master node
            cnt = addrefcount(c.handle, -1)
            cnt <= 0 && @async free!(c.handle)
        elseif c isa Thunk
            free_chunks(c.inputs)
        end
    end
end

function free!(x::DArray)
    freed = Bool(Threads.atomic_cas!(x.freed, UInt8(0), UInt8(1)))
    !freed && @async Dagger.free_chunks(x.chunks)
    nothing
end

# mainly for backwards-compatibility
DArray{T, N}(domain, subdomains, chunks) where {T,N} = DArray(T, domain, subdomains, chunks)

"""
`DArray(T, domain, subdomains, chunks, [concat=cat])`

Creates a distributed array of element type T.

- `T`: element type

rest of the arguments are the same as the DArray constructor.
"""
function DArray(T, domain::ArrayDomain{N},
             subdomains::AbstractArray{ArrayDomain{N}, N},
             chunks::AbstractArray{<:Any, N}, concat=cat) where N
    DArray{T, N, typeof(concat)}(domain, subdomains, chunks, concat)
end

domain(d::DArray) = d.domain
chunks(d::DArray) = d.chunks
domainchunks(d::DArray) = d.subdomains
size(x::DArray) = size(domain(x))
stage(ctx, c::DArray) = c

function collect(ctx::Context, d::DArray; tree=false)
    a = compute(ctx, d)

    if isempty(d.chunks)
        return Array{eltype(d)}(undef, size(d)...)
    end

    dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)]
    if tree
        collect(treereduce_nd(delayed.(dimcatfuncs), a.chunks))
    else
        treereduce_nd(dimcatfuncs, asyncmap(collect, a.chunks))
    end
end

function (==)(x::ArrayOp, y::ArrayOp)
    x === y || reduce((a,b)->a&&b, map(==, x, y))
end

function Base.hash(x::ArrayOp, i::UInt)
    7*objectid(x)-2
end

function Base.isequal(x::ArrayOp, y::ArrayOp)
    x === y
end

"""
`view` of a `DArray` chunk returns a `DArray` of thunks
"""
function Base.view(c::DArray, d)
    subchunks, subdomains = lookup_parts(chunks(c), domainchunks(c), d)
    d1 = alignfirst(d)
    DArray(eltype(c), d1, subdomains, subchunks)
end

function group_indices(cumlength, idxs,at=1, acc=Any[])
    at > length(idxs) && return acc
    f = idxs[at]
    fidx = searchsortedfirst(cumlength, f)
    current_block = (get(cumlength, fidx-1,0)+1):cumlength[fidx]
    start_at = at
    end_at = at
    for i=(at+1):length(idxs)
        if idxs[i] in current_block
            end_at += 1
            at += 1
        else
            break
        end
    end
    push!(acc, fidx=>idxs[start_at:end_at])
    group_indices(cumlength, idxs, at+1, acc)
end

function group_indices(cumlength, idx::Int)
    group_indices(cumlength, [idx])
end

function group_indices(cumlength, idxs::AbstractRange)
    f = searchsortedfirst(cumlength, first(idxs))
    l = searchsortedfirst(cumlength, last(idxs))
    out = cumlength[f:l]
    isempty(out) && return []
    out[end] = last(idxs)
    map(=>, f:l, map(UnitRange, vcat(first(idxs), out[1:end-1].+1), out))
end

_cumsum(x::AbstractArray) = length(x) == 0 ? Int[] : cumsum(x)
function lookup_parts(ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}) where N
    groups = map(group_indices, subdmns.cumlength, indexes(d))
    sz = map(length, groups)
    pieces = Array{Union{Chunk,Thunk}}(undef, sz)
    for i = CartesianIndices(sz)
        idx_and_dmn = map(getindex, groups, i.I)
        idx = map(x->x[1], idx_and_dmn)
        dmn = ArrayDomain(map(x->x[2], idx_and_dmn))
        pieces[i] = delayed(getindex)(ps[idx...], project(subdmns[idx...], dmn))
    end
    out_cumlength = map(g->_cumsum(map(x->length(x[2]), g)), groups)
    out_dmn = DomainBlocks(ntuple(x->1,Val(N)), out_cumlength)
    pieces, out_dmn
end

"""
A DArray object may contain a thunk in it, in which case
we first turn it into a Thunk object and then compute it.
"""
function compute(ctx, x::DArray; persist=true)
    thunk = thunkize(ctx, x, persist=persist)
    if isa(thunk, Thunk)
        compute(ctx, thunk)
    else
        x
    end
end

"""
If a DArray tree has a Thunk in it, make the whole thing a big thunk
"""
function thunkize(ctx, c::DArray; persist=true)
    if any(istask, chunks(c))
        thunks = chunks(c)
        sz = size(thunks)
        dmn = domain(c)
        dmnchunks = domainchunks(c)
        if persist
            foreach(persist!, thunks)
        end
        Thunk(thunks...; meta=true) do results...
            t = eltype(results[1])
            DArray(t, dmn, dmnchunks,
                                  reshape(Union{Chunk,Thunk}[results...], sz))
        end
    else
        c
    end
end

global _stage_cache = WeakKeyDict{Context, Dict}()
"""
A memoized version of stage. It is important that the
tasks generated for the same DArray have the same
identity, for example:

    A = rand(Blocks(100,100), Float64, 1000, 1000)
    compute(A+A')

must not result in computation of A twice.
"""
function cached_stage(ctx, x)
    cache = if !haskey(_stage_cache, ctx)
        _stage_cache[ctx] = Dict()
    else
        _stage_cache[ctx]
    end

    if haskey(cache, x)
        cache[x]
    else
        cache[x] = stage(ctx, x)
    end
end

Base.@deprecate_binding Cat DArray
Base.@deprecate_binding ComputedArray DArray

export Distribute, distribute

struct Distribute{T, N} <: ArrayOp{T, N}
    domainchunks
    data::AbstractArray{T,N}
end

size(x::Distribute) = size(domain(x.data))

export BlockPartition, Blocks

struct Blocks{N}
    blocksize::NTuple{N, Int}
end
Blocks(xs::Int...) = Blocks(xs)

Base.@deprecate BlockPartition Blocks

Distribute(p::Blocks, data::AbstractArray) =
    Distribute(partition(p, domain(data)), data)

function stage(ctx, d::Distribute)
    if isa(d.data, ArrayOp)
        # distributing a dsitributed array
        x = cached_stage(ctx, d.data)
        if d.domainchunks == domainchunks(x)
            return x # already properly distributed
        end
        Nd = ndims(x)
        T = eltype(d.data)
        concat = x.concat
        cs = map(d.domainchunks) do idx
            chunks = cached_stage(ctx, x[idx]).chunks
            shape = size(chunks)
            (delayed() do shape, parts...
                if prod(shape) == 0
                    return Array{T}(undef, shape)
                end
                dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:length(shape)]
                ps = reshape(Any[parts...], shape)
                collect(treereduce_nd(dimcatfuncs, ps))
            end)(shape, chunks...)
        end
    else
        cs = map(c -> delayed(identity)(d.data[c]), d.domainchunks)
    end

    DArray(
           eltype(d.data),
           domain(d.data),
           d.domainchunks,
           cs
    )
end

function distribute(x::AbstractArray, dist)
    compute(Distribute(dist, x))
end

function distribute(x::AbstractArray{T,N}, n::NTuple{N}) where {T,N}
    p = map((d, dn)->ceil(Int, d / dn), size(x), n)
    distribute(x, Blocks(p))
end

function distribute(x::AbstractVector, n::Int)
    distribute(x, (n,))
end

function distribute(x::AbstractVector, n::Vector{<:Integer})
    distribute(x, DomainBlocks((1,), (cumsum(n),)))
end

function Base.:(==)(x::ArrayOp{T,N}, y::AbstractArray{S,N}) where {T,S,N}
    collect(x) == y
end

function Base.:(==)(x::AbstractArray{T,N}, y::ArrayOp{S,N}) where {T,S,N}
    return collect(x) == y
end


"""
spin18.jl - sampling robustness
"""

WDIR = joinpath(@__DIR__, "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin18"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const SAMPLE_COUNT = 8
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
const ASTATE_SIZE = ASTATE_SIZE_BASE + SAMPLE_COUNT * HDIM_ISO
const ACONTROL_SIZE = CONTROL_COUNT
const INITIAL_STATE1 = [1., 0, 0, 0]
const INITIAL_STATE2 = [0., 1, 0, 0]
const INITIAL_STATE3 = [1., 0, 0, 1] ./ sqrt(2)
const INITIAL_STATE4 = [1., -1, 0, 0] ./ sqrt(2)
# state indices
const STATE1_IDX = 1:HDIM_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const S1_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO
const S2_IDX = S1_IDX[end] + 1:S1_IDX[end] + HDIM_ISO
const S3_IDX = S2_IDX[end] + 1:S2_IDX[end] + HDIM_ISO
const S4_IDX = S3_IDX[end] + 1:S3_IDX[end] + HDIM_ISO
const S5_IDX = S4_IDX[end] + 1:S4_IDX[end] + HDIM_ISO
const S6_IDX = S5_IDX[end] + 1:S5_IDX[end] + HDIM_ISO
const S7_IDX = S6_IDX[end] + 1:S6_IDX[end] + HDIM_ISO
const S8_IDX = S7_IDX[end] + 1:S7_IDX[end] + HDIM_ISO
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT

# model
struct Model <: AbstractModel
    namp::Float64
end
@inline RD.state_dim(::Model) = ASTATE_SIZE
@inline RD.control_dim(::Model) = ACONTROL_SIZE

# This cost puts a gate error cost on
# the sample states and a LQR cost on the other terms.
# The hessian w.r.t the state and controls is constant.
struct Cost{N,M,T} <: TO.CostFunction
    Q::Diagonal{T, SVector{N,T}}
    R::Diagonal{T, SVector{M,T}}
    q::SVector{N, T}
    c::T
    hess_astate::Symmetric{T, SMatrix{N,N,T}}
    target_states::Array{SVector{HDIM_ISO, T}, 1}
    q_ss1::T
    q_ss2::T
    q_ss3::T
    q_ss4::T
end

function Cost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
              xf::SVector{N,T}, target_states::Array{SVector{HDIM_ISO}, 1},
              q_ss1::T, q_ss2::T, q_ss3::T, q_ss4::T) where {N,M,T}
    q = -Q * xf
    c = 0.5 * xf' * Q * xf
    hess_astate = zeros(N, N)
    # For reasons unknown to the author, throwing a -1 in front
    # of the gate error Hessian makes the cost function work.
    # This is strange, because the gate error Hessian has been
    # checked against autodiff.
    hess_state1 = -1 * q_ss1 * hessian_gate_error_iso2(target_states[1])
    hess_state2 = -1 * q_ss2 * hessian_gate_error_iso2(target_states[2])
    hess_state3 = -1 * q_ss3 * hessian_gate_error_iso2(target_states[3])
    hess_state4 = -1 * q_ss4 * hessian_gate_error_iso2(target_states[4])
    hess_astate[S1_IDX, S1_IDX] = hess_state1
    hess_astate[S2_IDX, S2_IDX] = hess_state2
    hess_astate[S3_IDX, S3_IDX] = hess_state3
    hess_astate[S4_IDX, S4_IDX] = hess_state4
    hess_astate[S5_IDX, S5_IDX] = hess_state1
    hess_astate[S6_IDX, S6_IDX] = hess_state2
    hess_astate[S7_IDX, S7_IDX] = hess_state3
    hess_astate[S8_IDX, S8_IDX] = hess_state4
    hess_astate += Q
    hess_astate = Symmetric(SMatrix{N, N}(hess_astate))
    return Cost{N,M,T}(Q, R, q, c, hess_astate, target_states, q_ss1, q_ss2, q_ss3, q_ss4)
end

@inline TO.state_dim(cost::Cost{N,M,T}) where {N,M,T} = N
@inline TO.control_dim(cost::Cost{N,M,T}) where {N,M,T} = M
@inline Base.copy(cost::Cost{N,M,T}) where {N,M,T} = Cost{N,M,T}(
    cost.Q, cost.R, cost.q, cost.c, cost.hess_astate,
    cost.target_states, cost.q_ss1, cost.q_ss2, cost.q_ss3, cost.q_ss4
)

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}) where {N,M,T} = (
    0.5 * astate' * cost.Q * astate + cost.q'astate + cost.c
    + cost.q_ss1 * gate_error_iso2(astate, cost.target_states[1], S1_IDX[1] - 1)
    + cost.q_ss2 * gate_error_iso2(astate, cost.target_states[2], S2_IDX[1] - 1)
    + cost.q_ss3 * gate_error_iso2(astate, cost.target_states[3], S3_IDX[1] - 1)
    + cost.q_ss4 * gate_error_iso2(astate, cost.target_states[4], S4_IDX[1] - 1)
    + cost.q_ss1 * gate_error_iso2(astate, cost.target_states[1], S5_IDX[1] - 1)
    + cost.q_ss2 * gate_error_iso2(astate, cost.target_states[2], S6_IDX[1] - 1)
    + cost.q_ss3 * gate_error_iso2(astate, cost.target_states[3], S7_IDX[1] - 1)
    + cost.q_ss4 * gate_error_iso2(astate, cost.target_states[4], S8_IDX[1] - 1)
)

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}, acontrol::SVector{M}) where {N,M,T} = (
    TO.stage_cost(cost, astate) + 0.5 * acontrol' * cost.R * acontrol
)

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T}) where {N,M,T}
    E.q = (cost.Q * astate + cost.q + [
        @SVector zeros(ASTATE_SIZE_BASE);
        cost.q_ss1 * jacobian_gate_error_iso2(astate, cost.target_states[1], S1_IDX[1] - 1);
        cost.q_ss2 * jacobian_gate_error_iso2(astate, cost.target_states[2], S2_IDX[1] - 1);
        cost.q_ss3 * jacobian_gate_error_iso2(astate, cost.target_states[3], S3_IDX[1] - 1);
        cost.q_ss4 * jacobian_gate_error_iso2(astate, cost.target_states[4], S4_IDX[1] - 1);
        cost.q_ss1 * jacobian_gate_error_iso2(astate, cost.target_states[1], S5_IDX[1] - 1);
        cost.q_ss2 * jacobian_gate_error_iso2(astate, cost.target_states[2], S6_IDX[1] - 1);
        cost.q_ss3 * jacobian_gate_error_iso2(astate, cost.target_states[3], S7_IDX[1] - 1);
        cost.q_ss4 * jacobian_gate_error_iso2(astate, cost.target_states[4], S8_IDX[1] - 1);
    ])
    return false
end

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                      acontrol::SVector{M,T}) where {N,M,T}
    TO.gradient!(E, cost, astate)
    E.r = cost.R * acontrol
    E.c = 0
    return false
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T}) where {N,M,T}
    E.Q = cost.hess_astate
    return true
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                     acontrol::SVector{M,T}) where {N,M,T}
    TO.hessian!(E, cost, astate)
    E.R = cost.R
    E.H .= 0
    return true
end

# dynamics
function RD.discrete_dynamics(::Type{RD.RK3}, model::Model, astate::StaticVector,
                              acontrols::StaticVector, time::Real, dt::Real) where {SC}
    camp = astate[CONTROLS_IDX[1]]
    negi_hc = camp * NEGI_H1_ISO
    h_prop = exp((FQ_NEGI_H0_ISO + negi_hc) * dt)
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrols[D2CONTROLS_IDX[1]]

    hp_prop = exp((FQ_NEGI_H0_ISO + (camp + model.namp) * NEGI_H1_ISO) * dt)
    hn_prop = exp((FQ_NEGI_H0_ISO + (camp - model.namp) * NEGI_H1_ISO) * dt)
    s1 = hp_prop * astate[S1_IDX]
    s2 = hp_prop * astate[S2_IDX]
    s3 = hp_prop * astate[S3_IDX]
    s4 = hp_prop * astate[S4_IDX]
    s5 = hn_prop * astate[S5_IDX]
    s6 = hn_prop * astate[S6_IDX]
    s7 = hn_prop * astate[S7_IDX]
    s8 = hn_prop * astate[S8_IDX]

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
        s1; s2; s3; s4; s5; s6; s7; s8;
    ]

    return astate_
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=56.8, solver_type=altro,
                  sqrtbp=false, integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e0, 1e0, 1e0, 1e0, 1e-1],
                  dt_inv=Int64(1e1), smoke_test=false, constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true, max_iterations=Int64(2e5),
                  namp=NAMP_PREFACTOR)
    # model configuration
    model = Model(namp)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.
    tf = evolution_time

    # initial state
    x0 = SVector{n}([
        INITIAL_STATE1;
        INITIAL_STATE2;
        zeros(3 * CONTROL_COUNT);
        repeat([INITIAL_STATE1; INITIAL_STATE2; INITIAL_STATE3; INITIAL_STATE4], 2);
    ])
    # target state
    gate_unitary = GT_GATE[gate_type]
    target_states = Array{SVector{HDIM_ISO}, 1}(undef, 4)
    target_states[1] = gate_unitary * INITIAL_STATE1
    target_states[2] = gate_unitary * INITIAL_STATE2
    target_states[3] = gate_unitary * INITIAL_STATE3
    target_states[4] = gate_unitary * INITIAL_STATE4
    xf = SVector{n}([
        target_states[1];
        target_states[2];
        zeros(3 * CONTROL_COUNT);
        repeat([target_states[1]; target_states[2];
                target_states[3]; target_states[4]], 2);
    ])

    # control amplitude constraint
    x_max = SVector{n}([
        fill(Inf, STATE_COUNT * HDIM_ISO);
        fill(Inf, CONTROL_COUNT);
        fill(MAX_CONTROL_NORM_0, 1); # control
        fill(Inf, CONTROL_COUNT);
        fill(Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    x_min = SVector{n}([
        fill(-Inf, STATE_COUNT * HDIM_ISO);
        fill(-Inf, CONTROL_COUNT);
        fill(-MAX_CONTROL_NORM_0, 1); # control
        fill(-Inf, CONTROL_COUNT);
        fill(-Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    # control amplitude constraint at boundary
    x_max_boundary = SVector{n}([
        fill(Inf, STATE_COUNT * HDIM_ISO);
        fill(Inf, CONTROL_COUNT);
        fill(0, 1); # control
        fill(Inf, CONTROL_COUNT);
        fill(Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    x_min_boundary = SVector{n}([
        fill(-Inf, STATE_COUNT * HDIM_ISO);
        fill(-Inf, CONTROL_COUNT);
        fill(0, 1); # control
        fill(-Inf, CONTROL_COUNT);
        fill(-Inf, SAMPLE_COUNT * HDIM_ISO);
    ])

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
    ]) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], 1); # ∫a
        fill(qs[3], 1); # a
        fill(qs[4], 1); # ∂a
        fill(0, SAMPLE_COUNT * HDIM_ISO);
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[9], CONTROL_COUNT);
    ]))
    cost_k = Cost(Q, R, xf, target_states, qs[5], qs[6], qs[7], qs[8])
    cost_f = Cost(Qf, R, xf, target_states, N * qs[5], N * qs[6], N * qs[7], N * qs[8])
    objective = TO.Objective(cost_k, cost_f, N)

    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idxs) for idxs in (
        STATE1_IDX, STATE2_IDX, S1_IDX, S2_IDX, S3_IDX, S4_IDX, S5_IDX, S6_IDX, S7_IDX, S8_IDX,
    )]
    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N);
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints, x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations)
    Altro.solve!(solver)

    # post-process
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = Altro.iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "Q" => Q_arr,
        "Qf" => Qf_arr,
        "R" => R_arr,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "sample_count" => SAMPLE_COUNT,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "integrator_type" => Integer(integrator_type),
        "iterations" => iterations_,
        "max_iterations" => max_iterations,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end


struct GeneralUnstructuredMesh_2D <: UnstructuredMesh_2D
    name::String
    points::Vector{Point_2D}
    edges::Vector{<:SVector{L, UInt32} where {L}}
    materialized_edges::Vector{<:Edge_2D}
    faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}}
    materialized_faces::Vector{<:Face_2D}
    edge_face_connectivity::Vector{SVector{2, UInt32}}
    face_edge_connectivity::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}}
    boundary_edges::Vector{Vector{UInt32}}
    face_sets::Dict{String, Set{UInt32}}
end

function GeneralUnstructuredMesh_2D(;
        name::String = "DefaultMeshName",
        points::Vector{Point_2D} = Point_2D[],
        edges::Vector{<:SVector{L, UInt32} where {L}} = SVector{2, UInt32}[],
        materialized_edges::Vector{<:Edge_2D} = LineSegment_2D[],
        faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}} = SVector{4, UInt32}[],
        materialized_faces::Vector{<:Face_2D} = Triangle_2D[],
        edge_face_connectivity::Vector{SVector{2, UInt32}} = SVector{2, UInt32}[],
        face_edge_connectivity ::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}} = SVector{3, UInt32}[],
        boundary_edges::Vector{Vector{UInt32}} = Vector{UInt32}[],
        face_sets::Dict{String, Set{UInt32}} = Dict{String, Set{UInt32}}()
    ) where {F <: AbstractFloat, U <: Unsigned}
        return GeneralUnstructuredMesh_2D(name,
                                          points,
                                          edges,
                                          materialized_edges,
                                          faces,
                                          materialized_faces,
                                          edge_face_connectivity,
                                          face_edge_connectivity,
                                          boundary_edges,
                                          face_sets,
                                         )
end

Base.broadcastable(mesh::GeneralUnstructuredMesh_2D) = Ref(mesh)

# Return a mesh with boundary edges
function add_boundary_edges(mesh::M; boundary_shape="Unknown"
    ) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edge_face_connectivity)
        mesh = add_connectivity(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = boundary_edges(mesh, boundary_shape),
             face_sets = mesh.face_sets
            )
end

# Return a mesh with face/edge connectivity and edge/face connectivity
function add_connectivity(mesh::UnstructuredMesh_2D)
    return add_edge_face_connectivity(mesh)
end

# Return a mesh with its edges
function add_edges(mesh::M) where {M <: UnstructuredMesh_2D}
    return M(name = mesh.name,
             points = mesh.points,
             edges = edges(mesh),
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with edge/face connectivity
function add_edge_face_connectivity(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.face_edge_connectivity)
        mesh = add_face_edge_connectivity(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = edge_face_connectivity(mesh),
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with every field created
function add_everything(mesh::UnstructuredMesh_2D)
    return add_materialized_faces(
             add_materialized_edges(
               add_boundary_edges(mesh, boundary_shape = "Rectangle")))
end

# Return a mesh with face/edge connectivity
function add_face_edge_connectivity(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edges)
        mesh = add_edges(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = face_edge_connectivity(mesh),
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with materialized edges
function add_materialized_edges(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edges)
        mesh = add_edges(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = materialize_edges(mesh),
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with materialized faces
function add_materialized_faces(mesh::M) where {M <: UnstructuredMesh_2D}
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = materialize_faces(mesh),
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a vector of the faces adjacent to the face of ID face
function adjacent_faces(face::UInt32, mesh::UnstructuredMesh_2D)
    edges = mesh.face_edge_connectivity[face]
    the_adjacent_faces = UInt32[]
    for edge in edges
        faces = mesh.edge_face_connectivity[edge]
        for face_id in faces
            if face_id != face && face_id != 0
                push!(the_adjacent_faces, face_id)
            end
        end
    end
    return the_adjacent_faces
end

# Area of face
function area(face::SVector{N, UInt32}, points::Vector{Point_2D}) where {N}
    return area(materialize_face(face, points))
end

# Return the area of a face set
function area(mesh::UnstructuredMesh_2D, face_set::Set{UInt32})
    if 0 < length(mesh.materialized_faces)
        return mapreduce(x->area(mesh.materialized_faces[x]), +, face_set)
    else
        return mapreduce(x->area(mesh.faces[x], mesh.points), +, face_set)
    end
end

# Return the area of a face set by name
function area(mesh::UnstructuredMesh_2D, set_name::String)
    return area(mesh, mesh.face_sets[set_name])
end

# Bounding box of a vector of points
function boundingbox(points::Vector{Point_2D})
    x = getindex.(points, 1)
    y = getindex.(points, 2)
    return Rectangle_2D(minimum(x), minimum(y), maximum(x), maximum(y))
end

# Bounding box of a vector of points
function boundingbox(points::SVector{L, Point_2D}) where {L}
    x = getindex.(points, 1)
    y = getindex.(points, 2)
    return Rectangle_2D(minimum(x), minimum(y), maximum(x), maximum(y))
end

# Return a vector containing vectors of the edges in each side of the mesh's bounding shape, e.g.
# For a rectangular bounding shape the sides are North, East, South, West. Then the output would
# be [ [e1, e2, e3, ...], [e17, e18, e18, ...], ..., [e100, e101, ...]]
function boundary_edges(mesh::UnstructuredMesh_2D, boundary_shape::String)
    # edges which have face 0 in their edge_face connectivity are boundary edges
    the_boundary_edges = UInt32.(findall(x->x[1] === 0x00000000, mesh.edge_face_connectivity))
    if boundary_shape == "Rectangle"
        # Sort edges into NESW
        bb = boundingbox(mesh.points)
        y_north = bb.ymax
        x_east  = bb.xmax
        y_south = bb.ymin
        x_west  = bb.xmin
        p_NW = Point_2D(x_west, y_north)
        p_NE = bb.tr
        p_SE = Point_2D(x_east, y_south)
        p_SW = bb.bl
        edges_north = UInt32[]
        edges_east = UInt32[]
        edges_south = UInt32[]
        edges_west = UInt32[]
        # Insert edges so that indices move from NW -> NE -> SE -> SW -> NW
        for edge ∈  the_boundary_edges
            epoints = edgepoints(mesh.edges[edge], mesh.points)
            if all(x->abs(x[2] - y_north) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_NW, edges_north, mesh)
            elseif all(x->abs(x[1] - x_east) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_NE, edges_east, mesh)
            elseif all(x->abs(x[2] - y_south) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_SE, edges_south, mesh)
            elseif all(x->abs(x[1] - x_west) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_SW, edges_west, mesh)
            else
                @error "Edge $edge could not be classified as NSEW"
            end
        end
        return [ edges_north, edges_east, edges_south, edges_west ]
    else
        return [ convert(Vector{UInt32}, the_boundary_edges) ]
    end
end

# SVector of MVectors of point IDs representing the 3 edges of a triangle
function edges(face::SVector{3, UInt32})
    edges = SVector( MVector{2, UInt32}(face[1], face[2]),
                     MVector{2, UInt32}(face[2], face[3]),
                     MVector{2, UInt32}(face[3], face[1]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 4 edges of a quadrilateral
function edges(face::SVector{4, UInt32})
    edges = SVector( MVector{2, UInt32}(face[1], face[2]),
                     MVector{2, UInt32}(face[2], face[3]),
                     MVector{2, UInt32}(face[3], face[4]),
                     MVector{2, UInt32}(face[4], face[1]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 3 edges of a quadratic triangle
function edges(face::SVector{6, UInt32})
    edges = SVector( MVector{3, UInt32}(face[1], face[2], face[4]),
                     MVector{3, UInt32}(face[2], face[3], face[5]),
                     MVector{3, UInt32}(face[3], face[1], face[6]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 4 edges of a quadratic quadrilateral
function edges(face::SVector{8, UInt32})
    edges = SVector( MVector{3, UInt32}(face[1], face[2], face[5]),
                     MVector{3, UInt32}(face[2], face[3], face[6]),
                     MVector{3, UInt32}(face[3], face[4], face[7]),
                     MVector{3, UInt32}(face[4], face[1], face[8]) )
    # Order th linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# The unique edges from a vector of triangles or quadrilaterals represented by point IDs
function edges(mesh::UnstructuredMesh_2D)
    edges_filtered = sort(unique(reduce(vcat, edges.(mesh.faces))))
    return [ SVector(e.data) for e in edges_filtered ]
end

# A vector of length 2 SVectors, denoting the face ID each edge is connected to. If the edge
# is a boundary edge, face ID 0 is returned
function edge_face_connectivity(mesh::UnstructuredMesh_2D)
    # Each edge should only border 2 faces if it is an interior edge, and 1 face if it is
    # a boundary edge.
    # Loop through each face in the face_edge_connectivity vector and mark each edge with
    # the faces that it borders.
    if length(mesh.edges) === 0
        @error "Does not have edges!"
    end
    if length(mesh.face_edge_connectivity) === 0
        @error "Does not have face/edge connectivity!"
    end
    edge_face = [MVector{2, UInt32}(0, 0) for _ in eachindex(mesh.edges)]
    for (iface, edges) in enumerate(mesh.face_edge_connectivity)
        for iedge in edges
            # Add the face id in the first non-zero position of the edge_face conn. vec.
            if edge_face[iedge][1] == 0
                edge_face[iedge][1] = iface
            elseif edge_face[iedge][2] == 0
                edge_face[iedge][2] = iface
            else
                @error "Edge $iedge seems to have 3 faces associated with it!"
            end
        end
    end
    return [SVector(sort(two_faces).data) for two_faces in edge_face]
end

# Return an SVector of the points in the edge (Linear)
function edgepoints(edge::SVector{2, UInt32}, points::Vector{Point_2D})
    return SVector(points[edge[1]], points[edge[2]])
end

# Return an SVector of the points in the edge (Quadratic)
function edgepoints(edge::SVector{3, UInt32}, points::Vector{Point_2D})
    return SVector(points[edge[1]], points[edge[2]], points[edge[3]])
end

# Return an SVector of the points in the edge
function edgepoints(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return edgepoints(mesh.edges[edge_id], mesh.points)
end

# Return an SVector of the points in the face (Triangle)
function facepoints(face::SVector{3, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]])
end

# Return an SVector of the points in the face (Quadrilateral)
function facepoints(face::SVector{4, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]], points[face[4]])
end

# Return an SVector of the points in the face (Triangle6)
function facepoints(face::SVector{6, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]],
                   points[face[4]], points[face[5]], points[face[6]])
end

# Return an SVector of the points in the face (Quadrilateral8)
function facepoints(face::SVector{8, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]], points[face[4]],
                   points[face[5]], points[face[6]], points[face[7]], points[face[8]])
end

# Return an SVector of the points in the face
function facepoints(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return facepoints(mesh.faces[edge_id], mesh.points)
end

# Find the faces which share the vertex of ID v.
function faces_sharing_vertex(v::Integer, mesh::UnstructuredMesh_2D)
    shared_faces = UInt32[]
    for i ∈ 1:length(mesh.faces)
        if v ∈  mesh.faces[i]
            push!(shared_faces, UInt32(i))
        end
    end
    return shared_faces
end

# Return the face containing point p.
function findface(p::Point_2D, mesh::UnstructuredMesh_2D)
    if 0 < length(mesh.materialized_faces)
        return UInt32(findface_explicit(p, mesh.materialized_faces))
    else
        return UInt32(findface_implicit(p, mesh.faces, mesh.points))
    end
end

# Find the face containing the point p, with explicitly represented faces
function findface_explicit(p::Point_2D, faces::Vector{<:Face_2D})
    for i ∈ 1:length(faces)
        if p ∈  faces[i]
            return i
        end
    end
    return 0
end

# Return the face containing the point p, with implicitly represented faces
function findface_implicit(p::Point_2D,
                            faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}},
                            points::Vector{Point_2D})
    for i ∈ 1:length(faces)
        bool = p ∈  materialize_face(faces[i], points)
        if bool
            return i
        end
    end
    return 0
end

# Return the intersection algorithm that will be used for l ∩ mesh
function get_intersection_algorithm(mesh::UnstructuredMesh_2D)
    if length(mesh.materialized_edges) !== 0
        return "Edges - Explicit"
    elseif length(mesh.edges) !== 0
        return "Edges - Implicit"
    elseif length(mesh.materialized_faces) !== 0
        return "Faces - Explicit"
    else
        return "Faces - Implicit"
    end
end

# Insert the boundary edge into the correct place in the vector of edge indices, based on
# the distance from some reference point
function insert_boundary_edge!(edge_index::UInt32, p_ref::Point_2D, edge_indices::Vector{UInt32},
                               mesh::UnstructuredMesh_2D)
    # Compute the minimum distance from the edge to be inserted to the reference point
    insertion_distance = minimum(distance.(Ref(p_ref),
                                           edgepoints(mesh.edges[edge_index], mesh.points)))
    # Loop through the edge indices until an edge with greater distance from the reference point
    # is found, then insert
    nindices = length(edge_indices)
    for i ∈ 1:nindices
        epoints = edgepoints(mesh.edges[edge_indices[i]], mesh.points)
        edge_distance = minimum(distance.(Ref(p_ref), epoints))
        if insertion_distance < edge_distance
            insert!(edge_indices, i, edge_index)
            return nothing
        end
    end
    insert!(edge_indices, nindices+1, edge_index)
    return nothing
end

# Intersect a line with the mesh. Returns a vector of intersection points, sorted based
# upon distance from the line's start point
function intersect(l::LineSegment_2D, mesh::UnstructuredMesh_2D)
    # Edges are faster, so they are the default
    if length(mesh.edges) !== 0
        if 0 < length(mesh.materialized_edges)
            return intersect_edges_explicit(l, mesh.materialized_edges)
        else
            return intersect_edges_implicit(l, mesh.edges, mesh.points)
        end
    else
        if 0 < length(mesh.materialized_faces)
            return intersect_faces_explicit(l, mesh.materialized_faces)
        else
            return intersect_faces_implicit(l, mesh.faces, mesh.points)
        end
    end
end

# Intersect a line with an implicitly defined edge
function intersect_edge_implicit(l::LineSegment_2D,
                                 edge::SVector{L, UInt32} where {L},
                                 points::Vector{Point_2D})
    return l ∩ materialize_edge(edge, points)
end

# Intersect a line with materialized edges
function intersect_edges_explicit(l::LineSegment_2D, edges::Vector{LineSegment_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    for edge in edges
        npoints, point = l ∩ edge
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            push!(intersection_points, point)
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with implicitly defined edges
function intersect_edges_explicit(l::LineSegment_2D, edges::Vector{QuadraticSegment_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    for edge in edges
        npoints, points = l ∩ edge
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, points[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with a vector of implicitly defined linear edges
function intersect_edges_implicit(l::LineSegment_2D,
                                  edges::Vector{SVector{2, UInt32}},
                                  points::Vector{Point_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for edge in edges
        npoints, point = intersect_edge_implicit(l, edge, points)
        if 0 < npoints
            push!(intersection_points, point)
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with a vector of implicitly defined quadratic edges
function intersect_edges_implicit(l::LineSegment_2D,
                                  edges::Vector{SVector{3, UInt32}},
                                  points::Vector{Point_2D})
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for edge in edges
        npoints, ipoints = intersect_edge_implicit(l, edge, points)
        if 0 < npoints
            append!(intersection_points, ipoints[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with an implicitly defined face
function intersect_face_implicit(l::LineSegment_2D,
                                 face::SVector{L, UInt32} where {L},
                                 points::Vector{Point_2D})
    return l ∩ materialize_face(face, points)
end

# Intersect a line with explicitly defined linear faces
function intersect_faces_explicit(l::LineSegment_2D, faces::Vector{<:Face_2D} )
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    for face in faces
        npoints, points = l ∩ face
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, points[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with implicitly defined faces
function intersect_faces_implicit(l::LineSegment_2D,
                                  faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}},
                                  points::Vector{Point_2D})
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for face in faces
        npoints, ipoints = intersect_face_implicit(l, face, points)
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, ipoints[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# If a point is a vertex
function isvertex(p::Point_2D, mesh::UnstructuredMesh_2D)
    for point in mesh.points
        if p ≈ point
            return true
        end
    end
    return false
end

# Return a LineSegment_2D from the point IDs in an edge
function materialize_edge(edge::SVector{2, UInt32}, points::Vector{Point_2D})
    return LineSegment_2D(edgepoints(edge, points))
end

# Return a QuadraticSegment_2D from the point IDs in an edge
function materialize_edge(edge::SVector{3, UInt32}, points::Vector{Point_2D})
    return QuadraticSegment_2D(edgepoints(edge, points))
end

# Return a LineSegment_2D or QuadraticSegment_2D
function materialize_edge(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return materialize_edge(mesh.edges[edge_id], mesh.points)
end

# Return a materialized edge for each edge in the mesh
function materialize_edges(mesh::UnstructuredMesh_2D)
    return materialize_edge.(mesh.edges, Ref(mesh.points))
end

# Return a Triangle_2D from the point IDs in a face
function materialize_face(face::SVector{3, UInt32}, points::Vector{Point_2D})
    return Triangle_2D(facepoints(face, points))
end

# Return a Quadrilateral_2D from the point IDs in a face
function materialize_face(face::SVector{4, UInt32}, points::Vector{Point_2D})
    return Quadrilateral_2D(facepoints(face, points))
end

# Return a Triangle6_2D from the point IDs in a face
function materialize_face(face::SVector{6, UInt32}, points::Vector{Point_2D})
    return Triangle6_2D(facepoints(face, points))
end

# Return a Quadrilateral8_2D from the point IDs in a face
function materialize_face(face::SVector{8, UInt32}, points::Vector{Point_2D})
    return Quadrilateral8_2D(facepoints(face, points))
end

# Return an SVector of the points in the edge
function materialize_face(face_id::UInt32, mesh::UnstructuredMesh_2D)
    return materialize_face(mesh.faces[face_id], mesh.points)
end

# Return a materialized face for each face in the mesh
function materialize_faces(mesh::UnstructuredMesh_2D)
    return materialize_face.(mesh.faces, Ref(mesh.points))
end

# Return the number of edges in a face
function num_edges(face::SVector{L, UInt32}) where {L}
    if L % 3 === 0
        return 0x00000003
    elseif L % 4 === 0
        return 0x00000004
    else
        # Error
        return 0x00000000
    end
end

function remap_points_to_hilbert(points::Vector{Point_2D})
    bb = boundingbox(points)
    npoints = length(points)
    # Generate a Hilbert curve
    hilbert_points = hilbert_curve(bb, npoints)
    nhilbert_points = length(hilbert_points)
    # For each point, get the index of the hilbert points that is closest
    point_indices = Vector{Int64}(undef, npoints)
    for i ∈ 1:npoints
        min_distance = 1.0e10
        for j ∈ 1:nhilbert_points
            pdistance = distance²(points[i], hilbert_points[j])
            if pdistance < min_distance
                min_distance = pdistance
                point_indices[i] = j
            end
        end
    end
    return sortperm(point_indices)
end

function reorder_points_to_hilbert!(mesh::UnstructuredMesh_2D)
    # Points
    # point_map     maps  new_points[i] == mesh.points[point_map[i]]
    # point_map_inv maps mesh.points[i] == new_points[point_map_inv[i]]
    point_map  = remap_points_to_hilbert(mesh.points)
    point_map_inv = UInt32.(sortperm(point_map))
    # reordered to resemble a hilbert curve
    permute!(mesh.points, point_map)

    # Adjust face indices
    # Point IDs have changed, so we need to change the point IDs referenced by the faces
    new_faces_vec = [ point_map_inv[face] for face in mesh.faces]
    for i in 1:length(mesh.faces)
        mesh.faces[i] = SVector(point_map_inv[mesh.faces[i]])
    end
    return nothing
end

function reorder_faces_to_hilbert!(mesh::UnstructuredMesh_2D)
    face_map  = UInt32.(remap_points_to_hilbert(centroid.(materialize_faces(mesh))))
    permute!(mesh.faces, face_map)
    return nothing
end

function reorder_to_hilbert!(mesh::UnstructuredMesh_2D)
    reorder_points_to_hilbert!(mesh)
    reorder_faces_to_hilbert!(mesh)
    return nothing
end

# Return the ID of the edge shared by two adjacent faces
function shared_edge(face1::UInt32, face2::UInt32, mesh::UnstructuredMesh_2D)
    for edge1 in mesh.face_edge_connectivity[face1]
        for edge2 in mesh.face_edge_connectivity[face2]
            if edge1 == edge2
                return edge1
            end
        end
    end
    return 0x00000000
end

# How to display a mesh in REPL
function Base.show(io::IO, mesh::UnstructuredMesh_2D)
    mesh_type = typeof(mesh)
    println(io, mesh_type)
    println(io, "  ├─ Name      : $(mesh.name)")
    size_MB = Base.summarysize(mesh)/1E6
    if size_MB < 1
        size_KB = size_MB*1000
        println(io, "  ├─ Size (KB) : $size_KB")
    else
        println(io, "  ├─ Size (MB) : $size_MB")
    end
    println(io, "  ├─ Points    : $(length(mesh.points))")
    nedges = length(mesh.edges)
    println(io, "  ├─ Edges     : $nedges")
    println(io, "  │  ├─ Linear         : $(count(x->x isa SVector{2, UInt32},  mesh.edges))")
    println(io, "  │  ├─ Quadratic      : $(count(x->x isa SVector{3, UInt32},  mesh.edges))")
    println(io, "  │  └─ Materialized?  : $(length(mesh.materialized_edges) !== 0)")
    println(io, "  ├─ Faces     : $(length(mesh.faces))")
    println(io, "  │  ├─ Triangle       : $(count(x->x isa SVector{3, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Quadrilateral  : $(count(x->x isa SVector{4, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Triangle6      : $(count(x->x isa SVector{6, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Quadrilateral8 : $(count(x->x isa SVector{8, UInt32},  mesh.faces))")
    println(io, "  │  └─ Materialized?  : $(length(mesh.materialized_faces) !== 0)")
    println(io, "  ├─ Connectivity")
    println(io, "  │  ├─ Edge/Face : $(0 < length(mesh.edge_face_connectivity))")
    println(io, "  │  └─ Face/Edge : $(0 < length(mesh.face_edge_connectivity))")
    println(io, "  ├─ Boundary edges")
    nsides = length(mesh.boundary_edges)
    if nsides !== 0
        println(io, "  │  ├─ Edges : $(mapreduce(x->length(x), +, mesh.boundary_edges))")
    else
        println(io, "  │  ├─ Edges : 0")
    end
    println(io, "  │  └─ Sides : $nsides")
    println(io, "  └─ Face sets : $(length(keys(mesh.face_sets)))")
end

# Sort intersection points, deleting points that are less than minimum_segment_length apart
function sort_intersection_points!(p::Point_2D, points::Vector{Point_2D})
    if 2 <= length(points)
        sortpoints!(p, points)
        # Eliminate any points for which the distance between consecutive points
        # is less than the minimum segment length
        delete_ids = Int64[]
        id_start = 1
        for id_stop ∈ 2:length(points)
            if distance²(points[id_start], points[id_stop]) < minimum_segment_length^2
                push!(delete_ids, id_stop)
            else
                id_start = id_stop
            end
        end
        deleteat!(points, delete_ids)
    else
        return points
    end
end

# Return a mesh with name name, composed of the faces in the set face_ids
function submesh(name::String, mesh::M) where {M <: UnstructuredMesh_2D}
    # Setup faces and get all vertex ids
    face_ids = mesh.face_sets[name]
    submesh_faces = Vector{Vector{UInt32}}(undef, length(face_ids))
    vertex_ids = Set{UInt32}()
    for (i, face_id) in enumerate(face_ids)
        face_vec = collect(mesh.faces[face_id].data)
        submesh_faces[i] = face_vec
        union!(vertex_ids, Set(face_vec))
    end
    # Need to remap vertex ids in faces to new ids
    vertex_ids_sorted = sort(collect(vertex_ids))
    vertex_map = Dict{UInt32, UInt32}()
    for (i,v) in enumerate(vertex_ids_sorted)
        vertex_map[v] = i
    end
    submesh_points = Vector{Point_2D}(undef, length(vertex_ids_sorted))
    for (i, v) in enumerate(vertex_ids_sorted)
        submesh_points[i] = mesh.points[v]
    end
    # remap vertex ids in faces
    for face in submesh_faces
        for (i, v) in enumerate(face)
            face[i] = vertex_map[v]
        end
    end
    # At this point we have points, faces, & name.
    # Just need to get the face sets
    submesh_face_sets = Dict{String, Set{UInt32}}()
    for face_set_name in keys(mesh.face_sets)
        set_intersection = mesh.face_sets[face_set_name] ∩ face_ids
        if length(set_intersection) !== 0
            submesh_face_sets[face_set_name] = set_intersection
        end
    end
    # Need to remap face ids in face sets
    face_map = Dict{UInt32, UInt32}()
    for (i,f) in enumerate(face_ids)
        face_map[f] = i
    end
    for face_set_name in keys(submesh_face_sets)
        new_set = Set{UInt32}()
        for fid in submesh_face_sets[face_set_name]
            union!(new_set, face_map[fid])
        end
        submesh_face_sets[face_set_name] = new_set
    end
    return M(name = name,
             points = submesh_points,
             faces = [SVector{length(f), UInt32}(f) for f in submesh_faces],
             face_sets = submesh_face_sets
            )
end

# Plot
# -------------------------------------------------------------------------------------------------
if enable_visualization
    function convert_arguments(LS::Type{<:LineSegments}, mesh::UnstructuredMesh_2D)
        if 0 < length(mesh.materialized_edges)
            return convert_arguments(LS, mesh.materialized_edges)
        elseif 0 < length(mesh.edges)
            return convert_arguments(LS, materialize_edges(mesh))
        else
            return convert_arguments(LS, materialize_faces(mesh))
        end
    end

    function convert_arguments(P::Type{<:Mesh}, mesh::UnstructuredMesh_2D)
        if 0 < length(mesh.materialized_faces)
            return convert_arguments(P, mesh.materialized_faces)
        else
            return convert_arguments(P, materialize_faces(mesh))
        end
    end
end


function build!(m::EnergyModel)
    # If subnetworks are not set, compute them
    isempty(m.subnetworks) && determine_subnetworks!(m)

    # Clear jump model
    empty!(m.jumpobjects)
    MOI.empty!(backend(m.jumpmodel))
    for vartype in [:fix, :integrality, :lower_bound, :upper_bound, :zero_one]
        empty!(getfield(m.jumpmodel, Symbol(:variable_to_, vartype)))
    end

    # Fill it up again
    addto!(m.jumpmodel, m)
end

function addto!(jm::JuMP.AbstractModel, m::EnergyModel)
    # @info("* Equations for individual devices")
    for d = devices(m)
        # @info("  - $(naming(d)) ($(naming(typeof(d))))")
        addto!(jm, m, d)
    end
    # @info("* Equations for energy balance")
    for d = buses(m)
        addto!(jm, m, d)
    end
    # @info("* Equations for subnetworks")
    for sn = subnetworks(m)
        addto!(jm, m, sn)
    end
    # @info("* Cost minimization objective")
    @objective(jm, Min, sum(cost(d) for d = devices(m)))

    jm
end


using ProgressMeter
using Base.Threads

"""
    filter_positions(snap_file::String, corner_lowerleft::Array{<:Real}, corner_upperright::Array{<:Real},
                          parttype::Integer)

Reads positions from `snap_file` and returns the indices of particles contained in a box defined by `corner_lowerleft` and `corner_upperright`.
"""
function filter_positions(snap_file::String, corner_lowerleft::Array{<:Real}, corner_upperright::Array{<:Real},
                          parttype::Integer)

    pos = read_snap(snap_file, "POS", parttype)
    sel = findall( @views ( ( corner_lowerleft[1] .<= pos[1,:] .< corner_upperright[1] ) .&
                            ( corner_lowerleft[2] .<= pos[2,:] .< corner_upperright[2] ) .&
                            ( corner_lowerleft[3] .<= pos[3,:] .< corner_upperright[3] )) )
    return sel
end

"""
    filter_ids(snap_file::String, selected_ids::Vector{<:Integer})

Reads IDs from `snap_file` and returns the indices of particles matching the `selected_ids`
"""
function filter_ids(snap_file::String, selected_ids::Vector{<:Integer}, parttype::Integer)

    # read ids in file
    ids = read_snap(snap_file, "ID", parttype)

    # check which IDs match
    return filter_by_ids( selected_ids, ids)
end

"""
    read_blocks_over_all_files( snap_base::String, blocks::Array{String};
                                filter_function::Union{Function, Nothing}=nothing,
                                read_positions::Union{Dict, Nothing}=nothing,
                                parttype::Integer=0, verbose::Bool=true )

Reads the specified blocks from all distributed files where particles pass the `filter_function`, or are given by a `Dict` of `read_positions`.
For `read_positions` please see [`find_read_positions`](@ref).
"""
function read_blocks_over_all_files(snap_base::String, blocks::Array{String};
                                    filter_function::Union{Function, Nothing}=nothing,
                                    read_positions::Union{Dict, Nothing}=nothing,
                                    parttype::Integer=0, verbose::Bool=true )

    # default behaviour is to find the particles to read.
    # If this number is known in advance it can be given as an input parameter.
    if isnothing(read_positions)

        # this only works if a filter function is provided!
        if isnothing(filter_function)
            error("Please provide either a dictionary with read positions or a filter function!")
        end

        # check how many particles fulfill the filter criterion and find the positions in the blocks.
        # -> needed for array pre-allocation to speed up read-in.
        read_positions = find_read_positions(snap_base, filter_function, verbose=verbose)
    end

    if verbose
        @info "Allocating storage dictionary..."
        t1 = time_ns()
    end

    # select current file
    filename = select_file(snap_base, 0 )

    # read info block
    snap_info = read_info(filename)

    # pre-allocate all data arrays in a dictionary
    d = allocate_data_dict(blocks, read_positions["N_part"], snap_info, false)

    if verbose
        t2 = time_ns()
        @info "  elapsed: $(output_time(t1,t2)) s"
    end

    # store the number of particles that have been read
    N_read = 0

    # get the relevant files from the dictionary
    key_entries = collect(keys(read_positions))
    files = key_entries[ key_entries .!= "N_part"]

    sort!(files)

    if verbose
        @info "Reading $(size(files,1)) snapshots..."
        t1 = time_ns()
    end

    for file ∈ files

        # select current file
        filename = select_file(snap_base, file )

        # read header block
        h = read_header(filename)

        # read info block
        snap_info = read_info(filename)

        # find the location of the blocks in the current file
        file_block_positions = get_block_positions(filename)

        # read the blocks
        @threads for block ∈ blocks

            # read the info ot the curent block
            block_info = snap_info[getfield.(snap_info, :block_name) .== block][1]

            # add offset of particle types that should not be read
            offset = 0
            for i = 1:size(h.npart,1)
                if (block_info.is_present[i] > 0) & (h.npart[i] > 0) & ( i < parttype + 1)
                    offset += h.npart[i]
                end
            end

            # store the relevant data in the `d[block]` array
            read_block_with_offset!(d[block], N_read, filename,
                                    file_block_positions[block],
                                    block_info, offset,
                                    read_positions[file]["index"],
                                    read_positions[file]["n_to_read"])

        end
        # count the number of particles already read
        N_read += sum(read_positions[file]["n_to_read"])

        if verbose
            @info "Read $N_read / $(read_positions["N_part"]) particles"
        end

    end

    if verbose
        t2 = time_ns()
        @info "  elapsed: $(output_time(t1,t2)) s"
    end

    return d
end


using HorizonSideRoots
#include("")
r = Robot()


module Wflow

using Dates
using TOML
using LightGraphs
using NCDatasets
using StaticArrays
using Statistics
using UnPack
using CSV
using Random
using DataFrames

mutable struct Clock
    time::DateTime
    iteration::Int
    Δt::Second
end

include("io.jl")

"""
    Model{N,L,V,R,W}

Composite type that represents all different aspects of a Wflow Model, such as the
network, parameters, clock, configuration and input and output.
"""
struct Model{N,L,V,R,W}
    config::Config  # all configuration options
    network::N  # connectivity information, directed graph
    lateral::L  # lateral model that holds lateral state, moves along network
    vertical::V  # vertical model that holds vertical state, independent of each other
    clock::Clock  # to keep track of simulation time
    reader::R  # provides the model with dynamic input
    writer::W  # writes model output
end

# prevent a large printout of model components and arrays
Base.show(io::IO, m::Model) = print(io, "model of type ", typeof(m))

include("horizontal_process.jl")
include("hbv.jl")
include("sbm.jl")
include("reservoir_lake.jl")
include("hbv_model.jl")
include("sbm_model.jl")
include("flow.jl")
include("vertical_process.jl")
include("groundwater/connectivity.jl")
include("groundwater/aquifer.jl")
include("groundwater/boundary_conditions.jl")
include("sbm_gwf_model.jl")
include("utils.jl")

"""
    run_simulation(tomlpath::String)
    run_simulation(config::Config)
    run_simulation(model::Model)

Run an entire simulation starting either from a path to a TOML settings file,
a prepared `Config` object, or an initialized `Model` object. This allows more flexibility
if you want to for example modify a `Config` before initializing the `Model`.
"""
function run_simulation(tomlpath)
    config = Config(tomlpath)
    run_simulation(config)
end

function run_simulation(config::Config)
    modeltype = config.model.type

    model = if modeltype == "sbm"
        initialize_sbm_model(config)
    elseif modeltype == "sbm_gwf"
        initialize_sbm_gwf_model(config)
    elseif modeltype == "hbv"
        initialize_hbv_model(config)
    else
        error("unknown model type")
    end

    run_simulation(model)
end

function run_simulation(model::Model; close_files = true)
    @unpack network, config, writer, clock = model

    # in the case of sbm_gwf it's currently a bit hard to use dispatch
    update_func = config.model.type == "sbm_gwf" ? update_sbm_gwf : update

    while true
        model = update_func(model)
        is_finished(clock, config) && break
    end

    # write output state NetCDF
    # undo the clock advance at the end of the last iteration, since there won't
    # be a next step, and then the output state falls on the correct time
    rewind!(clock)
    write_netcdf_timestep(model, writer.state_dataset, writer.state_parameters)

    reset_clock!(model.clock, config)

    # option to support running function twice without re-initializing
    # and thus opening the NetCDF files
    if close_files
        Wflow.close_files(model, delete_output = false)
    end
    return model
end

end # module


#using Morbit
#using Test

@testset "batch_objectives" begin
global counter_1 = 0;
global counter_2 = 0;

f = x -> sum( x.^2 )

# this function does NOT handle array of vector inputs and uses standard broadcasting
function objf1( x_vec )
    global counter_1
    counter_1 += 1;
    f(x_vec)
end

# this function handles batch evaluation all by itself,
# and broadcasting should be overwritten internally
function objf2( x_vec_or_array )
    global counter_2
    counter_2 += 1;
    if isa( x_vec_or_array , Vector{<:Real} )
        return f(x_vec_or_array)
    elseif isa(x_vec_or_array , Vector{<:Vector{<:Real}} )
        return [ f(x) for x ∈ x_vec_or_array ]
    end
end

n_vars = 3
sites = [ rand(n_vars) for i = 1 : 5 ];

mop = MixedMOP(n_vars)
add_objective!(mop, objf1, :expensive, 1, false)
add_objective!(mop, objf2, :expensive, 1, true)

Y = Morbit.eval_all_objectives.(mop, sites)

@test counter_1 == 5
@test counter_2 == 1
@test all( y[1] == y[2] for y in Y )
end


############################################################
# judiExtendedSource ## ####################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export judiExtendedSource, judiExtendedSourceException

############################################################

struct judiExtendedSource{vDT<:Number} <: joAbstractLinearOperator{vDT,vDT}
    name::String
    m::Integer
    n::Integer
    info::Info
    wavelet
    weights
end

mutable struct judiExtendedSourceException <: Exception
    msg :: String
end

############################################################

## outer constructors

"""
    judiExtendedSource
        name::String
        m::Integer
        n::Integer
        info::Info
        wavelet
        weights
Abstract sparse vector for an extended source. The `judiRHS` vector has the\\
dimensions of the full time history of the wavefields, but contains only the wavelet, \\
as well as the weights of the extended source.
Constructor
==========
    judiExtendedSource(info, wavelet, weights)
Examples
========
Assuming `Pw` ia a projection operators of type `judiLRWF` and `w` is a seismic weight \\
vector of type `judiWeights`, then a `judiExtendedSource` vector can be created as follows:
    q_ext = adjoint(Pw)*w    # abstract extended source
"""
function judiExtendedSource(info,wavelet,weights;vDT::DataType=Float32)
    vDT == Float32 || throw(judiExtendedSourceException("Domain type not supported"))
    # length of vector
    m = info.n * sum(info.nt)
    n = 1
    return judiExtendedSource{Float32}("judiExtendedSource",m,n,info,wavelet,weights)
end

####################################################################
## overloaded Base functions

# conj(jo)
conj(A::judiExtendedSource{vDT}) where vDT =
    judiExtendedSource{vDT}("conj("*A.name*")",A.m,A.n,A.info,A.wavelet,A.weights)

# transpose(jo)
transpose(A::judiExtendedSource{vDT}) where vDT =
    judiExtendedSource{vDT}(""*A.name*".'",A.n,A.m,A.info,A.wavelet,A.weights)

# adjoint(jo)
adjoint(A::judiExtendedSource{vDT}) where vDT =
    judiExtendedSource{vDT}(""*A.name*"'",A.n,A.m,A.info,A.wavelet,A.weights)

####################################################################

# +(judiExtendedSource,judiExtendedSource)
function +(A::judiExtendedSource{avDT}, B::judiExtendedSource{bvDT}) where {avDT, bvDT}
    # Error checking
    size(A) == size(B) || throw(judiExtendedSourceException("Shape mismatch: A:$(size(A)), v: $(size(B))"))
    compareInfo(A.info, B.info) == true || throw(judiExtendedSourceException("info mismatch"))

    # Size
    m = A.info.n * sum(A.info.nt)
    n = 1

    # wavelet and weights
    A.wavelet == B.wavelet ? nothing : throw(judiExtendedSourceException("Can only add two extended
                                                                          sources with same wavelet"))
    wavelet = A.wavelet
    weights = A.weights .+ B.weights

    nvDT = promote_type(avDT,bvDT)

    return judiExtendedSource{nvDT}("judiRHS",m,n,A.info, wavelet, weights)
end

# -(judiExtendedSource,judiExtendedSource)
function -(A::judiExtendedSource{avDT}, B::judiExtendedSource{bvDT}) where {avDT, bvDT}

        # Error checking
        size(A) == size(B) || throw(judiExtendedSourceException("Shape mismatch: A:$(size(A)), v: $(size(B))"))
        compareInfo(A.info, B.info) == true || throw(judiExtendedSourceException("info mismatch"))

        # Size
        m = A.info.n * sum(A.info.nt)
        n = 1

        # wavelet and weights
        A.wavelet == B.wavelet ? nothing : throw(judiExtendedSourceException("Can only add two extended
                                                                            sources with same wavelet"))
        wavelet = A.wavelet
        weights = A.weights .- B.weights

        nvDT = promote_type(avDT,bvDT)

        return judiExtendedSource{nvDT}("judiRHS",m,n,A.info, wavelet, weights)
    end

function subsample(a::judiExtendedSource{avDT},srcnum) where avDT
    info = Info(a.info.n, length(srcnum), a.info.nt[srcnum])
    return judiExtendedSource(info,a.wavelet[srcnum],a.weights[srcnum];vDT=avDT)
end

getindex(x::judiExtendedSource,a) = subsample(x,a)


"""
Module providing parallelised versions of [`QuantumOptics.timeevolution.mcwf`](@ref).

`ParallelMCWF` exports functions:
[`pmcwf`](@ref), [`load_trajs`](@ref), [`save_trajs`](@ref),
[`kets_to_dm`](@ref), [`kets_to_obs`](@ref)
"""
module ParallelMCWF

using Distributed, Base.Threads
if myid() == 1
    @info "ParallelMCWF loaded with $(nthreads()) threads per process."
    @info "ParallelMCWF loaded with $(nprocs()) workers."
    #@info "Processes must be added BEFORE using `using ParallelMCWF`, see Julia issue #3674."
end
using ProgressMeter, JLD2
import OrdinaryDiffEq
using QuantumOptics.bases, QuantumOptics.states, QuantumOptics.operators
using QuantumOptics.operators_dense, QuantumOptics.operators_sparse
using QuantumOptics.timeevolution
using QuantumOptics.operators_lazysum, QuantumOptics.operators_lazytensor, QuantumOptics.operators_lazyproduct
@everywhere using QuantumOptics.timeevolution.timeevolution_mcwf, QuantumOptics.timeevolution
const DecayRates = Union{Vector{Float64}, Matrix{Float64}, Nothing}
Base.@pure pure_inference(fout,T) = Core.Compiler.return_type(fout, T)

include("trajs_IO.jl")
export load_trajs, save_trajs
include("trajs_processing.jl")
export kets_to_dm, kets_to_obs
include("pmcwf.jl")
export pmcwf

end # module

using Distributed
@everywhere using QuantumOptics # TO DO: Can be avoided


"""
    PointCloud{T<:AbstractFloat, I<:Integer}(dataset::DataSet{T, I}, hyperrect::HyperRectVolume{T}, searchpts::Bool = false)::PointCloud

*AHMI-internal, not part of stable public API.*

creates a point cloud by searching the data tree for points which are inside the hyper-rectangle
The parameter searchpts determines if an array of the point IDs is created as well
"""
function PointCloud(
    dataset::DataSet{T, I},
    hyperrect::HyperRectVolume{T},
    searchpts::Bool)::PointCloud{T, I} where {T<:AbstractFloat, I<:Integer}

    result = PointCloud(T, I)

    PointCloud!(result, dataset, hyperrect, searchpts)

    result
end

function PointCloud!(
    cloud::PointCloud{T, I},
    dataset::DataSet{T, I},
    hyperrect::HyperRectVolume{T},
    searchpts::Bool) where {T<:AbstractFloat, I<:Integer}

    search!(cloud.searchres, dataset, hyperrect, searchpts)

    cloud.points = cloud.searchres.points

    resize!(cloud.pointIDs, cloud.points)
    copyto!(cloud.pointIDs, cloud.searchres.pointIDs)

    cloud.maxWeightProb = cloud.searchres.maxWeightProb
    cloud.minWeightProb = cloud.searchres.minWeightProb
    cloud.maxLogProb = cloud.searchres.maxLogProb
    cloud.minLogProb = cloud.searchres.minLogProb

    cloud.probfactor = exp(cloud.maxLogProb - cloud.minLogProb)
    cloud.probweightfactor = exp(cloud.maxWeightProb - cloud.minWeightProb)

    nothing
end

function create_pointweights(
    dataset::DataSet{T, I},
    volumes::Vector{IntegrationVolume{T, I, V}},
    ids::Array{I})::Vector{T} where {T<:AbstractFloat, I<:Integer, V<:SpatialVolume}

    pweights = zeros(T, dataset.N)

    for i in ids
        for p in eachindex(volumes[i].pointcloud.pointIDs)
            pweights[volumes[i].pointcloud.pointIDs[p]] += 1
        end
    end

    pweights
end

function Base.copy!(
    target::PointCloud{T, I},
    src::PointCloud{T, I}) where {T<:AbstractFloat, I<:Integer}

    target.points = src.points

    resize!(target.pointIDs, length(src.pointIDs))
    copyto!(target.pointIDs, src.pointIDs)

    target.maxLogProb = src.maxLogProb
    target.minLogProb = src.minLogProb

    target.maxWeightProb = src.maxWeightProb
    target.minWeightProb = src.minWeightProb

    target.probfactor = src.probfactor
    target.probweightfactor = src.probweightfactor

    nothing
end


using JSON

mutable struct Randomizer
    randomization_config::Dict
    default_config::Dict
    key_s::Set
end

function Randomizer(randomization_config_fp="default_dr.json", default_config_fp="default.json")
    randomization_config, default_config = nothing, nothing

    try
        open(get_file_path("src/randomization/config", randomization_config_fp, "json"), "r") do f
            randomization_config = json.load(f)
        end
    catch
        #logger.warning("Couldn't find {} in randomization/config subdirectory".format(randomization_config_fp))
        randomization_config = Dict()
    end

    open(get_file_path("src/randomization/config", default_config_fp, "json"), "r") do f
        default_config = JSON.parse(f)
    end

    key_s = Set(append!(collect.(keys.((randomization_config, default_config)))...))

    Randomizer(randomization_config, default_config, key_s)
end

function randomize(randomizer::Randomizer)
    ##Returns a dictionary of randomized parameters, with key: parameter name and value: randomized
    #value
    ##
    randomization_settings = Dict()

    for k in randomizer.key_s
        setting = nothing
        if k in keys(randomizer.randomization_config)
            randomization_definition = randomizer.randomization_config[k]

            if randomization_definition["type"] == "int"
                try
                    low = randomization_definition["low"]
                    high = randomization_definition["high"]
                    size = get(randomization_definition, "size", 1)
                catch
                    throw(KeyError("Please check your randomization definition for: $k"))
                end
                setting = rand(low:high, size)

            elseif randomization_definition["type"] == "uniform"
                try
                    low = randomization_definition["low"]
                    high = randomization_definition["high"]
                    size = get(randomization_definition, "size", 1)
                catch
                    throw(KeyError("Please check your randomization definition for: $k"))
                end

                setting = rand(low:high, size)

            elseif randomization_definition["type"] == "normal"
                try
                    loc = randomization_definition["loc"]
                    scale = randomization_definition["scale"]
                    size = get(randomization_definition, "size", 1)
                catch
                    throw(KeyError("Please check your randomization definition for: $k"))
                end

                setting = rand(Normal(loc, scale), size)

            else
                throw(NotImplementedError("You've specified an unsupported distribution type"))
            end

        elseif k in keys(randomizer.default_config)
            randomization_definition = randomizer.default_config[k]
            setting = randomization_definition["default"]
        end

        randomization_settings[k] = setting
    end
    return randomization_settings
end


using LightDarkLPDM
using Base.Test
using POMDPs
using POMDPToolbox
using StaticArrays
using Plots
using ParticleFilters

#Removed test code from here since it was for LightDarkPOMDPs.


############################################################
# judiWeights ##############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: June 2019

# Updated by Ziyi Yin (ziyi.yin@gatech.edu), Nov 2020

export judiWeights

############################################################

# structure for seismic data as an abstract vector
mutable struct judiWeights{T<:Number} <: judiMultiSourceVector{T}
    nsrc::Integer
    data::Vector{Array{T, N}} where N
end

# Bypass mismatch in naming and fields for backward compat
Base.getproperty(obj::judiWeights, sym::Symbol) = sym == :weights ? getfield(obj, :data) : getfield(obj, sym)

############################################################

## outer constructors

"""
    judiWeights
        nsrc
        weights
Abstract vector for weighting an extended source, which is injected at every grid point, as weighted by this vector.
Constructors
============
Construct vector cell array of weights. The `weights` keyword\\
can also be a single (non-cell) array, in which case the weights are the same for all source positions:
    judiWeights(weights; nsrc=1)
"""
function judiWeights(weights::Array{T, N}; nsrc=1) where {T<:AbstractFloat, N}
    weights = convert(Array{Float32}, weights)
    # length of vector
    weightsCell = [deepcopy(weights) for j=1:nsrc]
    return judiWeights{Float32}(nsrc, weightsCell)
end

# constructor if weights are passed as a cell array
judiWeights(weights::Vector{Array}) =
    judiWeights(convert.(Array{Float32, length(size(weights[1]))}, weights))

function judiWeights(weights::Vector{Array{T, N}}) where {T<:Number, N}
    nsrc = length(weights)
    weights = convert.(Array{Float32, N}, weights)
    return judiWeights{Float32}(nsrc, weights)
end

############################################################
# JOLI conversion
jo_convert(::Type{T}, jw::judiWeights{T}, ::Bool) where {T<:AbstractFloat} = jw
jo_convert(::Type{T}, jw::judiWeights{vT}, B::Bool) where {T<:AbstractFloat, vT} = judiWavefield{T}(jv.nsrc, jo_convert.(T, jw.weights, B))
zero(::Type{T}, v::judiWeights{vT}; nsrc::Integer=v.nsrc) where {T, vT} = judiWeights{T}(nsrc, T(0) .* v.data[1:nsrc])
(w::judiWeights)(x::Vector{<:Array}) = judiWeights(x)

function copy!(jv::judiWeights, jv2::judiWeights)
    jv.data .= jv2.data
    jv
end
copyto!(jv::judiWeights, jv2::judiWeights) = copy!(jv, jv2)

function push!(a::judiWeights{T}, b::judiWeights{T}) where T
	append!(a.weights, b.weights)
	a.nsrc += b.nsrc
end

make_input(w::judiWeights) = w.data[1]

# getindex weights container
"""
    getindex(x,source_numbers)
getindex seismic weights vectors or matrix-free linear operators and extract the entries that correspond\\
to the shot positions defined by `source_numbers`. Works for inputs of type `judiWeights`, `judiModeling`, \\
`judiProjection`, `judiJacobian`, `Geometry`, `judiRHS`, `judiPDE`, `judiPDEfull`.
Examples
========
(1) Extract 2 shots from `judiWeights` vector:
    dsub = getindex(dobs,[1,2])
(2) Extract geometry for shot location 100:
    geometry_sub = getindex(dobs.geometry,100)
(3) Extract Jacobian for shots 10 and 20:
    Jsub = getindex(J,[10,20])
"""
getindex(a::judiWeights{avDT}, srcnum::RangeOrVec) where avDT = judiWeights{avDT}(length(srcnum), a.data[srcnum])


"""Evaluates the stability of a set of SCM configurations for a parameter vector and possibly draws a new one."""

using ArgParse
using Distributions
using CalibrateEDMF
using CalibrateEDMF.DistributionUtils
using CalibrateEDMF.ReferenceModels
using CalibrateEDMF.ReferenceStats
using CalibrateEDMF.TurbulenceConvectionUtils
src_dir = dirname(pathof(CalibrateEDMF))
using CalibrateEDMF.HelperFuncs
# Import EKP modules
using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
using JLD2

# Read iteration number of ensemble to be recovered
s = ArgParseSettings()
@add_arg_table s begin
    "--version"
    help = "Calibration process number"
    arg_type = Int
    "--job_dir"
    help = "Job output directory"
    arg_type = String
    default = "output"
end
parsed_args = parse_args(ARGS, s)
version = parsed_args["version"]
outdir_path = parsed_args["job_dir"]
include(joinpath(outdir_path, "config.jl"))
config = get_config()
namelist_args = get_entry(config["scm"], "namelist_args", nothing)

scm_args = load(scm_init_path(outdir_path, version))
priors = deserialize_prior(load(joinpath(outdir_path, "prior.jld2")))
ekobj = load(ekobj_path(outdir_path, 1))["ekp"]

# Preconditioning ensemble methods in unconstrained space
if isa(ekobj.process, Inversion) || isa(ekobj.process, Sampler)
    model_evaluator = precondition(scm_args["model_evaluator"], priors, namelist_args = namelist_args)
    batch_indices = scm_args["batch_indices"]
    rm(scm_init_path(outdir_path, version))
    jldsave(scm_init_path(outdir_path, version); model_evaluator, version, batch_indices)
end


#
# Example of a set of reference values that StateCheck will use to test current vlaues against.
# These are generated by the function scprintref() and parsed by the scdocheck()
#

#! format: off
varr = [
 [ "My gradients", "ν∇u[1]",  1.34348869323730468750e-04,  9.84732866287231445313e-01,  5.23545503616333007813e-01,  3.08209930764271777814e-01 ],
 [ "My gradients", "ν∇u[2]",  1.16317868232727050781e-01,  9.92088317871093750000e-01,  4.83800649642944335938e-01,  2.83350456014221541157e-01 ],
 [ "My gradients", "ν∇u[3]",  1.05845928192138671875e-03,  9.51775908470153808594e-01,  4.65474426746368408203e-01,  2.73615551085745090099e-01 ],
 [ "My gradients", "ν∇u[4]",  5.97668886184692382813e-02,  9.68048095703125000000e-01,  5.42618036270141601563e-01,  2.81570862027933854765e-01 ],
 [ "My gradients", "ν∇u[5]",  8.31030607223510742188e-02,  9.35931921005249023438e-01,  5.05405902862548828125e-01,  2.46073509972619536290e-01 ],
 [ "My gradients", "ν∇u[6]",  3.09681892395019531250e-02,  9.98341441154479980469e-01,  4.54375565052032470703e-01,  3.09461067853178561915e-01 ],
 [ "My gradients", "κ∇θ[1]",  8.47448110580444335938e-02,  9.94180679321289062500e-01,  5.27157366275787353516e-01,  2.92455951648181833313e-01 ],
 [ "My gradients", "κ∇θ[2]",  1.20514631271362304688e-02,  9.93527650833129882813e-01,  4.71063584089279174805e-01,  2.96449027197666359346e-01 ],
 [ "My gradients", "κ∇θ[3]",  8.14980268478393554688e-02,  9.55443382263183593750e-01,  5.05038917064666748047e-01,  2.77201022741208891187e-01 ],
 [    "My fields",   "u[1]",  4.31410233294131639781e-02,  9.97140933049696531754e-01,  4.62139750850942054861e-01,  3.23076684924287371725e-01 ],
 [    "My fields",   "u[2]",  1.01416659908237782872e-02,  9.14712023896926407218e-01,  4.76160523012988778913e-01,  2.71443440757963339038e-01 ],
 [    "My fields",   "θ[1]",  6.58965491052394547467e-02,  9.73216404386510802738e-01,  4.60007166313864512830e-01,  2.87310472114545079059e-01 ],
]
parr = [
 [ "My gradients", "ν∇u[1]",    16,     7,    16,     0 ],
 [ "My gradients", "ν∇u[2]",    16,     7,    16,     0 ],
 [ "My gradients", "ν∇u[3]",    16,     7,    16,     0 ],
 [ "My gradients", "ν∇u[4]",    16,     7,    16,     0 ],
 [ "My gradients", "ν∇u[5]",    16,     7,    16,     0 ],
 [ "My gradients", "ν∇u[6]",    16,     7,    16,     0 ],
 [ "My gradients", "κ∇θ[1]",    16,    16,    16,     0 ],
 [ "My gradients", "κ∇θ[2]",    16,    16,    16,     0 ],
 [ "My gradients", "κ∇θ[3]",    16,    16,    16,     0 ],
 [    "My fields",   "u[1]",    16,    16,    16,     0 ],
 [    "My fields",   "u[2]",    16,    16,    16,     0 ],
 [    "My fields",   "θ[1]",    16,    16,    16,     0 ],
]
#! format: on


module BERT

export
    BertPreTraining,
    BertClassification,
    BertConfig,
    load_from_torch_base,
    load_from_torch_pretraining,
    load_from_torch_classification,
    BertAdam,
    bert_tokenize

using Knet, SpecialFunctions, LinearAlgebra
include("model.jl")
include("optimizer.jl")
include("preprocess.jl")

end # module


using Oceananigans.Grids: topology
using CUDA

include("utils_for_runtests.jl")
include("data_dependencies.jl")

archs = test_architectures()

function get_fields_from_checkpoint(filename)
    file = jldopen(filename)

    tracers = keys(file["tracers"])
    tracers = Tuple(Symbol(c) for c in tracers)

    velocity_fields = (u = file["velocities/u/data"],
                       v = file["velocities/v/data"],
                       w = file["velocities/w/data"])

    tracer_fields =
        NamedTuple{tracers}(Tuple(file["tracers/$c/data"] for c in tracers))

    current_tendency_velocity_fields = (u = file["timestepper/Gⁿ/u/data"],
                                        v = file["timestepper/Gⁿ/v/data"],
                                        w = file["timestepper/Gⁿ/w/data"])

    current_tendency_tracer_fields =
        NamedTuple{tracers}(Tuple(file["timestepper/Gⁿ/$c/data"] for c in tracers))

    previous_tendency_velocity_fields = (u = file["timestepper/G⁻/u/data"],
                                         v = file["timestepper/G⁻/v/data"],
                                         w = file["timestepper/G⁻/w/data"])

    previous_tendency_tracer_fields =
        NamedTuple{tracers}(Tuple(file["timestepper/G⁻/$c/data"] for c in tracers))

    close(file)

    solution = merge(velocity_fields, tracer_fields)
    Gⁿ = merge(current_tendency_velocity_fields, current_tendency_tracer_fields)
    G⁻ = merge(previous_tendency_velocity_fields, previous_tendency_tracer_fields)

    return solution, Gⁿ, G⁻
end

include("regression_tests/thermal_bubble_regression_test.jl")
include("regression_tests/rayleigh_benard_regression_test.jl")
include("regression_tests/ocean_large_eddy_simulation_regression_test.jl")
include("regression_tests/hydrostatic_free_turbulence_regression_test.jl")

@testset "Regression" begin
    @info "Running regression tests..."

    for arch in archs
        for grid_type in [:regular, :vertically_unstretched]
            @testset "Thermal bubble [$(typeof(arch)), $grid_type grid]" begin
                @info "  Testing thermal bubble regression [$(typeof(arch)), $grid_type grid]"
                run_thermal_bubble_regression_test(arch, grid_type)
            end

            @testset "Rayleigh–Bénard tracer [$(typeof(arch)), $grid_type grid]]" begin
                @info "  Testing Rayleigh–Bénard tracer regression [$(typeof(arch)), $grid_type grid]"
                run_rayleigh_benard_regression_test(arch, grid_type)
            end

            for closure in (AnisotropicMinimumDissipation(ν=1.05e-6, κ=1.46e-7), SmagorinskyLilly(C=0.23, Cb=1, Pr=1, ν=1.05e-6, κ=1.46e-7))
                closurename = string(typeof(closure).name.wrapper)
                @testset "Ocean large eddy simulation [$(typeof(arch)), $closurename, $grid_type grid]" begin
                    @info "  Testing oceanic large eddy simulation regression [$(typeof(arch)), $closurename, $grid_type grid]"
                    run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure)
                end
            end
        end

        explicit_free_surface = ExplicitFreeSurface(gravitational_acceleration=1.0)
        implicit_free_surface = ImplicitFreeSurface(gravitational_acceleration = 1.0,
                                                    solver_method = :PreconditionedConjugateGradient,
                                                    tolerance = 1e-15)

        x_bounded_lat_lon_grid  = RegularLatitudeLongitudeGrid(size = (160, 60, 3),
                                                               longitude = (-160, 160),
                                                               latitude = (-60, 60),
                                                               z = (-90, 0),
                                                               halo = (2, 2, 2))

        x_periodic_lat_lon_grid  = RegularLatitudeLongitudeGrid(size = (180, 60, 3),
                                                                longitude = (-180, 180),
                                                                latitude = (-60, 60),
                                                                z = (-90, 0),
                                                                halo = (2, 2, 2))

        for grid in [x_bounded_lat_lon_grid, x_periodic_lat_lon_grid]
            for free_surface in [explicit_free_surface, implicit_free_surface]

                # GPU + ExplicitFreeSurface is broken. See:
                # https://github.com/CliMA/Oceananigans.jl/pull/1985
                # if !(arch isa GPU && topology(grid, 1) === Periodic && free_surface isa ExplicitFreeSurface)

                free_surface_str = string(typeof(free_surface).name.wrapper)

                @testset "Hydrostatic free turbulence regression [$(typeof(arch)), $(topology(grid, 1)) longitude, $free_surface_str]" begin
                    @info "  Testing Hydrostatic free turbulence [$(typeof(arch)), $(topology(grid, 1)) longitude, $free_surface_str]"
                    run_hydrostatic_free_turbulence_regression_test(grid, free_surface, arch)
                end
            end
	    end
	end
end