# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).

"""
    VectorOfArrays{T,N,M} <: AbstractVector{<:AbstractArray{T,N}}

An `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may
differ in size). Internally, `VectorOfArrays` stores all elements of all
arrays in a single flat vector. `M` must equal `N - 1`

The `VectorOfArrays` itself supports `push!`, `unshift!`, etc., but the size
of each individual array in the vector is fixed. `resize!` can be used to
shrink, but not to grow, as the size of the additional element arrays in the
vector would be unknown. However, memory space for up to `n` arrays with a
maximum size `s` can be reserved via
`sizehint!(A::VectorOfArrays, n, s::Dims{N})`.

Constructors:

```julia
VectorOfArrays{T,N}()

VectorOfArrays(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T}(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T,N}(A::AbstractVector{<:AbstractArray})

VectorOfArrays(
    data::AbstractVector,
    elem_ptr::AbstractVector{Int},
    kernel_size::AbstractVector{<:Dims}
    checks::Function = ArraysOfArrays.full_consistency_checks
)
```

Other suitable values for `checks` are `ArraysOfArrays.simple_consistency_checks`
and `ArraysOfArrays.no_consistency_checks`.

`VectorOfVectors` is defined as an type alias:

```julia
`VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}`
```
"""
struct VectorOfArrays{
    T, N, M,
    VT<:AbstractVector{T},
    VI<:AbstractVector{Int},
    VD<:AbstractVector{Dims{M}}
} <: AbstractVector{Array{T,N}}
    data::VT
    elem_ptr::VI
    kernel_size::VD

    function VectorOfArrays{T,N}() where {T,N}
        M = length(Base.front(ntuple(_ -> 0, Val{N}())))
        data = Vector{T}()
        elem_ptr = [firstindex(data)]
        kernel_size = Vector{Dims{M}}()

        new{
            T, N, M,
            typeof(data),
            typeof(elem_ptr),
            typeof(kernel_size)
        }(data, elem_ptr, kernel_size)
    end

    function VectorOfArrays(
        data::VT,
        elem_ptr::VI,
        kernel_size::VD,
        checks::Function = full_consistency_checks
    ) where {
        T, M,
        VT<:AbstractVector{T},
        VI<:AbstractVector{Int},
        VD<:AbstractVector{Dims{M}}
    }
        N = length((ntuple(_ -> 0, Val{M}())..., 0))
        A = new{T,N,M,VT,VI,VD}(data, elem_ptr, kernel_size)
        checks(A)
        A
    end
end

export VectorOfArrays

function VectorOfArrays{T,N}(A::AbstractVector{<:AbstractArray{U,N}}) where {T,N,U}
    R = VectorOfArrays{T,N}()
    append!(R, A)
end

VectorOfArrays{T}(A::AbstractVector{<:AbstractArray{U,N}}) where {T,U,N} = VectorOfArrays{T,N}(A)

VectorOfArrays(A::AbstractVector{<:AbstractArray{T,N}}) where {T,N} = VectorOfArrays{T,N}(A)

Base.convert(VA::Type{VectorOfArrays{T,N}}, A::AbstractVector{AbstractArray{U,N}}) where {T,N,U} = VA(A)
Base.convert(VA::Type{VectorOfArrays}, A::AbstractVector{AbstractArray{T,N}}) where {T,N} = VA(A)

"""
    internal_element_ptr(A::VectorOfArrays)

Returns the internal element pointer vector of `A`.

Do *not* change modify the returned vector in any way, as this would break the
inner consistency of `A`.

Use with care, see [`element_ptr`](@ref) for a safe version of this function.
"""
internal_element_ptr(A::VectorOfArrays) = A.elem_ptr

"""
    element_ptr(A::VectorOfArrays)

Returns a copy of the internal element pointer vector of `A`.
"""
element_ptr(A::VectorOfArrays) = deepcopy(internal_element_ptr(A))

function full_consistency_checks(A::VectorOfArrays)
    simple_consistency_checks(A)
    all(eachindex(A.kernel_size)) do i
        len = A.elem_ptr[i+1] - A.elem_ptr[i]
        klen = prod(A.kernel_size[i])
        len >= 0 && (klen == 1 || mod(len, klen) == 0)
    end || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent"))
    nothing
end

function simple_consistency_checks(A::VectorOfArrays{T,N,M}) where {T,N,M}
    M == N - 1 || throw(ArgumentError("VectorOfArrays{T,N,M} inconsistent: M must equal N - 1"))
    firstindex(A.elem_ptr) == firstindex(A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible indexing"))
    size(A.elem_ptr, 1) == size(A.kernel_size, 1) + 1 || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible size"))
    first(A.elem_ptr) >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices"))
    last(A.elem_ptr) - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices"))
    nothing
end

function no_consistency_checks(A::VectorOfArrays)
    nothing
end

Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer)
    elem_ptr = A.elem_ptr

    from = elem_ptr[i]
    until = elem_ptr[i+1]
    to = until - 1
    len = until - from

    ksize = A.kernel_size[i]
    klen = prod(ksize)
    len_p, klen_p = promote(len, klen)
    sz_lastdim = len == 0 ? len_p : div(len_p, klen_p)
    sz = (ksize..., sz_lastdim)

    (from:to, sz)
end

import Base.==
(==)(A::VectorOfArrays, B::VectorOfArrays) =
    A.data == B.data && A.elem_ptr == B.elem_ptr && A.kernel_size == B.kernel_size

"""
    flatview(A::VectorOfArrays{T})::Vector{T}

Returns the internal serialized representation of all element arrays of `A` as
a single vector. Do *not* change the length of the returned vector, as it
would break the inner consistency of `A`.
"""
flatview(A::VectorOfArrays) = A.data

Base.size(A::VectorOfArrays) = size(A.kernel_size)

Base.IndexStyle(A::VectorOfArrays) = IndexLinear()

Base.@propagate_inbounds _reshape_dataview(dataview::AbstractArray, s::NTuple{1,Integer}) = dataview

Base.@propagate_inbounds _reshape_dataview(dataview::AbstractArray, s::NTuple{N,Integer}) where {N} =
    Base.__reshape((dataview, IndexStyle(dataview)), s)

Base.@propagate_inbounds function Base.getindex(A::VectorOfArrays, i::Integer)
    @boundscheck checkbounds(A, i)
    r, s = _elem_range_size(A, i)
    dataview = view(A.data, r)
    _reshape_dataview(dataview, s)
end

Base.@propagate_inbounds function Base._getindex(l::IndexStyle, A::VectorOfArrays, idxs::AbstractUnitRange{<:Integer})
    from = first(idxs)
    to = last(idxs)
    elem_ptr = A.elem_ptr[from:(to+1)]
    kernel_size = A.kernel_size[from:to]
    data = A.data[first(elem_ptr):(last(elem_ptr) - 1)]
    broadcast!(+, elem_ptr, elem_ptr, firstindex(data) - first(elem_ptr))
    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

Base.@propagate_inbounds function Base._getindex(l::IndexStyle, A::VectorOfArrays, idxs::AbstractVector{<:Integer})
    @boundscheck checkbounds(A, idxs)

    A_ep = A.elem_ptr
    A_data = A.data

    elem_ptr = similar(A_ep, length(eachindex(idxs)) + 1)
    delta_i = firstindex(elem_ptr) - firstindex(idxs)

    elem_ptr[firstindex(elem_ptr)] = firstindex(A_data)
    for i in eachindex(idxs)
        idx = idxs[i]
        l = A_ep[idx + 1] - A_ep[idx]
        elem_ptr[i + 1 + delta_i] = elem_ptr[i + delta_i] + l
    end

    data = similar(A_data, last(elem_ptr) - first(elem_ptr))
    if firstindex(data) != firstindex(A_data)
        @assert firstindex(data) != first(elem_ptr)
        broadcast!(+, elem_ptr, elem_ptr, firstindex(data) - first(elem_ptr))
    end

    for i in eachindex(idxs)
        idx = idxs[i]
        l = A_ep[idx + 1] - A_ep[idx]

        # Sanity check:
        @assert l == elem_ptr[i + 1 + delta_i] - elem_ptr[i + delta_i]

        copyto!(data, elem_ptr[i + delta_i], A_data, A_ep[idx], l)
    end

    kernel_size = A.kernel_size[idxs]

    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

Base.@propagate_inbounds function Base.setindex!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}, i::Integer) where {T,N,U}
    r, s = _view_reshape_spec(A, i)
    @boundscheck s == size(x) || throw(DimensionMismatch("Can't assign array to element $i of VectorOfArrays, array size is incompatible"))
    A.data[rng] = x
    A
end

Base.length(A::VectorOfArrays) = length(A.kernel_size)

@inline function Base.resize!(A::VectorOfArrays{T,N,M}, n::Integer) where {T,M,N}
    old_n = length(A)
    if n > old_n
        throw(ArgumentError("Cannot resize VectorOfArrays from length $old_n to $n, can only shrink, not grow"))
    elseif n < old_n
        resize!(A.data, A.elem_ptr[n+1] - 1)
        resize!(A.elem_ptr, n + 1)
        resize!(A.kernel_size, n)
    end
    A
end

function append_elemptr!(A::AbstractVector{<:Integer}, B::AbstractVector{<:Integer})
    idxs_A = LinearIndices(A)
    idxs_B = LinearIndices(B)
    length_B = length(idxs_B)

    A_from = last(idxs_A) + 1
    A_to = A_from - 1 + length_B - 1
    resize!(A, length(first(idxs_A):A_to))

    B_from = first(idxs_B) + 1
    B_to = last(idxs_B)

    A_idxs_offs = A_from - B_from

    checkindex(Bool, eachindex(B), B_from:B_to)
    checkindex(Bool, eachindex(A), (B_from:B_to) .+ A_idxs_offs)
    @inbounds begin
        value_offset = A[A_from - 1] - B[B_from - 1]
        @simd for i in B_from:B_to
            A[i + A_idxs_offs] = B[i] + value_offset
        end
    end

    A
end

function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U}
    if !isempty(B)
        append!(A.data, B.data)
        append_elemptr!(A.elem_ptr, B.elem_ptr)
        append!(A.kernel_size, B.kernel_size)
    end
    A
end

function Base.append!(A::VectorOfArrays{T,N}, B::AbstractVector{<:AbstractArray{U,N}}) where {T,N,U}
    if !isempty(B)
        n_A = length(eachindex(A))
        n_B = length(eachindex(B))
        datalen_A = length(eachindex(A.data))
        datalen_B = zero(Int)
        for i in eachindex(B)
            datalen_B += Int(length(eachindex(B[i])))
        end

        sizehint!(A.data, datalen_A + datalen_B)
        sizehint!(A.elem_ptr, n_A + n_B + 1)
        sizehint!(A.kernel_size, n_A + n_B)

        for i in eachindex(B)
            push!(A, B[i])
        end
    end
    A
end

Base.vcat(V::VectorOfArrays) = V

function Base.vcat(Vs::(VectorOfArrays{U,N} where U)...) where {N}
    data = vcat(map(x -> x.data, Vs)...)

    elem_ptr = similar(Vs[1].elem_ptr, 1)
    elem_ptr[1] = firstindex(data)
    @inbounds for i in eachindex(Vs)
        append_elemptr!(elem_ptr, Vs[i].elem_ptr)
    end

    kernel_size = vcat(map(x -> x.kernel_size, Vs)...)

    VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks)
end

function Base.copy(V::VectorOfArrays)
    VectorOfArrays(copy(V.data), copy(V.elem_ptr), copy(V.kernel_size), no_consistency_checks)
end

Base.@propagate_inbounds function Base.unsafe_view(A::VectorOfArrays, idxs::AbstractUnitRange{<:Integer})
    from = first(idxs)
    to = last(idxs)
    VectorOfArrays(
        A.data,
        view(A.elem_ptr, from:(to+1)),
        view(A.kernel_size, from:to),
        no_consistency_checks
    )
end

function Base.sizehint!(A::VectorOfArrays{T,N}, n, s::Dims{N}) where {T,N}
    sizehint!(A.data, n * mul(s))
    sizehint!(A.elem_ptr, n + 1)
    sizehint!(A.kernel_size, n)
    A
end

function Base.push!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}) where {T,N,U}
    @assert last(A.elem_ptr) == lastindex(A.data) + 1
    append!(A.data, x)
    push!(A.elem_ptr, lastindex(A.data) + 1)
    push!(A.kernel_size, Base.front(size(x)))
    A
end

function UnsafeArrays.uview(A::VectorOfArrays)
    VectorOfArrays(
        uview(A.data),
        uview(A.elem_ptr),
        uview(A.kernel_size),
        no_consistency_checks
    )
end

function innermap(f::Base.Callable, A::VectorOfArrays)
    new_data = map(f, A.data)
    VectorOfArrays(new_data, A.elem_ptr, A.kernel_size, simple_consistency_checks)
end

function deepmap(f::Base.Callable, A::VectorOfArrays)
    new_data = deepmap(f, A.data)
    VectorOfArrays(new_data, A.elem_ptr, A.kernel_size, simple_consistency_checks)
end

"""
    VectorOfVectors{T,...} = VectorOfArrays{T,1,...}

Constructors:

```julia
VectorOfVectors(A::AbstractVector{<:AbstractVector})
VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T}

VectorOfVectors(
    data::AbstractVector, elem_ptr::AbstractVector{Int},
    checks::Function = full_consistency_checks
)

See also [VectorOfArrays](@ref).
```
"""
const VectorOfVectors{
    T,
    VT<:AbstractVector{T},
    VI<:AbstractVector{Int},
    VD<:AbstractVector{Dims{0}}
} = VectorOfArrays{T,1,0,VT,VI,VD}

export VectorOfVectors

VectorOfVectors{T}() where {T} = VectorOfArrays{T,1}()

VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A)
VectorOfVectors(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A)

VectorOfVectors(
    data::AbstractVector,
    elem_ptr::AbstractVector{Int},
    checks::Function = full_consistency_checks
) = VectorOfArrays(
    data,
    elem_ptr,
    similar(elem_ptr, Dims{0}, size(elem_ptr, 1) - 1),
    checks
)

"""
    consgrouped_ptrs(A::AbstractVector)

Compute an element pointer vector, suitable for creation of a
`VectorOfVectors` that implies grouping equal consecutive entries of
`A`.

Example:

```julia
    A = [1, 1, 2, 3, 3, 2, 2, 2]
    elem_ptr = consgrouped_ptrs(A)
    first.(VectorOfVectors(A, elem_ptr)) == [1, 2, 3, 2]
```
consgrouped_ptrs
Typically, `elem_ptr` will be used to apply the computed grouping to other
data:

```julia
    B = [1, 2, 3, 4, 5, 6, 7, 8]
    VectorOfVectors(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]]
```
"""
function consgrouped_ptrs end
export consgrouped_ptrs

function consgrouped_ptrs(A::AbstractVector)
    elem_ptr = Vector{Int}()
    idxs = eachindex(A)
    push!(elem_ptr, first(idxs))
    if !isempty(A)
        prev_0 = A[first(idxs)]
        prev::typeof(prev_0) = prev_0
        @inbounds for i in (first(idxs) + 1):last(idxs)
            curr = A[i]
            if (curr != prev)
                push!(elem_ptr, i)
                prev = curr
            end
        end
        push!(elem_ptr, last(idxs) + 1)
    end
    elem_ptr
end

"""
    consgroupedview(source::AbstractVector, target)

Compute a grouping of equal consecutive elements on `source` via
[`consgrouped_ptrs`](@ref) and apply the grouping to target, resp. each
element of `target`. `target` may be an vector or a named or unnamed tuple of
vectors. The result is a `VectorOfVectors`, resp. a tuple of such.

Example:

    A = [1, 1, 2, 3, 3, 2, 2, 2]
    B = [1, 2, 3, 4, 5, 6, 7, 8]
    consgroupedview(A, B) == [[1, 2], [3], [4, 5], [6, 7, 8]]

`consgroupedview` plays well with columnar tables, too:

```julia
    using Tables, TypedTables
    data = Table(
        a = [1, 1, 2, 3, 3, 2, 2, 2],
        b = [1, 2, 3, 4, 5, 6, 7, 8],
        c = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
    )

    result = Table(consgroupedview(data.a, Tables.columns(data)))
```

will return

```
     a          b          c
   ┌──────────────────────────────────────
 1 │ [1, 1]     [1, 2]     [1.1, 2.2]
 2 │ [2]        [3]        [3.3]
 3 │ [3, 3]     [4, 5]     [4.4, 5.5]
 4 │ [2, 2, 2]  [6, 7, 8]  [6.6, 7.7, 8.8]
```

without copying any data:

```
    flatview(result.a) === data.a
    flatview(result.b) === data.b
    flatview(result.c) === data.c
```
"""
function consgroupedview end
export consgroupedview

function consgroupedview(source::AbstractVector, target::AbstractVector)
    elem_ptr = consgrouped_ptrs(source)
    VectorOfVectors(target, elem_ptr)
end

function consgroupedview(source::AbstractVector, target::NTuple{N,AbstractVector}) where {N}
    elem_ptr = consgrouped_ptrs(source)
    map(X -> VectorOfVectors(X, elem_ptr), target)
end

function consgroupedview(source::AbstractVector, target::NamedTuple{syms,<:NTuple{N,AbstractVector}}) where {syms,N}
    elem_ptr = consgrouped_ptrs(source)
    map(X -> VectorOfVectors(X, elem_ptr), target)
end


function neighpairs(x, q, pairlist, pairnum, atomtype, ilist, alist)

dim = size(x,1);
ncq = size(q,1);
inum = length(ilist);
ijnum = sum(pairnum[ilist.+1]-pairnum[ilist]);

ai = zeros(eltype(ilist),ijnum);
aj = zeros(eltype(ilist),ijnum);
ti = zeros(eltype(ilist),ijnum);
tj = zeros(eltype(ilist),ijnum);
xij = zeros(dim,ijnum);
qi = zeros(ncq,ijnum);
qj = zeros(ncq,ijnum);

count = 0;
for ii=1:inum
    i = ilist[ii]; # atom i
    n1 = pairnum[i];
    n2 = pairnum[i+1];
    n = n2-n1;
    g = pairlist[(n1+1):n2]; # a list of neighbors around atom i
    j = alist[g];
    ind = (count+1):(count+n);
    count = count + n;

    ai[ind] .= i;
    aj[ind] .= j;
    ti[ind] .= atomtype[i];
    tj[ind] .= atomtype[j];
    xij[:,ind] = x[:,g] .- x[:,i];

    if ncq>0
        qi[ind] .= q[i];
        qj[ind] .= q[j];
    end
end

return xij, ai, aj, ti, tj, qi, qj

end

function neighpairs(x, pairlist, pairnum, atomtype, ilist, alist)

    dim = size(x,1);
    inum = length(ilist);
    ijnum = sum(pairnum[ilist.+1]-pairnum[ilist]);

    ai = zeros(eltype(ilist),ijnum);
    aj = zeros(eltype(ilist),ijnum);
    ti = zeros(eltype(ilist),ijnum);
    tj = zeros(eltype(ilist),ijnum);
    xij = zeros(dim,ijnum);

    count = 0;
    for ii=1:inum
        i = ilist[ii]; # atom i
        n1 = pairnum[i];
        n2 = pairnum[i+1];
        n = n2-n1;
        g = pairlist[(n1+1):n2]; # a list of neighbors around atom i
        j = alist[g];
        ind = (count+1):(count+n);
        count = count + n;

        ai[ind] .= i;
        aj[ind] .= j;
        ti[ind] .= atomtype[i];
        tj[ind] .= atomtype[j];
        xij[:,ind] = x[:,g] .- x[:,i];

    end

    return xij, ai, aj, ti, tj

end


function interpretwasm(f::Func, fs, args)
  params = [convert(jltype(typ), arg) for (typ, arg) in zip(f.params, args)]
  locals = [convert(jltype(typ), 0) for typ in f.locals]
  ms = vcat(params, locals)

  apInstr(f.body, ms, 0, fs)

  returns = popn!(ms, length(f.result))

  # Check to make sure the return types are all correct
  for (rtyp, r) in zip(f.result, returns)
    if jltype(rtyp) != typeof(r)
      error("ERROR: Return type mismatch")
    end
  end

  return returns
end

function toFunction(f::Func, fs)
  return ((args...) -> interpretwasm(f, fs, args))
end

function interpret_module(m::Module)
  d = Dict()
  [(g = toFunction(f, d); d[f.name] = (length(f.params), g); g) for f in m.funcs]
end

function runBody(body, ms, level, fs)
  for i in body
    if typeof(i) in [Branch, Return]
      level_ = apInstr(i, ms, level + 1)
      level_ <= level && return level_
    elseif typeof(i) in [If, Block, Loop]
      level_ = apInstr(i, ms, level + 1, fs)
      level_ <= level && return level_
    elseif i isa Call
      apInstr(i, ms, level + 1, fs)
    else
      apInstr(i, ms)
    end
  end
  return level
end

function popn!(xs, n)
  len = length(xs)
  return splice!(xs, len - n + 1:len)
end

apN(n, f) = ms -> push!(ms, f(popn!(ms, n)...))

# Functions to fix true -> 1 and false -> 0, but that's the case anyway
# function MComp(comparator)
#   return (args...) -> comparator(args...) ? 1 : 0
# end
#
# function cpN(n, f)
#   return apN(n, MComp(f))
# end

unsign(x::Int64) = reinterpret(UInt64, x)
unsign(x::Int32) = reinterpret(UInt32, x)
resign(x::UInt64) = reinterpret(Int64, x)
resign(x::UInt32) = reinterpret(Int32, x)

unsign(x) = x
resign(x) = x

apN_U(n, f) = ms -> push!(ms, resign(f(unsign(popn!(ms, n))...)))

operations =
  Dict(:lt_s   => apN(2, <)
      ,:lt_u   => apN_U(2, <)
      ,:le_s   => apN(2, <=)
      ,:gt_s   => apN(2, >)
      ,:eq     => apN(2, ==)
      ,:ne     => apN(2, !=)
      ,:gt     => apN(2, >)
      ,:eqz    => apN(1, x -> x == 0)
      ,:sub    => apN(2, -)
      ,:add    => apN(2, +)
      ,:div    => apN(2, /)
      ,:mul    => apN(2, *)
      ,:ctz    => apN(1, trailing_zeros)
      ,:shr_s  => apN(2, >>)
      ,:shr_u  => apN(2, >>>)
      ,:shl    => apN(2, <<)
      ,:or     => apN(2, |)
      ,:and    => apN(2, &)
      ,:rem_s  => apN(2, rem)
      ,:rem_u  => apN_U(2, rem)
      ,:div_s  => apN(2, div)
      ,:div_u  => apN_U(2, div)
      )

# Level Agnostic Functions

apInstr(i::Nop,         ms) = nothing
apInstr(i::Local,       ms) = push!(ms,ms[i.id + 1]);
apInstr(i::Const,       ms) = push!(ms,value(i));
apInstr(i::Unreachable, ms) = error("Unreachable")
apInstr(i::Convert,     ms) = push!(ms, convert(jltype(i.to), float(pop!(ms))))
apInstr(i::Op,          ms) = operations[Symbol(split(string(i.name), ".")[2])](ms)

# apInstr(i::SetLocal,    ms) = i.id + 1 == length(ms) ? ms[i.id + 1] = i.tee ? last(ms) : pop!(ms) : ms
apInstr(i::SetLocal,    ms) = ms[i.id + 1] = i.tee ? last(ms) : pop!(ms)

function apInstr(i::Select, ms)
  arg2, cond = popn!(ms, 2)
  if cond == 0
    pop!(ms)
    push!(ms, arg2)
  end
end

# Level Based Functions

apInstr(i::Return, ms, l) = 0
apInstr(i::Branch, ms, l) = i.cond && pop!(ms) != 0 || !i.cond ? l - i.level - 1 : l

# Instructions dependent on in scope functions

apInstr(i::Call,   ms, l, fs) = (f = fs[i.name]; push!(ms, f[2]((popn!(ms, f[1]))...)...))

apInstr(i::If,     ms, l, fs) = pop!(ms) != 0 ? runBody(i.t, ms, l, fs) : runBody(i.f, ms, l, fs)
apInstr(i::Block,  ms, l, fs) = runBody(i.body, ms, l, fs)
function apInstr(loop::Loop, ms, level, fs)
  level_ = level
  while level_ == level
    level_ = runBody(loop.body, ms, level, fs)
  end
  return level_
end


using DimensionalData, Test, Unitful, Adapt

using DimensionalData.LookupArrays, DimensionalData.Dimensions

struct CustomArray{T,N} <: AbstractArray{T,N}
    arr::Array
end

CustomArray(x::Array{T,N}) where {T,N} = CustomArray{T,N}(x)
Adapt.adapt_storage(::Type{<:CustomArray}, xs::Array) = CustomArray(xs)

Base.size(x::CustomArray, y...) = size(x.arr, y...)
Base.getindex(x::CustomArray, y...) = getindex(x.arr, y...)
Base.count(x::CustomArray) = count(x.arr)

@testset "Metadata" begin
    @test adapt(CustomArray, Metadata(:a=>"1", :b=>"2")) == NoMetadata()
end

@testset "LookupArray" begin
    l = Sampled([1:10...]; metadata=Metadata(:a=>"1", :b=>"2"))
    l1 = Adapt.adapt(CustomArray, l)
    @test parent(parent(l1)) isa CustomArray
    @test parent(parent(l1)).arr == [1:10...]
    @test metadata(l1) == NoMetadata()
    l = Categorical('a':1:'n'; metadata=Metadata(:a=>"1", :b=>"2"))
    l1 = Adapt.adapt(CustomArray, l)
    @test parent(parent(l1)) isa StepRange
    @test metadata(l1) == NoMetadata()
    l = Categorical(['a', 'e', 'n']; metadata=Metadata(:a=>"1", :b=>"2"))
    l1 = Adapt.adapt(CustomArray, l)
    @test parent(parent(l1)) isa CustomArray
    @test metadata(l1) == NoMetadata()
    l = NoLookup(Base.OneTo(10))
    l1 = Adapt.adapt(CustomArray, l)
    @test parent(parent(l1)) isa Base.OneTo
    l = AutoLookup()
    l1 = Adapt.adapt(CustomArray, l)
    @test parent(l1) isa AutoIndex
end

@testset "Dimension" begin
    d = X(Sampled([1:10...]; metadata=Metadata(:a=>"1", :b=>"2")))
    d1 = Adapt.adapt(CustomArray, d)
    @test parent(parent(d1)) isa CustomArray
    @test parent(parent(d1)).arr == [1:10...]
    @test metadata(d1) == NoMetadata()
end

@testset "DimArray" begin
    A = rand(4, 5)
    da = DimArray(A, (X, Y))
    da1 = adapt(CustomArray, da)
    @test parent(da1) isa CustomArray
    @test parent(da1).arr == A
    @test metadata(da1) == NoMetadata()
end

@testset "DimStack" begin
    A = rand(4, 5)
    B = rand(4, 5)
    ds = DimStack((a=A, b=B), (X, Y))
    ds1 = adapt(CustomArray, ds)
    @test parent(ds1[:a]) isa CustomArray
    @test parent(ds1[:a]).arr == A
    @test parent(ds1[:b]) isa CustomArray
    @test parent(ds1[:b]).arr == B
    @test metadata(ds) == NoMetadata()
end


module VectorizedRNG

using VectorizationBase, Random, UnPack
using VectorizationBase: simd_integer_register_size, gep, _Vec, ifelse, VecUnroll, AbstractSIMD,
    rotate_right, vadd, vsub, zero_offsets, vfmadd, vfmsub, vfnmadd, shufflevector, register_size,
    cache_linesize, StaticInt, pick_vector_width, data

using Distributed: myid

export local_rng, rand!, randn!#, randexp, randexp!

abstract type AbstractVRNG{N} <: Random.AbstractRNG end
abstract type AbstractState{N,W} end

@inline vloadu(ptr::Ptr) = VectorizationBase.__vload(ptr, VectorizationBase.False(), register_size())
@inline vloadu(ptr::Ptr, i) = VectorizationBase.__vload(ptr, i, VectorizationBase.False(), register_size())
@inline vloada(ptr::Ptr) = VectorizationBase.__vload(ptr, VectorizationBase.True(), register_size())
@inline vloada(ptr::Ptr, i) = VectorizationBase.__vload(ptr, i, VectorizationBase.True(), register_size())
@inline vloada(ptr::Ptr, i, m) = VectorizationBase.__vload(ptr, i, m, VectorizationBase.True(), register_size())
@inline vstorea!(ptr::Ptr, v) = VectorizationBase.__vstore!(ptr, v, VectorizationBase.True(), VectorizationBase.False(), VectorizationBase.False(), register_size())
@inline vstorea!(ptr::Ptr, v, i) = VectorizationBase.__vstore!(ptr, v, i, VectorizationBase.True(), VectorizationBase.False(), VectorizationBase.False(), register_size())
@inline vstorea!(ptr::Ptr, v, i, m) = VectorizationBase.__vstore!(ptr, v, i, m, VectorizationBase.True(), VectorizationBase.False(), VectorizationBase.False(), register_size())
@inline vstoreu!(ptr::Ptr, v, i) = VectorizationBase.__vstore!(ptr, v, i, VectorizationBase.False(), VectorizationBase.False(), VectorizationBase.False(), register_size())
@inline vstoreu!(ptr::Ptr, v) = VectorizationBase.__vstore!(ptr, v, VectorizationBase.False(), VectorizationBase.False(), VectorizationBase.False(), register_size())

include("masks.jl")
include("api.jl")
include("special_approximations.jl")
include("xoshiro.jl")
# const GLOBAL_vPCGs = Ref{Ptr{UInt64}}()

const GLOBAL_vRNGs = Ref{Ptr{UInt64}}()

const RNG_MEM_SIZE = (5(simd_integer_register_size()*XREGISTERS + 2048*3))
local_rng(i) = Xoshift{XREGISTERS}(i*(RNG_MEM_SIZE) + GLOBAL_vRNGs[])
local_rng() = local_rng(Base.Threads.threadid() - 1)

# include("precompile.jl")
# _precompile_()

function __init__()
    # ccall(:jl_generating_output, Cint, ()) == 1 && return
  nthreads = Base.Threads.nthreads()
  GLOBAL_vRNGs[] = ptr = VectorizationBase.valloc((RNG_MEM_SIZE ÷ 8)*nthreads, UInt64)
  nstreams = XREGISTERS * nthreads * simd_integer_register_size()
  initXoshift!(ptr, nstreams)
  for tid ∈ 0:nthreads-1
    rng = local_rng(tid)
    setrandu64counter!(rng, 0x00)
    # setrandn32counter!(rng, 0x00)
    setrand64counter!(rng, 0x00)
    setrandn64counter!(rng, 0x00)
  end
end

end # module


###############################################################################
#
# Plant Hydraulic system
#
###############################################################################
"""
    abstract type AbstractPlantOrganism{FT}

Hierachy of AbstractPlantOrganism
- [`GrassLikeOrganism`](@ref)
- [`PalmLikeOrganism`](@ref)
- [`TreeLikeOrganism`](@ref)
"""
abstract type AbstractPlantOrganism{FT<:AbstractFloat} end

"""
    mutable struct GrassLikeOrganism{FT}

A plant hydraulic system like a grass, which contains multiple root layers, and
multiple canopy layers. No trunk or branch system applies.

# Fields
$(TYPEDFIELDS)
"""
mutable struct GrassLikeOrganism{FT} <: AbstractPlantOrganism{FT}
    # structre information
    "Root Layers"
    n_root  ::Int
    "Canopy Layers"
    n_canopy::Int

    # Arrays of roots and leaves
    "Roots system"
    roots ::Array{RootHydraulics{FT},1}
    "Leaves"
    leaves::Array{LeafHydraulics{FT},1}

    # Root and canopy index in Soil and Atmosphere
    "Corresponding soil layer per root layer"
    root_index_in_soil ::Array{Int,1}
    "Corresponding air layer per canopy layer"
    canopy_index_in_air::Array{Int,1}

    # containers for root flows
    "Conductances for each layer at given flow"
    cache_k::Array{FT,1}
    "Pressure for each layer at given flow"
    cache_p::Array{FT,1}
    "Flow rate"
    cache_q::Array{FT,1}
end

"""
    mutable struct PalmLikeOrganism{FT}

A plant hydraulic system like a palm, which contains multiple root layers, one
trunk, and multiple canopy layers. No branch system applies.

# Fields
$(TYPEDFIELDS)
"""
mutable struct PalmLikeOrganism{FT} <: AbstractPlantOrganism{FT}
    # structre information
    "Root Layers"
    n_root  ::Int
    "Canopy Layers"
    n_canopy::Int

    # Arrays of roots and leaves
    "Roots system"
    roots ::Array{RootHydraulics{FT},1}
    "Trunk"
    trunk ::StemHydraulics{FT}
    "Leaves"
    leaves::Array{LeafHydraulics{FT},1}

    # Root and canopy index in Soil and Atmosphere
    "Corresponding soil layer per root layer"
    root_index_in_soil ::Array{Int,1}
    "Corresponding air layer per canopy layer"
    canopy_index_in_air::Array{Int,1}

    # containers for root flows
    "Conductances for each layer at given flow"
    cache_k::Array{FT,1}
    "Pressure for each layer at given flow"
    cache_p::Array{FT,1}
    "Flow rate"
    cache_q::Array{FT,1}
end

"""
    mutable struct TreeLikeOrganism{FT}

A plant hydraulic system like a tree, which contains multiple root layers, one
trunk, and multiple branch and canopy layers.

# Fields
$(TYPEDFIELDS)
"""
mutable struct TreeLikeOrganism{FT} <: AbstractPlantOrganism{FT}
    # structre information
    "Root Layers"
    n_root  ::Int
    "Canopy Layers"
    n_canopy::Int

    # Arrays of roots and leaves
    "Roots system"
    roots ::Array{RootHydraulics{FT},1}
    "Trunk"
    trunk ::StemHydraulics{FT}
    "Branch system"
    branch::Array{StemHydraulics{FT},1}
    "Leaves"
    leaves::Array{LeafHydraulics{FT},1}

    # Root and canopy index in Soil and Atmosphere
    "Corresponding soil layer per root layer"
    root_index_in_soil ::Array{Int,1}
    "Corresponding air layer per canopy layer"
    canopy_index_in_air::Array{Int,1}

    # containers for root flows
    "Conductances for each layer at given flow"
    cache_k::Array{FT,1}
    "Pressure for each layer at given flow"
    cache_p::Array{FT,1}
    "Flow rate"
    cache_q::Array{FT,1}
end

"""
    mutable struct TreeSimple{FT}

A plant hydraulic system with one root, one stem, and one leaf for testing
    purpose

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct TreeSimple{FT} <: AbstractPlantOrganism{FT}
    # Arrays of roots and leaves
    "Root"
    root::RootHydraulics{FT} = RootHydraulics{FT}()
    "Stem"
    stem::StemHydraulics{FT} = StemHydraulics{FT}()
    "Leaf"
    leaf::LeafHydraulics{FT} = LeafHydraulics{FT}()

    # Local container for tree information
    "Relative hydraulic conductance"
    krs::Array{FT,1} = ones(FT,4)
end


__precompile__()

module EMP

using JuMP
using ReSHOP
using MathOptInterface
MOI = MathOptInterface

using Compat
using Base.Iterators

export @vipair, @NLvipair, EMPmaster, solveEMP, _solveEMP, MathPrgm, EquilibriumProblem, BilevelProblem, addovf!, getsolution, status, termination_status, primal_status, dual_status, get_model_result_num, solve, getobjval

include("JuMPExtension.jl")

" Optimal Value Function (OVF) representation "
mutable struct OVF
    vidx::Int
    args::Vector{Int}
    name::String
    params::Dict{String,Any}
    # No idea how to enforce params to be a Dict{String,Any}
    function OVF(v::JuMP.VariableRef, args::Vector{JuMP.VariableRef}, name::String, params::Dict)
        args_idx = [arg.index.value for arg in args]
        new(v.index.value, args_idx, name, params)
    end
end

" EMP master object "
mutable struct EMPmaster
    backend
    mps::Vector{MathPrgm}
    equils::Vector{Vector{MathPrgm}}
    root::Union{Nothing,MathPrgm,Vector{MathPrgm},Vector{Vector{MathPrgm}}}
    ovfs::Vector{OVF}
    NLoffset::Int
end

"""
    Model([modeling_pkg = "JuMP", solver = "jams"])

Create an EMP master object, using the modeling package given as argument to store the variables and equations.
The solver argument is used to construct the ReSHOPSolver object.
"""
function EMPmaster(; modeling_pkg = "JuMP", solver = "jams")
    if modeling_pkg == "JuMP"
        backend = direct_model(ReSHOP.Optimizer(;solver=solver))
        emp = EMPmaster(backend)
    elseif modeling_pkg == "Convex"
        error("Convex.jl is WIP")
    else
        error("No valid modeling_ds value passed. It should be either ``JuMP'' or ``Convex''")
    end

    return emp
end

"""
    Model(backend)

Create an EMP master object and use the argument as the model data storage object
"""
function EMPmaster(backend)
    if !(backend.moi_backend isa ReSHOP.Optimizer)
        error("The JuMP model must be direct solver, use JuMP.direct_model")
    end
    emp = EMPmaster(backend, Vector{MathPrgm}(), Vector{Vector{MathPrgm}}(), nothing, Vector{OVF}(), -1)

    return emp
end

"""
    MathPrgm(m::EMPmaster)

Create a Mathematical Programm in the EMP master object
"""

" Create an Equilibrium object (not implemented) "
function Equilibrium(emp::EMPmaster, mps::Vector{MathPrgm})
    push!(emp.equils, mps)
end

"""
    EquilibriumProblem(emp::EMPmaster, mps::Vector{MathPrgm})

Define the EMP problem as an equilibrium problem
"""
function EquilibriumProblem(emp::EMPmaster, mps::Vector{MathPrgm})
    emp.root = mps
end

function BilevelProblem(emp::EMPmaster, upper::MathPrgm, lower::MathPrgm)
    emp.root = upper
    push!(upper.mps, lower)
end

include("helpers.jl")

#"""
#    addvar!(mp, var)
#
#Add a variable to a Mathematical Programm
#"""
#function addvar!(mp::MathPrgm, var::JuMP.VariableRef)
#    push!(mp.vars, var.index.value)
#end
#
#function addvar!(mp::MathPrgm, var::Array{JuMP.VariableRef})
#    map(x -> addvar!(mp, x), flatten(var))
#end
#
#function addvar!(mp::MathPrgm, var::JuMP.JuMPArray)
#    map(x-> addvar!(mp, x), values(var))
#end
#
#function addvar!(mp::MathPrgm, var::JuMP.JuMPDict)
#    map(x-> addvar!(mp, x), values(var))
#end

#function addvar!(mp::MathPrgm, var::AbstractArray)
#    map(x -> addvar!(mp, x), flatten(var))
#end

#"""
#    addequ!(mp, eqn)
#
#Add an equation to a Mathematical Programm. The argument eqn is a single or list of JuMP.ConstraintRef, not an expression
#"""
#function addequ!(mp::MathPrgm, eqn::JuMP.ConstraintRef, isNL::Bool)
#    push!(mp.equs, (eqn.index.value, isNL))
#end

#function addequ!(mp::MathPrgm, eqn::Vector{ConstraintRef{M,C}}) where {M <: JuMP.AbstractModel, C <: JuMP.AbstractConstraint}
#    map(x -> addequ!(mp, x), eqn)
#end

#function addequ!(mp::MathPrgm, eqn::Vector{JuMP.ConstraintRef})
#    map(x -> addequ!(mp, x), eqn)
#end

#function addequ!(mp::MathPrgm, eqn::AbstractArray, isNL::Bool)
#    map(x -> addequ!(mp, x, isNL), flatten(eqn))
#
#end
#
#function addequ!(mp::MathPrgm, eqn::Vector, isNL::Bool)
#    map(x-> addequ!(mp, x, isNL), values(eqn))
#end

"""
     addovf!(emp, v, args, name, params)

Tag a variable as a supported OVF

# Arguments
- `emp::EMPmaster`: the EMP object
- `v::JuMP.VariableRef`: the variable
- `args::Vector{JuMP.VariableRef}`: the arguments for the OVF problem
- `name::String`: the name of the OVF
- `params::Dict`: the parameters for defining this OVF
"""
function addovf!(emp::EMPmaster, v::JuMP.VariableRef, args::Vector{JuMP.VariableRef}, name::String, params::Dict)
    push!(emp.ovfs, OVF(v, args, name, params))
end

"""
    solveEMP(emp)

Solve an EMP model
"""
function solveEMP(emp::EMPmaster)
    m = getReSHOPModel(emp)

    if (m.mdl == C_NULL)
        m.mdl = ReSHOP.reshop_alloc(m.ctx)
    end

    ReSHOP.emp_init(m.mdl)

    # This is needed to get the metadata initialized
    if (length(emp.mps) > 0)
        ReSHOP.reshop_set_modeltype(m.ctx, ReSHOP.emp)
    end

    # This is needed to get the NL part loaded, before emptree is called
    if emp.backend.nlp_data !== nothing
        MOI.set(emp.backend, MOI.NLPBlock(), JuMP._create_nlp_block_data(emp.backend))
        empty!(emp.backend.nlp_data.nlconstr_duals)
    end

    emp.NLoffset = m.start_nl_cons

    # Deal with the emptree
    emptree(emp, m.mdl)

    # Deal with the OVF functions
    for ovf in emp.ovfs
        ReSHOP.reshop_ovf(m.mdl, ovf)
    end

    optimize!(emp.backend)
end

#########################################################################################################
# End of public API
#########################################################################################################

include("emptree.jl")

function _solveEMP(emp::EMPmaster)

    # Steps:
    # 1. Define the full models and agents
    # 2. Solve
    # 3. Report values

    reshop_model = getReSHOPModel(emp)

    ReSHOP.make_var_index!(reshop_model)
#    ReSHOP.make_con_index!(reshop_model)

    ########
    #
    # We have to mirror the behavior of ReSHOP, where the nonlinear equation are put before
    # the linear ones.
    #
    # TODO(xhub) fix this mess ...
    #
    #####

    if (length(emp.mps) > 0)
        make_equ_index!(emp.equs, emp.backend.internalModel)
    else
        ReSHOP.make_con_index!(reshop_model)
    end

    ctx = ReSHOP.create_reshop_ctx(reshop_model)
    reshop_model.reshop_ctx = ctx

    reshop_emp = ReSHOP.emp_create(ctx)

    if (length(emp.mps) > 0)
        ReSHOP.reshop_set_modeltype(ctx, ReSHOP.emp)
    else
        ReSHOP.reshop_set_modeltype(reshop_model)
    end

    # Define the EMP tree
    emptree(emp, reshop_emp, ctx)

    for ovf in emp.ovfs
        ReSHOP.reshop_ovf(reshop_emp, ovf)
    end

    setvarnames(ctx, emp.backend)

    gams_ctx, gams_dir = ReSHOP.reshop_setup_gams()
    reshop_model.reshop_ctx_dest = gams_ctx
    reshop_model.gams_dir = gams_dir

    reshop_model.solve_exitcode = ReSHOP.reshop_solve(ctx, gams_ctx, reshop_model.solver_name, reshop_emp)

    ###############################################################################################
    #
    #  Report values and codes
    #
    ###############################################################################################

    if reshop_model.solve_exitcode == 0
        ReSHOP.emp_report_values(reshop_emp, ctx)
        ReSHOP.report_results_common(reshop_model)
        _setvalues(emp.backend, getsolution(emp))
    else
        println("DEBUG: solver failed with status $(reshop_model.solve_exitcode)")
        reshop_model.status = :Error
        reshop_model.solution = fill(NaN,reshop_model.nvar)
        reshop_model.solve_result = "failure"
        reshop_model.solve_result_num = 999
    end

    #    TODO(xhub) move that elsewhere
    # This is left here for now, otherwise, we can't remove the test folder
    # This also plugs a memleak on a context
    ReSHOP.emp_delete(reshop_emp)
end

###################################################################################################
#
#   Accessor and such
#
###################################################################################################

getsolution(emp::EMPmaster) = copy(getReSHOPModel(emp).solution)
getsolvetime(emp::EMPmaster) = getReSHOPModel(emp).solve_time
status(emp::EMPmaster) = status(getReSHOPModel(emp))

function getobjval(mp::MathPrgm)
    reshop_model = getReSHOPModel(mp.emp)
    objvar = ReSHOP.emp_mp_getobjvar(mp.solverobj)
    ReSHOP.ctx_getvarval(reshop_model.reshop_ctx, objvar)
end

# Access to solve results
termination_status(emp::EMPmaster) = JuMP.termination_status(emp.backend)
primal_status(emp::EMPmaster) = JuMP.primal_status(emp.backend)
dual_status(emp::EMPmaster) = JuMP.dual_status(emp.backend)

function Base.show(io::IO, m::EMPmaster)
    println(io, "EMP data structure")
end

end


# test deprecated functions

T = Float32
fs = 100.  # sampling rate
N = 2^14   # number of points
Nwin = 10
freqmin = 10. # minimum frequency
freqmax = 20. # maximum frequency
corners = 4   # number of corners
zerophase = true # do zerophase filtering
maxlag = 20.     # maximum lag time in coda
starttime = d2u(DateTime(Date(now()))) # starttime for SeisChanel / SeisData

# test FFTData creation with compute_fft
@testset "compute fft" begin

    cc_len = 10.24
    cc_step = 5.12
    C = SeisChannel()
    C.x = rand(T,N)
    C.fs = fs
    C.id = "MZ.KTFD.08.ENZ"
    C.t = [1 (starttime + 0.004) * 1e6;N 0]
    R = RawData(C,cc_len,cc_step)
    R.gain = 1.234234e6
    F = compute_fft(R)
    @test F.name == C.id == R.name
    @test F.id == R.id
    @test F.loc == R.loc
    @test F.fs == R.fs
    @test F.gain == R.gain
    @test F.freqmin == R.freqmin
    @test F.freqmax == F.freqmax
    @test F.cc_len == R.cc_len == cc_len
    @test F.cc_step == R.cc_step == cc_step
    @test F.whitened == R.whitened
    @test F.time_norm == R.time_norm
    @test F.resp == R.resp
    @test F.t == R.t
    @test length(F.t) == size(F.fft,2)
end

# test correlation with FFTData
@testset "compute cc" begin
    cc_len = 40.96 # length of noise window
    cc_step = 40.96 # step between windows
    Ch = SeisIO.RandSeis.randSeisChannel(c=false,s=true)
    ungap!(Ch)
    Ch.x = rand(T,Int(cc_len*Nwin*fs +1)) .- T(0.5)
    Ch.fs = fs
    Ch.t = [1 starttime * 1e6;length(Ch.x) 0]
    Ch.loc = GeoLoc()
    R = RawData(Ch,cc_len,cc_step)
    F = rfft(R)
    C = compute_cc(F,F,maxlag)
    @test isa(C,CorrData) # test return type
    @test C.comp == repeat(Ch.id[end],2) # test component name
    @test size(C.corr) == (Int(2 * maxlag * fs) + 1,Nwin) # test size
    @test eltype(C.corr) == eltype(Ch.x) # test type stability
    t = -maxlag:1/fs:maxlag
    maxinds = argmax(C.corr,dims=1)
    @test all([t[maxinds[ii][1]] == 0 for ii in length(maxinds)]) # test max args
    @test C.rotated == false # test rotation
    @test C.corr_type == "CC" # test cross-correlation

    # test with no intersecting windows
    F1 = deepcopy(F)
    F2 = deepcopy(F)
    F2.t = rand(eltype(F2.t),size(F2.fft,2))
    @test_throws ArgumentError compute_cc(F1,F2,maxlag)

end


import GPLinearODEMaker.powers_of_negative_one

"""
    pp_kernel(hyperparameters, δ, dorder; shift_ind=0)

Created by kernel_coder(). Requires 1 hyperparameters.
Likely created using pp_kernel_base() as an input.
Use with include("src/kernels/pp_kernel.jl").

# Arguments
- `hyperparameters::Vector`: The hyperparameter values. For this kernel, they should be `["λ"]`
- `δ::Real`: The difference between the inputs (e.g. `t1 - t2`)
- `dorder::Vector{<:Integer}`: How many times to differentiate with respect to the inputs and the `hyperparameters` (e.g. `dorder=[0, 1, 0, 2]` would correspond to differentiating once w.r.t the second input and twice w.r.t `hyperparameters[2]`)
- `shift_ind::Integer=0`: If changed, the index of which hyperparameter is the `δ` shifting one
"""
function pp_kernel(
    hyperparameters::Vector{<:Real},
    δ::Real,
    dorder::Vector{<:Integer};
    shift_ind::Integer=0)

    include_shift = shift_ind!=0
    @assert length(hyperparameters)==1+Int(include_shift) "hyperparameters is the wrong length"
    dorder_len = 3 + Int(include_shift)
    @assert length(dorder)==dorder_len "dorder is the wrong length"
    @assert maximum(dorder) < 3 "No more than two derivatives for each time or hyperparameter can be calculated"

    dorder2 = dorder[2]
    dorder[2] += dorder[1]

    if include_shift
        dorder[2] += dorder[shift_ind+2]
        δ += hyperparameters[shift_ind]
        hyperparameters_view = view(hyperparameters, 1:2 .!= shift_ind)
        dorder_view = view(dorder, 2:dorder_len)
        dorder_view = view(dorder_view, 1:(dorder_len-1) .!= (shift_ind+1))
    else
        hyperparameters_view = hyperparameters
        dorder_view = view(dorder, 2:dorder_len)
    end

    λ = hyperparameters_view[1]

    if abs(δ) > λ
        dorder[2] = dorder2
        return 0
    end

    dabs_δ = powers_of_negative_one(δ < 0)  # store derivative of abs()
    abs_δ = abs(δ)

    if dorder_view==[6, 2]
        func = -403200*abs_δ*dabs_δ^6/λ^9 + 1209600*dabs_δ^6*(1 - abs_δ/λ)/λ^8 - 21600*dabs_δ^5*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^7 + 7200*dabs_δ^5*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^6 - 720*dabs_δ^5*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4)/λ^5
    end

    if dorder_view==[5, 2]
        func = -19200*abs_δ^2*dabs_δ^5/λ^9 - 288000*dabs_δ^5*(1 - abs_δ/λ)^2/λ^7 - 120*(48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*dabs_δ^5/λ^5 + 1200*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^5/λ^6 - 3600*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^5/λ^7 + 230400*abs_δ*dabs_δ^5*(1 - abs_δ/λ)/λ^8 - 6000*abs_δ*dabs_δ^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^7 + 1200*abs_δ*dabs_δ^4*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^6 + 12000*dabs_δ^4*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^6 - 4800*dabs_δ^4*(1 - abs_δ/λ)*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^5 + 600*dabs_δ^4*(1 - abs_δ/λ)*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4)/λ^4
    end

    if dorder_view==[4, 2]
        func = 38400*dabs_δ^4*(1 - abs_δ/λ)^3/λ^6 + 11520*abs_δ^2*dabs_δ^4*(1 - abs_δ/λ)/λ^8 - 57600*abs_δ*dabs_δ^4*(1 - abs_δ/λ)^2/λ^7 - 480*abs_δ^2*dabs_δ^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^7 - 2880*dabs_δ^3*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5 + 1440*dabs_δ^3*(1 - abs_δ/λ)^2*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^4 - 240*dabs_δ^3*(1 - abs_δ/λ)^2*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4)/λ^3 + 120*(48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*dabs_δ^4*(1 - abs_δ/λ)/λ^4 + 240*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*abs_δ*dabs_δ^4/λ^6 - 960*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^4*(1 - abs_δ/λ)/λ^5 - 1200*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^4/λ^7 + 2400*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^4*(1 - abs_δ/λ)/λ^6 + 3840*abs_δ*dabs_δ^3*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^6 - 960*abs_δ*dabs_δ^3*(1 - abs_δ/λ)*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^5
    end

    if dorder_view==[3, 2]
        func = -2880*dabs_δ^3*(1 - abs_δ/λ)^4/λ^5 - 2880*abs_δ^2*dabs_δ^3*(1 - abs_δ/λ)^2/λ^7 + 7680*abs_δ*dabs_δ^3*(1 - abs_δ/λ)^3/λ^6 + 360*dabs_δ^2*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4 - 240*dabs_δ^2*(1 - abs_δ/λ)^3*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^3 + 60*dabs_δ^2*(1 - abs_δ/λ)^3*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4)/λ^2 - 60*(48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^3 + 360*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^4 - 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ^2*dabs_δ^3/λ^7 - 720*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^5 + 360*abs_δ^2*dabs_δ^2*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^6 - 1080*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5 + 360*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^4 - 240*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*abs_δ*dabs_δ^3*(1 - abs_δ/λ)/λ^5 + 960*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^3*(1 - abs_δ/λ)/λ^6
    end

    if dorder_view==[2, 2]
        func = 96*dabs_δ^2*(1 - abs_δ/λ)^5/λ^4 + 320*abs_δ^2*dabs_δ^2*(1 - abs_δ/λ)^3/λ^6 - 480*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^4/λ^5 - 20*dabs_δ*(1 - abs_δ/λ)^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^3 + 20*dabs_δ*(1 - abs_δ/λ)^4*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^2 - 10*dabs_δ*(1 - abs_δ/λ)^4*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4)/λ + 20*(48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^2 - 80*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^3 + 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^4 - 120*abs_δ^2*dabs_δ*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5 + 160*abs_δ*dabs_δ*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4 - 80*abs_δ*dabs_δ*(1 - abs_δ/λ)^3*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^3 + 120*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2/λ^4 + 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ^2*dabs_δ^2*(1 - abs_δ/λ)/λ^6 - 360*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2/λ^5
    end

    if dorder_view==[1, 2]
        func = (1 - abs_δ/λ)^5*(10*dabs_δ/λ^3 + 96*abs_δ*dabs_δ/λ^4) + 20*abs_δ^2*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4 - 10*abs_δ*(1 - abs_δ/λ)^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^3 + 10*abs_δ*(1 - abs_δ/λ)^4*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^2 - 5*(48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*dabs_δ*(1 - abs_δ/λ)^4/λ + 10*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ*(1 - abs_δ/λ)^4/λ^2 - 10*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ*(1 - abs_δ/λ)^4/λ^3 - 40*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*abs_δ*dabs_δ*(1 - abs_δ/λ)^3/λ^3 - 60*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ^2*dabs_δ*(1 - abs_δ/λ)^2/λ^5 + 80*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ*(1 - abs_δ/λ)^3/λ^4
    end

    if dorder_view==[0, 2]
        func = (48*abs_δ^2/λ^4 + 10*abs_δ/λ^3)*(1 - abs_δ/λ)^5 + 10*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*abs_δ*(1 - abs_δ/λ)^4/λ^2 + 20*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ^2*(1 - abs_δ/λ)^3/λ^4 - 10*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*(1 - abs_δ/λ)^4/λ^3
    end

    if dorder_view==[6, 1]
        func = 28800*abs_δ*dabs_δ^6/λ^8 - 172800*dabs_δ^6*(1 - abs_δ/λ)/λ^7 + 3600*dabs_δ^5*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^6 - 720*dabs_δ^5*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^5
    end

    if dorder_view==[5, 1]
        func = 48000*dabs_δ^5*(1 - abs_δ/λ)^2/λ^6 - 120*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^5/λ^5 + 600*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^5/λ^6 - 19200*abs_δ*dabs_δ^5*(1 - abs_δ/λ)/λ^7 + 600*abs_δ*dabs_δ^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^6 - 2400*dabs_δ^4*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5 + 600*dabs_δ^4*(1 - abs_δ/λ)*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^4
    end

    if dorder_view==[4, 1]
        func = -7680*dabs_δ^4*(1 - abs_δ/λ)^3/λ^5 + 5760*abs_δ*dabs_δ^4*(1 - abs_δ/λ)^2/λ^6 + 720*dabs_δ^3*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4 - 240*dabs_δ^3*(1 - abs_δ/λ)^2*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^3 + 120*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^4*(1 - abs_δ/λ)/λ^4 + 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^4/λ^6 - 480*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^4*(1 - abs_δ/λ)/λ^5 - 480*abs_δ*dabs_δ^3*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5
    end

    if dorder_view==[3, 1]
        func = 720*dabs_δ^3*(1 - abs_δ/λ)^4/λ^4 - 960*abs_δ*dabs_δ^3*(1 - abs_δ/λ)^3/λ^5 - 120*dabs_δ^2*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^3 + 60*dabs_δ^2*(1 - abs_δ/λ)^3*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ^2 - 60*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^3 + 180*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^4 + 180*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4 - 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^3*(1 - abs_δ/λ)/λ^5
    end

    if dorder_view==[2, 1]
        func = -32*dabs_δ^2*(1 - abs_δ/λ)^5/λ^3 + 80*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^4/λ^4 + 10*dabs_δ*(1 - abs_δ/λ)^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^2 - 10*dabs_δ*(1 - abs_δ/λ)^4*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3)/λ + 20*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^2 - 40*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^3 - 40*abs_δ*dabs_δ*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^3 + 60*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ^2*(1 - abs_δ/λ)^2/λ^4
    end

    if dorder_view==[1, 1]
        func = (1 - abs_δ/λ)^5*(-5*dabs_δ/λ^2 - 32*abs_δ*dabs_δ/λ^3) + 5*abs_δ*(1 - abs_δ/λ)^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^2 - 5*(-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*dabs_δ*(1 - abs_δ/λ)^4/λ + 5*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ*(1 - abs_δ/λ)^4/λ^2 - 20*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*dabs_δ*(1 - abs_δ/λ)^3/λ^3
    end

    if dorder_view==[0, 1]
        func = (-16*abs_δ^2/λ^3 - 5*abs_δ/λ^2)*(1 - abs_δ/λ)^5 + 5*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*abs_δ*(1 - abs_δ/λ)^4/λ^2
    end

    if dorder_view==[6, 0]
        func = 28800*dabs_δ^6*(1 - abs_δ/λ)/λ^6 - 720*dabs_δ^5*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^5
    end

    if dorder_view==[5, 0]
        func = -9600*dabs_δ^5*(1 - abs_δ/λ)^2/λ^5 - 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^5/λ^5 + 600*dabs_δ^4*(1 - abs_δ/λ)*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^4
    end

    if dorder_view==[4, 0]
        func = 1920*dabs_δ^4*(1 - abs_δ/λ)^3/λ^4 - 240*dabs_δ^3*(1 - abs_δ/λ)^2*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^3 + 120*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^4*(1 - abs_δ/λ)/λ^4
    end

    if dorder_view==[3, 0]
        func = -240*dabs_δ^3*(1 - abs_δ/λ)^4/λ^3 + 60*dabs_δ^2*(1 - abs_δ/λ)^3*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ^2 - 60*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^3*(1 - abs_δ/λ)^2/λ^3
    end

    if dorder_view==[2, 0]
        func = 16*dabs_δ^2*(1 - abs_δ/λ)^5/λ^2 - 10*dabs_δ*(1 - abs_δ/λ)^4*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2)/λ + 20*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ^2*(1 - abs_δ/λ)^3/λ^2
    end

    if dorder_view==[1, 0]
        func = (1 - abs_δ/λ)^5*(5*dabs_δ/λ + 16*abs_δ*dabs_δ/λ^2) - 5*(1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*dabs_δ*(1 - abs_δ/λ)^4/λ
    end

    if dorder_view==[0, 0]
        func = (1 + 8*abs_δ^2/λ^2 + 5*abs_δ/λ)*(1 - abs_δ/λ)^5
    end

    dorder[2] = dorder2  # resetting dorder[2]
    return powers_of_negative_one(dorder2) * float(func)  # correcting for amount of t2 derivatives

end

return pp_kernel, 1  # the function handle and the number of kernel hyperparameters


using Oceananigans
using Oceananigans.Units
using OceanTurbulenceParameterEstimation
using Distributions
using DataDeps
using GLMakie

using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities:
    CATKEVerticalDiffusivity,
    SurfaceTKEFlux,
    MixingLength

include("best_catke_parameters.jl")

θ = [θ_constant_Pr, θ_variable_Pr, θ_conv_variable_Pr]
Nensemble = length(θ)

function perturbation_prior(θ★, ϵ=0.1)
    L = (1 - ϵ) * θ★
    U = (1 + ϵ) * θ★
    return ScaledLogitNormal(bounds=(L, U))
end

θ★ = θ_constant_Pr
priors = NamedTuple(name => perturbation_prior(θ★[name]) for name in keys(θ★))
free_parameters = FreeParameters(priors)

suites = [
          "two_day_suite_1m",
          "four_day_suite_1m",
          "six_day_suite_1m",
         ]

end_times = [
             2days,
             4days,
             6days
            ]

for (suite, end_time) in zip(suites, end_times)
    case_path(case) = @datadep_str("$suite/$(case)_instantaneous_statistics.jld2")

    cases = [
             "free_convection",
             "weak_wind_strong_cooling",
             "strong_wind_weak_cooling",
             "strong_wind",
             "strong_wind_no_rotation",
            ]

    Δz = 1
    times = [6hours, end_time]
    field_names = (:u, :v, :T, :e)
    regrid_size = (1, 1, Int(256/Δz))
    observations = [SyntheticObservations(case_path(case); field_names, regrid_size, times) for case in cases]
    obs = first(observations)
    α = obs.metadata.parameters.thermal_expansion_coefficient
    equation_of_state = LinearEquationOfState(; α)
    buoyancy = SeawaterBuoyancy(; equation_of_state, constant_salinity=true)

    mixing_length = MixingLength(Cᴬu   = 0.0,
                                 Cᴬc   = 0.0,
                                 Cᴬe   = 0.0,
                                 CᴷRiᶜ = 2.0,
                                 Cᴷuʳ  = 0.0,
                                 Cᴷcʳ  = 0.0,
                                 Cᴷeʳ  = 0.0)

    catke = CATKEVerticalDiffusivity(; mixing_length)

    simulation = ensemble_column_model_simulation(observations; Nensemble, buoyancy,
                                                  architecture = GPU(),
                                                  tracers = (:T, :e),
                                                  closure = catke)

    simulation.Δt = 1.0
    progress(sim) = @info "Iter: $(iteration(sim)), time: $(prettytime(sim))"
    simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

    Qᵘ = simulation.model.velocities.u.boundary_conditions.top.condition
    Qᵀ = simulation.model.tracers.T.boundary_conditions.top.condition
    dTdz = simulation.model.tracers.T.boundary_conditions.bottom.condition

    for (case, obs) in enumerate(observations)
        @show case cases[case]
        @show obs.metadata.parameters.momentum_flux
        @show obs.metadata.parameters.temperature_flux
        @show obs.metadata.parameters.buoyancy_flux
        @show f = obs.metadata.parameters.coriolis_parameter

        view(Qᵘ, :, case) .= obs.metadata.parameters.momentum_flux
        view(Qᵀ, :, case) .= obs.metadata.parameters.temperature_flux
        view(dTdz, :, case) .= obs.metadata.parameters.dθdz_deep
        view(simulation.model.coriolis, :, case) .= Ref(FPlane(f=f))
    end

    calibration = InverseProblem(observations, simulation, free_parameters)

    model = simulation.model

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, merge(model.velocities, model.tracers, model.diffusivity_fields),
                         prefix = string("catke_simulation_", suite[1:end-2], Δz, "m"),
                         schedule = SpecifiedTimes(times[1]:10minutes:end_time...),
                         force = true)

    @show simulation.output_writers[:fields]

    forward_run!(calibration, θ; suppress=false)
end


export restart, step_forward, get_state

function restart(roadway::Roadway, render::Bool)
    # Create scene
    scene = Scene()
    push!(scene,Vehicle(VehicleState(VecSE2(5.0 + rand()*10.0,randn()*0.5,0.0), roadway, 20.0),
                    VehicleDef(1, AgentClass.CAR, 4.826, 1.81)))
    push!(scene,Vehicle(VehicleState(VecSE2(0.,randn()*0.5,0.0), roadway, 10.0 + rand()*20.0),
                        VehicleDef(2, AgentClass.CAR, 4.826, 1.81)))

    # Specify models
    context = IntegratedContinuous(1/10,3) #1/framrate
    models = Dict{Int, DriverModel}()

    if render
    	model1 = LatLonSeparableDriver(context,
	                                ProportionalLaneTracker(σ=0.0, kp=5.0, kd=0.5),
	                                IntelligentDriverModel(σ=0.0, v_des=21.0))
	    model2 = LatLonSeparableDriver(context,
	                                ProportionalLaneTracker(σ=0.0, kp=5.0, kd=0.5),
	                                StaticLongitudinalDriver())
    else
	    model1 = LatLonSeparableDriver(context,
	                                ProportionalLaneTracker(σ=0.01, kp=5.0, kd=0.5),
	                                IntelligentDriverModel(σ=0.2, v_des=21.0))
	    model2 = LatLonSeparableDriver(context,
	                                ProportionalLaneTracker(σ=0.01, kp=5.0, kd=0.5),
	                                StaticLongitudinalDriver())
	end

    # Placeholder for actions
    return scene, model1, model2
end
function step_forward(scene::Scene, roadway::Roadway, model1::DriverModel, model2::DriverModel, action::Float64)

    observe!(model1, scene, roadway, 1)
    context = action_context(model1)
    scene[1].state = propagate(scene[1], rand(model1), context, roadway)

    model2.mlon.a = action
    observe!(model2, scene, roadway, 2)
    context = action_context(model2)
    scene[2].state = propagate(scene[2], rand(model2), context, roadway)

    return get_state(scene)
end
function get_state(scene::Scene)
	# Extract state values
	d = scene[1].state.posF.s - scene[2].state.posF.s - scene[1].def.length/2. - scene[2].def.length/2.
	r = scene[2].state.v - scene[1].state.v
    s = scene[1].state.v
    return d, r, s
end


module color

import Term: NAMED_COLORS, nospaces, COLORS_16b, remove_brackets, ANSICode, int, CODES, CODES_16BIT_COLORS

export NamedColor, BitColor, RGBColor, get_color

# ----------------------------- types definition ----------------------------- #
"""
    AbstractColor

Abstract color type.
"""
abstract type AbstractColor end

Base.show(io::IO, color::AbstractColor) = print(io, "$(typeof(color))")

struct NamedColor <: AbstractColor
    color::String
end

struct BitColor <: AbstractColor
    color::String
end

struct RGBColor <: AbstractColor
    r::Int
    g::Int
    b::Int
end

function RGBColor(s)
    to_number(x) = '.' ∈ x ? parse(Float64, x) : parse(Int, x)
    r, g, b = to_number.(match(RGB_REGEX, s).captures)
    if r < 1 || g < 1 || b < 1
        r *= 255
        g *= 255
        b *= 255
    end
    return RGBColor(int(r), int(g), int(b))
end

"""
    ANSICode(color; bg::Bool=false)

Create ANSI tags for colors.
"""
function ANSICode(color::NamedColor; bg::Bool = false)
    Δ = bg ? 40 : 30
    v = CODES[color.color]
    return ANSICode("\e[$(Δ + v)m", "\e[$(Δ+9)m")
end

function ANSICode(color::BitColor; bg::Bool = false)
    Δ = bg ? 48 : 38
    v = CODES_16BIT_COLORS[color.color]
    return ANSICode("\e[$Δ;5;$(v)m", "\e[$(Δ+1)m")
end

function ANSICode(color::RGBColor; bg::Bool = false)
        Δ = bg ? 48 : 38
        rgb = "$(color.r);$(color.g);$(color.b)"
        return ANSICode("\e[$Δ;2;$(rgb)m", "\e[$(Δ+1)m")
end

# --------------------------------- is color? -------------------------------- #

RGB_REGEX = r"\(\s*([\d\.]{1,3})\s*,\s*([\d\.]{1,3})\s*,\s*([\d\.]{1,3})\s*\)"
HEX_REGEX = r"#(?:[0-9a-fA-F]{3}){1,2}$"

"""
    is_named_color(string::String)::Bool

Check if a string represents a named color.
"""
is_named_color(string)::Bool = string ∈ NAMED_COLORS || string ∈ COLORS_16b

"""
    is_rgb_color(string::String)::Bool

Check if a string represents a RGB color.
"""
is_rgb_color(string)::Bool = !occursin("on_", string) && occursin(RGB_REGEX, string)

"""
    is_hex_color(string::String)::Bool

Check if a string represents a hex color.
"""
is_hex_color(string)::Bool = !occursin("on_", string) && occursin(HEX_REGEX, string)

"""
    is_color(string::String)::Bool

Check if a string represents color information, of any type.
"""
is_color(string)::Bool = is_named_color(string) || is_rgb_color(string) || is_hex_color(string)

"""
    is_background(string::String)::Bool

Check if a string represents background color information, of any type.
"""
function is_background(string)::Bool
    stripped = nospaces(string)
    length(stripped) < 3 && return false
    return stripped[1:3] == "on_" && is_color(stripped[4:end])
end

# --------------------------------- get color -------------------------------- #
"""
    hex2rgb(hex::String)

Converts a string hex color code to a RGB color
"""
function hex2rgb(hex)::RGBColor
    to_int(h) = parse(Int, h; base = 16)
    r, g, b = [to_int(hex[i:(i + 1)]) for i in (2, 4, 6)]
    return RGBColor(r, g, b)
end

"""
    get_color(string::String; bg=false)::AbstractColor

Extract a color type from a string with color information.
"""
function get_color(string; bg = false)::AbstractColor
    if bg
        string = nospaces(string)[4:end]
    end

    if is_named_color(string)
        if string ∈ COLORS_16b
            idx = findfirst((c) -> c == string, COLORS_16b)[1]
            return BitColor(COLORS_16b[idx])
        else
            idx = findfirst((c) -> c == string, NAMED_COLORS)[1]
            return NamedColor(NAMED_COLORS[idx])
        end
    elseif is_rgb_color(string)
        return RGBColor(string)
    else
        # convert hex to rgb
        return hex2rgb(string)
    end
end

"""
Convert HSL values to RGB values, return a markup color string.
"""
function hsl2rgb(h, s, l)
    C = (1 - abs(2*l - 1)) * s
    X = C * (1 - abs(mod(h/60, 2) - 1))
    M = l - C/2

    if 0 ≤ h < 60
        r, g, b = C, X, 0
    elseif 60 ≤ h < 120
        r, g, b = X, C, 0
    elseif 120 ≤ h < 180
        r, g, b = 0, C, X
    elseif 180 ≤ h < 240
        r, g, b = 0, X, C
    elseif 240 ≤ h < 300
        r, g, b = X, 0, C
    elseif 300 ≤ h ≤ 360
        r, g, b = C, 0, X
    end

    r = (Int ∘ round)((r+M) * 255)
    g = (Int ∘ round)((g+M) * 255)
    b = (Int ∘ round)((b+M) * 255)

    return "($r, $g, $b)"
end
end


module Metrics

using util
using DataStructures
using PyCall

export createMAP_RecallRate, createBinaryPredictionROC, reset, update_metric!, compute, MAP_RecallRate, BinaryPredictionROC

abstract type Metric end

mutable struct MAP_RecallRate <: Metric
    label::String
    k::Int
    sum_map::Float64
    hits_per_k::Dict{Int, Float64}
    n_duplicates::Int
    group_by_master::Bool
end

function createMAP_RecallRate(;label="MAP_RecallRate", k=20, group_by_master=true)
    return MAP_RecallRate(label, k, 0.0, Dict{Int, Float64}([ (i, 0.0) for i in 1:k]), 0, group_by_master)
end

function reset(metric::MAP_RecallRate)
    metric.hits_per_k = Dict{Int, Float64}([ (i, 0.0) for i in 1:k])
    metric.n_duplicates = 0
    metric.sum_map = 0.0
end

function update_metric!(metric::MAP_RecallRate, query::Report, recommendation_list::Vector{Tuple{Report, Float64}})
    if query.dup_id == query.id
        #println("Simgleton: $(query.id) $([ (r.id, s) for (r, s)  in recommendation_list[1:20]])")
        return
    end

    buckets = Set()
    pos = Inf
    correct_cand = Int64(-1)

    for (cand, score) in recommendation_list
        if metric.group_by_master
            query.dup_id in buckets && continue
            push!(buckets, cand.dup_id)
        else
            push!(buckets, cand.id)
        end

        if cand.dup_id == query.dup_id
            pos = Float64(length(buckets))

            # println("$(query.id) $(cand.id) $(pos) $([(r.id, s) for (r, s)  in recommendation_list[1:20]])")
            # @debug
            # print("{}, {}, {}, {}, {}".format(report_id, pos, cand_id, self.master_id_by_bug_id[report_id], list(
            #     zip(candidates[:30], map(lambda x: self.master_id_by_bug_id[x], candidates[:30]), scores[:30]))))
            correct_cand = cand.id
            break
        end
    end

    for k in 1:metric.k
        k < pos && continue

        metric.hits_per_k[k]+=1.0
    end

    metric.sum_map += 1.0/pos
    metric.n_duplicates += 1

    return pos
end

function compute(metric::MAP_RecallRate)
    recall_rate = Dict{Int,Float64}()

    for k in 1:metric.k
        recall_rate[k] = metric.hits_per_k[k] / metric.n_duplicates
    end

    return Dict("type" => "metric",
            "label" => metric.label,
            "hits_per_k" => metric.hits_per_k,
            "rr" => recall_rate,
            "sum_map" => metric.sum_map,
            "map"=> metric.sum_map / metric.n_duplicates,
            "total" => metric.n_duplicates,
            )
end

mutable struct BinaryPredictionROC <: Metric
    label::String
    queries::MutableLinkedList{Report}
    scores::MutableLinkedList{Float64}
end

function  createBinaryPredictionROC(;label="BinaryPredictionROC")
    return BinaryPredictionROC(label, MutableLinkedList{Report}(), MutableLinkedList{Float64}())
end

function reset!(metric::BinaryPredictionROC)
    metric.queries = MutableLinkedList{Report}()
    metric.scores = MutableLinkedList{Float64}()
end

function update_metric!(metric::BinaryPredictionROC, query::Report, recommendation_list::Vector{Tuple{Report, Float64}})
    top_candidate = recommendation_list[1][1]
    top_score = recommendation_list[1][2]

    push!(metric.queries, query)
    push!(metric.scores, top_score)

        # print("{} {} {} {}".format(query_id, top_candidate, top_score, self.master_id_by_bug_id[query_id] != query_id))

    # print("{}, {}, {}, {}, {}, {}, {}".format(
    #     self.master_id_by_bug_id[query_id] == query_id, query_id,
    #     self.master_id_by_bug_id[query_id], top_candidate, self.master_id_by_bug_id[top_candidate], top_score,
    #     [(k, s) for k, s in zip(candidates[:20], scores[:20])]))

end

function compute(metric::BinaryPredictionROC)
    SKLEARN_METRICS = pyimport("sklearn.metrics")
    y_true = [Int(query.dup_id != query.id) for query in metric.queries]
    fpr, tpr, thresholds = SKLEARN_METRICS.roc_curve(y_true, collect(metric.scores), pos_label=1)

    return Dict("type" => "metric",
                "label" => metric.label,
                "fpr" => fpr,
                "tpr" => tpr,
                "auc" => SKLEARN_METRICS.auc(fpr,tpr),
                "n_queries" =>  length(metric.queries),
                "threshold" => thresholds,
    )

end

end


function test_set_rxn()
    model = MetNets.toy_model()

    rxn = MetNets.Rxn("new_rxn";
        lb = -rand(), ub = rand(), c = rand())
    for met in model.mets
        push!(rxn.mets, met)
        push!(rxn.S, rand([0,-1, 1]))
    end

    rxni = 1
    MetNets.set_rxn!(model, rxni, rxn)

    # Test
    @test model.rxns[rxni] == rxn.id
    @test model.lb[rxni] == rxn.lb
    @test model.ub[rxni] == rxn.ub
    @test model.c[rxni] == rxn.c

    for (met, s) in zip(rxn.mets, rxn.S)
        meti = MetNets.metindex(model, met)
        @test model.S[meti, rxni] == s
    end
end
test_set_rxn()


#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

using Printf
using Base.Threads
import Base.Threads: @threads, nthreads, threadid

function solve!(
    mip::MIP;
    time_limit::Float64 = Inf,
    node_limit::Int = typemax(Int),
    gap_limit::Float64 = 1e-4,
    print_interval::Int = 5,
    initial_primal_bound::Float64 = Inf,
    detailed_output::Bool = false,
    branch_rule::VariableBranchingRule = HybridBranching(),
)::NodePool
    time_initial = time()
    pool = NodePool(mip = mip)
    pool.primal_bound = initial_primal_bound
    print_progress_header(detailed_output = detailed_output)

    root_node = _create_node(mip)
    if isempty(root_node.fractional_variables)
        println("root relaxation is integer feasible")
        pool.dual_bound = root_node.obj
        pool.primal_bound = root_node.obj
        return pool
    end

    offer(pool, parent_node = nothing, child_nodes = [root_node])
    @threads for t = 1:nthreads()
        child_one, child_zero, suggestions = nothing, nothing, Node[]
        while true
            time_elapsed = time() - time_initial
            if child_one !== nothing
                suggestions = Node[child_one, child_zero]
            end
            node = take(
                pool,
                suggestions = suggestions,
                time_remaining = time_limit - time_elapsed,
                node_limit = node_limit,
                gap_limit = gap_limit,
            )
            if node == :END
                break
            elseif node == :WAIT
                sleep(0.1)
                continue
            else
                ids = generate_indices(pool, 2)
                branch_variable = find_branching_var(branch_rule, node, pool)
                child_zero = _create_node(
                    mip,
                    index = ids[1],
                    parent = node,
                    branch_variable = branch_variable,
                    branch_value = 0.0,
                )
                child_one = _create_node(
                    mip,
                    index = ids[2],
                    parent = node,
                    branch_variable = branch_variable,
                    branch_value = 1.0,
                )
                offer(
                    pool,
                    parent_node = node,
                    child_nodes = [child_one, child_zero],
                    time_elapsed = time_elapsed,
                    print_interval = print_interval,
                    detailed_output = detailed_output,
                )
            end
        end
    end
    return pool
end

function _create_node(
    mip;
    index::Int = 0,
    parent::Union{Nothing,Node} = nothing,
    branch_variable::Union{Nothing,Variable} = nothing,
    branch_value::Union{Nothing,Float64} = nothing,
)::Node
    if parent === nothing
        branch_variables = Variable[]
        branch_values = Float64[]
        depth = 1
    else
        branch_variables = [parent.branch_variables; branch_variable]
        branch_values = [parent.branch_values; branch_value]
        depth = parent.depth + 1
    end
    set_bounds!(mip, branch_variables, branch_values, branch_values)
    status, obj = solve_relaxation!(mip)
    if status == :Optimal
        vals = values(mip, mip.binary_variables)
        fractional_indices =
            [j for j in 1:length(mip.binary_variables) if 1e-6 < vals[j] < 1 - 1e-6]
        fractional_values = vals[fractional_indices]
        fractional_variables = mip.binary_variables[fractional_indices]
    else
        fractional_variables = Variable[]
        fractional_values = Float64[]
    end
    n_branch = length(branch_variables)
    set_bounds!(mip, branch_variables, zeros(n_branch), ones(n_branch))
    return Node(
        mip,
        index,
        depth,
        obj,
        status,
        branch_variables,
        branch_values,
        fractional_variables,
        fractional_values,
        parent,
    )
end


@enum coordstype auto griddata pointdata

# Insertion point for defining my own functions. They are mainly abstract
# over the full type and it's view, so another possibility would be to
# implement with unions
abstract type SESpatialData{T<:EcoBase.AbstractLocationData} <: EcoBase.AbstractPlaces{T} end
abstract type SELocations{T<:EcoBase.AbstractLocationData} <: EcoBase.AbstractPlaces{T} end
abstract type SEThings{D <: Real} <: EcoBase.AbstractThings end
abstract type SEGrid <: EcoBase.AbstractGrid end
abstract type SEPoints <: EcoBase.AbstractPoints end

abstract type AbstractComMatrix{D<:Real} end

# I could implement sitestats as a Dict with several DataFrames to make space for big data sets, but I prefer to not do this now. Example below.

# I could do a lot more with immutable types if I had a clearer view/copy implementation
mutable struct GridTopology <: EcoBase.AbstractGrid
    xs::StepRangeLen{Float64}
    ys::StepRangeLen{Float64}
end

mutable struct Bbox
    xmin::Number
    xmax::Number
    ymin::Number
    ymax::Number
end

# Do I need sitenames here? I think so, they should match those in sitestats, and be separate
mutable struct PointData <: SEPoints
    coords::Matrix{Float64}
end

mutable struct GridData <: SEGrid
    indices::Matrix{Int}
    grid::GridTopology
end

mutable struct Locations{T<:Union{GridData, PointData}} <: SELocations{T}
    coords::T
    sitestats::DataFrames.DataFrame
    function Locations{T}(coords, sitestats = DataFrames.DataFrame(sites = string.(1:nsites(coords)))) where T
        DataFrames.nrow(sitestats) == nsites(coords) || throw(DimensionMismatch("Wrong number of rows in sitestat")) # a little check for the right number
        new(coords, sitestats)
    end
end

Locations(coords, sitestats) = Locations{typeof(coords)}(coords, sitestats)

#it's a question how many of these structs need to be mutable, as opposed to when you want to allocate a new object
mutable struct ComMatrix{D} <: AbstractComMatrix{D}
    occurrences::SparseMatrixCSC{D}
    speciesnames::Vector{String}
    sitenames::Vector{String}
    ComMatrix{D}(occ::SparseMatrixCSC{D}, spn::Vector{String}, sin::Vector{String}) where {D} = new(dropzeros!(occ), spn, sin)
end

# likewise, do I need a speciesnames here? Should traits have a :series field (like now) or all matching be done on the speciesnames?
mutable struct SpeciesData{D} <: SEThings{D}
    commatrix::ComMatrix{D}
    traits::DataFrames.DataFrame
    function SpeciesData{D}(commatrix::ComMatrix{D}, traits::DataFrames.DataFrame) where D <: Real
        DataFrames.nrow(traits) ==  nspecies(commatrix) || throw(DimensionMismatch("Wrong number of species in traits"))
        new(commatrix, traits)
    end
end
SpeciesData(commatrix::ComMatrix{D}, traits) where D<:Real = SpeciesData{D}(commatrix, traits)

# Not really sure what this type is for
mutable struct SiteData{T<:Union{GridData, PointData}} <: SESpatialData{T}
    site::T
end

abstract type SEAssemblage{D<:Real, T<:SEThings, P<:SELocations} <: EcoBase.AbstractAssemblage{D, T, P} end

mutable struct Assemblage{D<:Real, P<:Locations} <: SEAssemblage{D, SpeciesData{D}, P} # A type to keep subtypes together, ensuring that they are all aligned at all times
    site::P
    occ::SpeciesData{D}

    # inner constructor
    function Assemblage{D, P}(site::P, occ::SpeciesData{D}) where {P <: SELocations, D <: Real}
        size(occurrences(occ), 2) == size(coordinates(site), 1) || error("Length mismatch between occurrence matrix and coordinates")
        #TODO activate this # sitenames(occ) == sitenames(site) || error("sitenames do not match") #I need a constructor that matches them up actively
        new(site, occ)
    end
end


"""
# module Reciprocal

# Examples

```jldoctest
julia>
```
"""
module Reciprocal

using StaticArrays: SVector

export KPoint

struct KPoint
    coordinates::SVector{3, Float64}
    weight::Float64
end

end


include("../src/NeuralNetwork.jl")
include("../src/MNIST/src/MNIST.jl")

using .NeuralNetwork
using .MNIST

trd = traindata()
tsd = testdata()

net = read(joinpath(@__DIR__, "net1"), Chain)

test(net, tsd)


# ============================================================
# Structs of Initial & Boundary Conditions
# ============================================================

"""
    mutable struct IB{T} <: AbstractCondition
        fw::Function
        bc::T
    end

Initial & boundary conditions with no distribution function

"""
mutable struct IB{T<:Union{Function,AbstractArray}} <: AbstractCondition
    fw::Function
    bc::T
end

function IB(fw, gas::Scalar)
    γ = gas.a

    bc = function (args...)
        w = fw(args...)
        return ifelse(γ == 0, conserve_prim(w), conserve_prim(w, γ))
    end

    return IB{typeof(bc)}(fw, bc)
end

function IB(fw, gas::AbstractGas)
    bc = function (args...)
        w = fw(args...)
        prim = begin
            if ndims(w) == 1
                conserve_prim(w, gas.γ)
            else
                mixture_conserve_prim(w, gas.γ)
            end
        end

        return prim
    end

    return IB{typeof(bc)}(fw, bc)
end

"""
    mutable struct IB1F{T} <: AbstractCondition
        fw::Function
        ff::Function
        bc::T
    end

Initial & boundary condition with 1 distribution function

"""
mutable struct IB1F{T} <: AbstractCondition
    fw::Function
    ff::Function
    bc::T
end

function IB1F(fw, vs::AbstractVelocitySpace, gas::AbstractProperty)
    bc = function (args...)
        w = fw(args...)
        prim = begin
            if ndims(w) == 1
                conserve_prim(w, gas.γ)
            else
                mixture_conserve_prim(w, gas.γ)
            end
        end

        return prim
    end

    ff = function (args...)
        prim = bc(args...)

        M = begin
            if !isdefined(vs, :v)
                if ndims(prim) == 1
                    maxwellian(vs.u, prim)
                else
                    mixture_maxwellian(vs.u, prim)
                end
            elseif !isdefined(vs, :w)
                if ndims(prim) == 1
                    maxwellian(vs.u, vs.v, prim)
                else
                    mixture_maxwellian(vs.u, vs.v, prim)
                end
            else
                if ndims(prim) == 1
                    maxwellian(vs.u, vs.v, vs.w, prim)
                else
                    mixture_maxwellian(vs.u, vs.v, vs.w, prim)
                end
            end
        end

        return M
    end

    return IB1F{typeof(bc)}(fw, ff, bc)
end

"""
    mutable struct IB2F{T} <: AbstractCondition
        fw::Function
        ff::Function
        bc::T
    end

Initial & boundary condition with 2 distribution functions

"""
mutable struct IB2F{T} <: AbstractCondition
    fw::Function
    ff::Function
    bc::T
end

function IB2F(fw, vs::AbstractVelocitySpace, gas::AbstractProperty)
    bc = function (args...)
        w = fw(args...)
        prim = begin
            if ndims(w) == 1
                conserve_prim(w, gas.γ)
            else
                mixture_conserve_prim(w, gas.γ)
            end
        end

        return prim
    end

    ff = function (args...)
        prim = bc(args...)

        if !isdefined(vs, :v)
            if ndims(prim) == 1
                H = maxwellian(vs.u, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, prim)
                B = zero(H)
                for j in axes(B, 2)
                    B[:, j] = H[:, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        elseif !isdefined(vs, :w)
            if ndims(prim) == 1
                H = maxwellian(vs.u, vs.v, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, vs.v, prim)
                B = zero(H)
                for j in axes(B, 3)
                    B[:, :, j] = H[:, :, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        else
            if ndims(prim) == 1
                H = maxwellian(vs.u, vs.v, vs.w, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, vs.v, vs.w, prim)
                B = zero(H)
                for j in axes(B, 4)
                    B[:, :, :, j] = H[:, :, :, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        end

        return H, B
    end

    return IB2F{typeof(bc)}(fw, ff, bc)
end

"""
    mutable struct IB3F{T} <: AbstractCondition
        fw::Function
        ff::Function
        fE::Function
        fB::Function
        fL::Function
        bc::T
    end

Initial & boundary condition with 3 distribution functions

"""
mutable struct IB3F{T} <: AbstractCondition
    fw::Function
    ff::Function
    fE::Function
    fB::Function
    fL::Function
    bc::T
end

"""
    mutable struct IB4F{T} <: AbstractCondition
        fw::Function
        ff::Function
        fE::Function
        fB::Function
        fL::Function
        bc::T
    end

Initial & boundary condition with 4 distribution functions

"""
mutable struct IB4F{T} <: AbstractCondition
    fw::Function
    ff::Function
    fE::Function
    fB::Function
    fL::Function
    bc::T
end


using Base: @propagate_inbounds

using OffsetArrays
using JLD2

using Oceananigans.Architectures
using Oceananigans.Grids
using Oceananigans.Fields

using Oceananigans.Grids: topology, total_size, interior_parent_indices, parent_index_range
using Oceananigans.Fields: show_location, interior_view_indices, data_summary, reduced_location

import Oceananigans.Fields: Field, set!, interior, indices
import Oceananigans.Architectures: architecture

struct FieldTimeSeries{LX, LY, LZ, K, I, D, G, T, B, χ} <: AbstractField{LX, LY, LZ, G, T, 4}
                   data :: D
                   grid :: G
    boundary_conditions :: B
                indices :: I
                  times :: χ

    function FieldTimeSeries{LX, LY, LZ, K}(data::D, grid::G, bcs::B,
                                            times::χ, indices::I) where {LX, LY, LZ, K, D, G, B, χ, I}
        T = eltype(data)
        return new{LX, LY, LZ, K, I, D, G, T, B, χ}(data, grid, bcs, indices, times)
    end
end

architecture(fts::FieldTimeSeries) = architecture(fts.grid)

#####
##### Constructors
#####

"""
    FieldTimeSeries{LX, LY, LZ}(grid, times, [FT=eltype(grid);]
                                indices = (:, :, :),
                                boundary_conditions = nothing)

Return a `FieldTimeSeries` at location `(LX, LY, LZ)`, on `grid`, at `times`.
"""
function FieldTimeSeries{LX, LY, LZ}(grid, times, FT=eltype(grid);
                                     indices = (:, :, :),
                                     boundary_conditions = nothing) where {LX, LY, LZ}

    Nt = length(times)
    arch = architecture(grid)
    loc = (LX, LY, LZ)
    space_size = total_size(loc, grid, indices)
    underlying_data = zeros(FT, arch, space_size..., Nt)
    data = offset_data(underlying_data, grid, loc, indices)

    return FieldTimeSeries{LX, LY, LZ, InMemory}(data, grid, boundary_conditions, times, indices)
end

"""
    FieldTimeSeries(path, name;
                    backend = InMemory(),
                    grid = nothing,
                    iterations = nothing,
                    times = nothing)

Returns a `FieldTimeSeries` for the field `name` describing a field's time history from a JLD2 file
located at `path`.

Keyword arguments
=================

- `backend`: `InMemory()` to load data into a 4D array or `OnDisk()` to lazily load data from disk
             when indexing into `FieldTimeSeries`.

- `grid`: A grid to associated with data, in the case that the native grid was not serialized
          properly.

- `iterations`: Iterations to load. Defaults to all iterations found in the file.

- `times`: Save times to load, as determined through an approximate floating point
           comparison to recorded save times. Defaults to times associated with `iterations`.
           Takes precedence over `iterations` if `times` is specified.
"""
FieldTimeSeries(path, name; backend=InMemory(), kwargs...) =
    FieldTimeSeries(path, name, backend; kwargs...)

#####
##### InMemory time serieses
#####

const InMemoryFieldTimeSeries{LX, LY, LZ} = FieldTimeSeries{LX, LY, LZ, InMemory}

struct UnspecifiedBoundaryConditions end

function FieldTimeSeries(path, name, backend::InMemory;
                         architecture = nothing,
                         grid = nothing,
                         location = nothing,
                         boundary_conditions = UnspecifiedBoundaryConditions(),
                         iterations = nothing,
                         times = nothing)

    file = jldopen(path)

    # Defaults
    isnothing(iterations)   && (iterations = parse.(Int, keys(file["timeseries/t"])))
    isnothing(times)        && (times      = [file["timeseries/t/$i"] for i in iterations])
    isnothing(location)     && (location   = file["timeseries/$name/serialized/location"])

    if boundary_conditions isa UnspecifiedBoundaryConditions
        boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"]
    end

    indices = try
        file["timeseries/$name/serialized/indices"]
    catch
        (:, :, :)
    end

    isnothing(grid) && (grid = file["serialized/grid"])

    # Default to CPU if neither architecture nor grid is specified
    architecture = isnothing(architecture) ? (isnothing(grid) ? CPU() : Architectures.architecture(grid)) : architecture

    # This should be removed in a month or two (4/5/2022).
    grid = try
        on_architecture(architecture, grid)
    catch err # Likely, the grid has CuArrays in it...
        if grid isa RectilinearGrid # we can try...
            Nx = file["grid/Nx"]
            Ny = file["grid/Ny"]
            Nz = file["grid/Nz"]
            Hx = file["grid/Hx"]
            Hy = file["grid/Hy"]
            Hz = file["grid/Hz"]
            xᶠᵃᵃ = file["grid/xᶠᵃᵃ"]
            yᵃᶠᵃ = file["grid/yᵃᶠᵃ"]
            zᵃᵃᶠ = file["grid/zᵃᵃᶠ"]
            x = file["grid/Δxᶠᵃᵃ"] isa Number ? (xᶠᵃᵃ[1], xᶠᵃᵃ[Nx+1]) : xᶠᵃᵃ
            y = file["grid/Δyᵃᶠᵃ"] isa Number ? (yᵃᶠᵃ[1], yᵃᶠᵃ[Ny+1]) : yᵃᶠᵃ
            z = file["grid/Δzᵃᵃᶠ"] isa Number ? (zᵃᵃᶠ[1], zᵃᵃᶠ[Nz+1]) : zᵃᵃᶠ
            topo = topology(grid)

            # Reduce for Flat dimensions
            domain = NamedTuple((:x, :y, :z)[i] => (x, y, z)[i] for i=1:3 if topo[i] !== Flat)
            size = Tuple((Nx, Ny, Nz)[i] for i=1:3 if topo[i] !== Flat)
            halo = Tuple((Hx, Hy, Hz)[i] for i=1:3 if topo[i] !== Flat)

            RectilinearGrid(architecture; size, halo, topology=topo, domain...)
        else
            throw(err)
        end
    end

    close(file)

    LX, LY, LZ = location
    time_series = FieldTimeSeries{LX, LY, LZ}(grid, times; indices, boundary_conditions)

    set!(time_series, path, name)

    return time_series
end

Base.parent(fts::FieldTimeSeries) = parent(fts.data)

function Base.getindex(fts::InMemoryFieldTimeSeries, n::Int)
    underlying_data = view(parent(fts), :, :, :, n)
    data = offset_data(underlying_data, fts.grid, location(fts), fts.indices)
    boundary_conditions = fts.boundary_conditions
    indices = fts.indices
    return Field(location(fts), fts.grid; data, boundary_conditions, indices)
end

#####
##### set!
#####

"""
    Field(location, path, name, iter;
          grid = nothing,
          architecture = nothing,
          indices = (:, :, :),
          boundary_conditions = nothing)

Load a field called `name` saved in a JLD2 file at `path` at `iter`ation.
Unless specified, the `grid` is loaded from `path`.
"""
function Field(location, path::String, name::String, iter;
               grid = nothing,
               architecture = nothing,
               indices = (:, :, :),
               boundary_conditions = nothing)

    file = jldopen(path)

    # Default to CPU if neither architecture nor grid is specified
    architecture = isnothing(architecture) ?
        (isnothing(grid) ? CPU() : Architectures.architecture(grid)) :
        architecture

    grid = isnothing(grid) ?
        on_architecture(architecture, file["serialized/grid"]) : grid

    raw_data = arch_array(architecture, file["timeseries/$name/$iter"])

    close(file)

    data = offset_data(raw_data, grid, location, indices)

    return Field(location, grid; boundary_conditions, indices, data)
end

function set!(time_series::InMemoryFieldTimeSeries, path::String, name::String)

    file = jldopen(path)
    file_iterations = parse.(Int, keys(file["timeseries/t"]))
    file_times = [file["timeseries/t/$i"] for i in file_iterations]
    close(file)

    for (n, time) in enumerate(time_series.times)
        file_index = findfirst(t -> t ≈ time, file_times)
        file_iter = file_iterations[file_index]

        field_n = Field(location(time_series), path, name, file_iter,
                        indices = time_series.indices,
                        boundary_conditions = time_series.boundary_conditions,
                        grid = time_series.grid)

        set!(time_series[n], field_n)
    end

    return nothing
end

function set!(fts::FieldTimeSeries, fields_vector::AbstractVector{<:AbstractField})
    raw_data = parent(fts)
    file = jldopen(path)

    for (n, field) in enumerate(fields_vector)
        raw_data[:, :, :, n] .= parent(field)
    end

    close(file)

    return nothing
end

function interior(fts::FieldTimeSeries)
    loc = location(fts)
    topo = topology(fts.grid)
    sz = size(fts.grid)
    halo_sz = halo_size(fts.grid)

    i_interior = interior_parent_indices.(loc, topo, sz, halo_sz)

    indices = fts.indices
    i_view = interior_view_indices.(indices, i_interior)

    return view(parent(fts), i_view..., :)
end

interior(fts::FieldTimeSeries, I...) = view(interior(fts), I...)

indices(fts::FieldTimeSeries) = fts.indices

#####
##### OnDisk time serieses
#####

struct OnDiskData
    path :: String
    name :: String
end

function FieldTimeSeries(path, name, backend::OnDisk; architecture=nothing, grid=nothing)
    file = jldopen(path)

    if isnothing(grid)
        grid = on_architecture(architecture, file["serialized/grid"])
    end

    iterations = parse.(Int, keys(file["timeseries/t"]))
    times = [file["timeseries/t/$i"] for i in iterations]

    data = OnDiskData(path, name)
    LX, LY, LZ = file["timeseries/$name/serialized/location"]
    bcs = file["timeseries/$name/serialized/boundary_conditions"]
    indices = file["timeseries/$name/serialized/indices"]

    close(file)

    return FieldTimeSeries{LX, LY, LZ, OnDisk}(data, grid, bcs, times, indices)
end

#####
##### Methods
#####

# Include the time dimension.
@inline Base.size(fts::FieldTimeSeries) = (size(location(fts), fts.grid, fts.indices)..., length(fts.times))

@propagate_inbounds Base.getindex(f::FieldTimeSeries{LX, LY, LZ, InMemory}, i, j, k, n) where {LX, LY, LZ} = f.data[i, j, k, n]

function Base.getindex(fts::FieldTimeSeries{LX, LY, LZ, OnDisk}, n::Int) where {LX, LY, LZ}
    # Load data
    arch = architecture(fts)
    file = jldopen(fts.data.path)
    iter = keys(file["timeseries/t"])[n]
    raw_data = arch_array(architecture(fts), file["timeseries/$(fts.data.name)/$iter"])
    close(file)

    # Wrap Field
    loc = (LX, LY, LZ)
    field_data = offset_data(raw_data, fts.grid, loc, fts.indices)

    return Field(loc, fts.grid; indices=fts.indices, boundary_conditions=fts.boundary_conditions, data=field_data)
end

Base.setindex!(fts::FieldTimeSeries, val, inds...) = Base.setindex!(fts.data, val, inds...)
Base.parent(fts::FieldTimeSeries{LX, LY, LZ, OnDisk}) where {LX, LY, LZ} = nothing

#####
##### Basic support for reductions
#####
##### TODO: support for reductions across _time_ (ie when 4 ∈ dims)
#####

const FTS = FieldTimeSeries

for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
    reduction! = Symbol(reduction, '!')

    @eval begin

        # Allocating
        function Base.$(reduction)(f::Function, fts::FTS; dims=:, kw...)
            if dims isa Colon
                return Base.$(reduction)($(reduction)(f, fts[n]; kw...) for n in 1:length(fts.times))
            else
                T = filltype(Base.$(reduction!), fts)
                loc = LX, LY, LZ = reduced_location(location(fts); dims)
                times = fts.times
                rts = FieldTimeSeries{LX, LY, LZ}(grid, times, T; indices=fts.indices)
                return Base.$(reduction!)(f, rts, fts; kw...)
            end
        end

        Base.$(reduction)(fts::FTS; kw...) = Base.$(reduction)(identity, fts; kw...)

        function Base.$(reduction!)(f::Function,rts::FTS, fts::FTS; dims=:, kw...)
            dims isa Tuple && 4 ∈ dims && error("Reduction across the time dimension (dim=4) is not yet supported!")
            times = rts.times
            for n = 1:length(times)
                Base.$(reduction!)(f, rts[i], fts[i]; dims, kw...)
            end
            return rts
        end

        Base.$(reduction!)(rts::FTS, fts::FTS; kw...) = Base.$(reduction!)(identity, rts, fts; kw...)
    end
end

#####
##### Show methods
#####

backend_str(::InMemory) = "InMemory"
backend_str(::OnDisk) = "OnDisk"

#####
##### show
#####

function Base.summary(fts::FieldTimeSeries{LX, LY, LZ, K}) where {LX, LY, LZ, K}
    arch = architecture(fts)
    A = typeof(arch)
    return string("$(join(size(fts), "×")) FieldTimeSeries{$(backend_str(K()))} located at ", show_location(fts), " on ", A)
end

function Base.show(io::IO, fts::FieldTimeSeries)
    prefix = string(summary(fts), '\n',
                   "├── grid: ", summary(fts.grid), '\n',
                   "├── indices: ", fts.indices, '\n')

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

    return print(io, prefix, suffix)
end