# lhs = left hand side; rhs = right hand side
# src = source, dst = destination
# order is: lhs, rhs, src, dst
function Arrow(lhs, rhs, src, dst; edge_fun = DirectedEdge, cross = false)
    if cross
        edges = edge_fun(src, dst)
    else
        edges = edge_fun.(src, dst)
    end
    edges = isa(edges, VecOrMat) ? edges : [edges]
    edges = isa(edges, Matrix) ? vec(edges) : edges
    edges = isa(lhs, Arrow) ? vcat(lhs.edges..., edges...) : edges
    edges = isa(rhs, Arrow) ? vcat(edges..., rhs.edges...) : edges
    Arrow(edges, lhs, rhs)
end

→(lhs, rhs) = Arrow(lhs, rhs, keep(lhs, Right, Node), keep(rhs, Left, AbstractNode))
⇒(lhs, rhs) = Arrow(lhs, rhs, keep(lhs, Right, Node), keep(rhs, Left, AbstractNode); cross = true)
←(lhs, rhs) = Arrow(lhs, rhs, keep(rhs, Left, Node), keep(lhs, Right, AbstractNode))
⇐(lhs, rhs) = Arrow(lhs, rhs, keep(rhs, Left, Node), keep(lhs, Right, AbstractNode); cross = true)
↔(lhs, rhs) = Arrow(lhs, rhs, keep(lhs, Right, Node), keep(rhs, Left, AbstractNode); edge_fun = UndirectedEdge)
⇔(lhs, rhs) = Arrow(lhs, rhs, keep(lhs, Right, Node), keep(rhs, Left, AbstractNode); edge_fun = UndirectedEdge, cross = true)

# if either lhs or rhs are a single node, broadcast arrow (barrow) behave like cross arrow (carrow)
barrows = [:→, :←, :↔]
carrows = [:⇒, :⇐, :⇔]
for (barrow, carrow) in zip(barrows,  carrows)
    @eval $barrow(lhs::Union{AbstractNode, Symbol}, rhs::Union{AbstractNode, Symbol}) = $carrow(lhs, rhs)
    @eval $barrow(lhs::Union{AbstractNode, Symbol}, rhs) = $carrow(lhs, rhs)
    @eval $barrow(lhs, rhs::Union{AbstractNode, Symbol}) = $carrow(lhs, rhs)
end


using DataValues
using Test

@testset "DataValueArray: Broadcast" begin

A1 = rand(Int, 10)
M1 = rand(Bool, 10)
n = rand(2:5)
dims = [ rand(2:5) for i in 1:n]
A2 = rand(10, dims...)
M2 = rand(Bool, 10, dims...)
C2 = Array{Float64}(undef, 10, dims...)
i = rand(2:5)
A3 = rand(10, [dims; i]...)
M3 = rand(Bool, 10, [dims; i]...)
C3 = Array{Float64}(undef, 10, [dims; i]...)

m1 = broadcast((x,y)->x, M1, M2)
m2 = broadcast((x,y)->x, M2, M3)
R2 = reshape(Bool[ (x->(m1[x] | M2[x]))(i) for i in 1:length(M2) ], size(M2))
R3 = reshape(Bool[ (x->(m2[x] | M3[x]))(i) for i in 1:length(M3) ], size(M3))
r2 = broadcast((x,y)->x, R2, R3)
Q3 = reshape(Bool[ (x->(r2[x] | R3[x]))(i) for i in 1:length(R3) ], size(M3))

U1 = DataValueArray(A1)
U2 = DataValueArray(A2)
U3 = DataValueArray(A3)
V1 = DataValueArray(A1, M1)
V2 = DataValueArray(A2, M2)
V3 = DataValueArray(A3, M3)
Z2 = DataValueArray{Float64}(10, dims...)
Z3 = DataValueArray{Float64}(10, [dims; i]...)

f() = 5
f(x) = 5 * x
f(x, y) = x * y
f(x, y, z) = x * y * z

for (dests, arrays, nullablearrays, mask) in
    ( ((C2, Z2), (A1, A2), (U1, U2), ()),
      ((C3, Z3), (A2, A3), (U2, U3), ()),
      ((C3, Z3), (A1, A2, A3), (U1, U2, U3), ()),

      ((C2, Z2), (A1, A2), (V1, V2), (R2,)),
      ((C3, Z3), (A2, A3), (V2, V3), (R3,)),
      ((C3, Z3), (A1, A2, A3), (V1, V2, V3), (Q3,)),
)

    # Base.broadcast!(f, B::DataValueArray, As::DataValueArray...)
    broadcast!(f, dests[1], arrays...)
    broadcast!(f, dests[2], nullablearrays...)
    @test isequal(dests[2], DataValueArray(dests[1], mask...))

    # Base.broadcast(f, As::DataValueArray...)
    D = broadcast(f, arrays...)
    X = broadcast(f, nullablearrays...)
    @test isequal(X, DataValueArray(D, mask...))
end

# Base.broadcast!(f, X::DataValueArray)
for (array, nullablearray, mask) in
    ( (A1, U1, ()), (A2, U2, ()), (A3, U3, ()),
      (A1, V1, (M1,)), (A2, V2, (M2,)), (A3, V3, (M3,)),
)
    broadcast!(f, array)
    broadcast!(f, nullablearray)
    @test isequal(nullablearray, DataValueArray(array, mask...))
end

# test broadcasted arithmetic operators
A = rand(10)
X1 = DataValueArray(A)
n = rand(2:5)
dims = rand(2:5, n)
B = rand(Float64, 10, dims...)
X2 = DataValueArray(B)
M = rand(Bool, 10, dims...)
Y = DataValueArray(B, M)

A = rand(Bool, 100)
B = rand(Bool, 100)
M1 = rand(Bool, 100)
M2 = rand(Bool, 100)
X = DataValueArray(A, M1)
Y = DataValueArray(B, M2)
# TODO This NullableArray test doesn't use 3VL
# @test isequal(broadcast(&, X, Y), DataValueArray(A .& B, M1 .| M2))
# @test isequal(broadcast(|, X, Y), DataValueArray(A .| B, M1 .| M2))

end # module


function get_b1_bot(obj::FieldNeQuickParameters)
    return jcall(obj, "getB1Bot", RealFieldElement, ())
end

function get_b1_top(obj::FieldNeQuickParameters)
    return jcall(obj, "getB1Top", RealFieldElement, ())
end

function get_b2_bot(obj::FieldNeQuickParameters)
    return jcall(obj, "getB2Bot", RealFieldElement, ())
end

function get_be_bot(obj::FieldNeQuickParameters)
    return jcall(obj, "getBEBot", RealFieldElement, ())
end

function get_be_top(obj::FieldNeQuickParameters)
    return jcall(obj, "getBETop", RealFieldElement, ())
end

function get_h0(obj::FieldNeQuickParameters)
    return jcall(obj, "getH0", RealFieldElement, ())
end

function get_hm_e(obj::FieldNeQuickParameters)
    return jcall(obj, "getHmE", RealFieldElement, ())
end

function get_hm_f1(obj::FieldNeQuickParameters)
    return jcall(obj, "getHmF1", RealFieldElement, ())
end

function get_hm_f2(obj::FieldNeQuickParameters)
    return jcall(obj, "getHmF2", RealFieldElement, ())
end

function get_layer_amplitudes(obj::FieldNeQuickParameters)
    return jcall(obj, "getLayerAmplitudes", Vector{RealFieldElement}, ())
end

function get_nm_f2(obj::FieldNeQuickParameters)
    return jcall(obj, "getNmF2", RealFieldElement, ())
end


#!/usr/bin/env julia

using PyCall
@pyimport matplotlib as mpl
mpl.rc("text", usetex=true)
mpl.use("Agg")
@pyimport matplotlib.ticker as mti
using PyPlot
using JLD2
using FileIO
using ArgParse
using Printf

o = 1
function singleplot(filename::String)
  d = load(filename*".jld2")
  x = d["x"]
  t = d["t"]
  n = d["n"]
  m = d["m"]
  println(m)
  mpl.rc("font", family="serif", size = 7)
  fig, ax = subplots(figsize = (2.5, 2.))
  col = ["black", "gray", "silver", "gray", "brown", "orange"]
  marker = [":s", ":o", ":v", ":<", ":>", ":d"]
  for i in 1:size(d["cumulants"], 2)
    if occursin("nblocks", filename)  | occursin("maxf", filename)
      str = replace(@sprintf("%.1e", t[i]), "0"=>"")
      tt = "\$"*replace(str, "e+"=>"\\times 10^{")*"}\$"
      ax[:plot](d[x], d["cumulants"][:,i], marker[i], label= "\$ t_{up} \$ = $tt", color = col[i], markersize=2.5, linewidth = 1)
    elseif occursin("nprocs", filename)
      str = replace(@sprintf("%.1e", t[i]), "0"=>"")
      tt = "\$"*replace(str, "e+"=>"\\times 10^{")*"}\$"
      y = d["cumulants"][:,i].\d["cumulants"][1,i]
      ax[:plot](d[x][1:6], y[1:6], marker[i], label= "\$ t_{up} \$ = $tt", color = col[i], markersize=2.5, linewidth = 1)
    else
      tt =  n[i]
      comptimes = d["cumulants"][:,i]./d["cumulants updat"][:,i]
      ax[:plot](d[x], comptimes, marker[i], label= "n = $tt", color = col[i], markersize=2.5, linewidth = 1)
      #i == size(d["cumulants"], 2)? ax[:plot](d[x], d["tm"], ":<", label= "theoretical", color = "black", markersize=2.5, linewidth = 1):()
    end
  end
  subplots_adjust(bottom = 0.16, top=0.92, left = 0.12, right = 0.92)
  fx = matplotlib[:ticker][:ScalarFormatter]()
  fx[:set_powerlimits]((-1, 4))
  ax[:xaxis][:set_major_formatter](fx)
  if occursin("nblocks", filename)
    subplots_adjust(left = 0.15)
    PyPlot.ylabel("computational time [s]", labelpad = 0.6)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:xaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    ax[:legend](fontsize = 5, loc = (m==5) ? 2 : 9, ncol = 1)
  elseif occursin("nprocs", filename)
    subplots_adjust(left = 0.15)
    PyPlot.ylabel("computational speedup", labelpad = 0)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:xaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    ax[:legend](fontsize = 5, loc = 2, ncol = 1)
  elseif occursin("maxf", filename)
    subplots_adjust(left = 0.18)
    PyPlot.ylabel("frequency [Hz]", labelpad = 0.)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:legend](fontsize = 5, loc = 1, ncol = 1)
  else
    PyPlot.ylabel("computational speedup", labelpad = -0.5)
    PyPlot.xlabel("\$ t_{up} \$", labelpad = 0.6)
    subplots_adjust(bottom = 0.16, top=0.92, left = 0.14, right = 0.92)
    if m == 6
      ax[:yaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    end
    ax[:legend](fontsize = 6, loc = 1, ncol = 1)
  end
  f = matplotlib[:ticker][:ScalarFormatter]()
  f[:set_powerlimits]((-3, 2))
  ax[:yaxis][:set_major_formatter](f)
  fig[:savefig]("cumulants"*filename*".eps")
end

function main(args)
  s = ArgParseSettings("description")
  @add_arg_table s begin
    "file"
    help = "the file name"
    arg_type = String
  end
  parsed_args = parse_args(s)
  filename = parsed_args["file"]
  singleplot(replace(filename, ".jld2"=>""))
end

main(ARGS)


## Abstract Markov Chain
## =====================

typealias AbstractChain{T<:Domain} Vector{T}

"Get point samples out of a abstract Markov chain"
point_sample(chain::AbstractChain) = [rand(state) for state in chain]


"""
    isstandard(kr::KRatio)::Boolean

Does this k-ratio have all the necessary basic properties required for use as a standard
(:TakeOffAngle, and :BeamEnergy for both `stdProps` and `unkProps` and :Composition for
`unkProps`.)
"""
function isstandard(std::KRatio)::Bool
    b1 = haskey(std.unkProps, :Composition)
    b1 || @warn "The :Composition property needs to be specified for the `unknown` in $std."
    b2 = haskey(std.unkProps, :BeamEnergy)
    b2 || @warn "The :BeamEnergy property needs to be specified for the `unknown` in $std."
    b3 = haskey(std.unkProps, :TakeOffAngle)
    b3 || @warn "The :TakeOffAngle property needs to be specified for the `unknown` in $std."
    b4 = haskey(std.stdProps, :BeamEnergy)
    b4 || @warn "The :BeamEnergy property needs to be specified for the `standard` in $std."
    b5 = haskey(std.stdProps, :TakeOffAngle)
    b5 || @warn "The :TakeOffAngle property needs to be specified for the `standard` in $std."
    return b1 && b2 && b3 && b4 && b5
end

"""
    matches(kr::Union{KRatio, KRatios}, std::Standard)::Bool

Is `std` a match for `kr`? (Same element, same standard, same lines, same :BeamEnergy & :TakeOffAngle )

stdProps[:Composition] items match if 1) are Materials with almost the same mass-fractions; 2) they are equivalent
AbstractString values.  The later is for
"""
function matches(unk::Union{KRatio, KRatios}, std::KRatio)::Bool
    up, sp = unk.stdProps, std.stdProps
    return element(unk)==element(std) && #  Same element
        isapprox(unk.standard, std.standard, atol=1.0e-5) && # Relative to same material
        issetequal(unk.xrays, std.xrays) && # Same lines
        isapprox(sp[:BeamEnergy], up[:BeamEnergy], rtol=0.001) && # Same E0
        isapprox(sp[:TakeOffAngle], up[:TakeOffAngle], atol=deg2rad(0.1)) # Same take-off angle
end

"""
    standardize(kr::KRatio, std::KRatio)::KRatio
    standardize(kr::KRatios, std::KRatio)::KRatios
    standardize(kratios::Union{AbstractVector{KRatio},AbstractVector{<:KRatios}}, stds::AbstractVector{KRatio})

If the `std::KRatio` is a suitable match for `kr` then `kr` is restandardized using `std`.  Otherwise, the original
`KRatio` or `KRatios` is returned.
"""
function standardize(kr::KRatio, std::KRatio)::KRatio
    div(n::AbstractFloat, d::AbstractFloat) = value(n)/value(d)
    div(n::UncertainValue, d::UncertainValue) = divide(n,d,0.0)
    @assert isstandard(std) "$std does not contain all the necessary properties (see details in the above warnings)."
    return if matches(kr, std)
        newkr = div(kr.kratio, std.kratio)
        KRatio(kr.xrays, kr.unkProps, std.unkProps, std.unkProps[:Composition], newkr)
    else
        kr
    end
end
function standardize(krs::KRatios, std::KRatio)::KRatios
    @assert isstandard(std) "$std does not contain all the necessary properties (see details in the above warnings)."
    div(n::AbstractFloat, d::AbstractFloat) = value(n)/value(d)
    div(n::UncertainValue, d::UncertainValue) = divide(n, d, 0.0)
    return if matches(krs, std)
        newkrs = map(krs.kratios) do kr
            div(kr, std.kratio)
        end
        KRatios(krs.xrays, krs.unkProps, std.unkProps, std.unkProps[:Composition], newkrs)
    else
        krs
    end
end
function standardize(kratios::Union{AbstractVector{KRatio},AbstractVector{<:KRatios}}, stds::AbstractVector{KRatio})
    return map(kratios) do kr
        i = findfirst(std->matches(kr,std), stds)
        isnothing(i) ? kr : standardize(kr, stds[i])
    end
end


module VisualiseTests

using Test
import MetaGraphs

using PicoQuant

@testset "Visualisation tests" begin
    @testset "Test contracting a network" begin

        qasm_str = """OPENQASM 2.0;
                    include "qelib1.inc";
                    qreg q[3];
                    h q[0];
                    cx q[0],q[1];
                    cx q[1],q[2];"""

        circ = load_qasm_as_circuit(qasm_str)
        tng = convert_qiskit_circ_to_network(circ, InteractiveBackend())
        add_input!(tng, "000")
        add_output!(tng, "000")
        g = PicoQuant.create_graph(tng)

        # Is there only one node left after the contraction?
        @test begin
            MetaGraphs.nv(g) == 9
        end

        @test begin
            MetaGraphs.ne(g) == 8
        end
    end
end
end


using KernelAbstractions: NoneEvent

using Oceananigans.Utils: arch_array
using Oceananigans.Grids: AbstractGrid

abstract type AbstractTimeDiscretization end

"""
    struct ExplicitTimeDiscretization <: AbstractTimeDiscretization

Represents fully-explicit time-discretization of a `TurbulenceClosure`.
"""
struct ExplicitTimeDiscretization <: AbstractTimeDiscretization end

Base.summary(::ExplicitTimeDiscretization) = "ExplicitTimeDiscretization"

"""
    struct VerticallyImplicitTimeDiscretization <: AbstractTimeDiscretization

Represents vertically-implicit time-discretization of a `TurbulenceClosure`.

This imples that a flux divergence such as ``∇ ⋅ q`` at the n-th timestep is
time-discretized as

```julia
[∇ ⋅ q]ⁿ = [explicit_flux_divergence]ⁿ + [∂z (κ ∂z c)]ⁿ⁺¹
```
"""
struct VerticallyImplicitTimeDiscretization <: AbstractTimeDiscretization end

Base.summary(::VerticallyImplicitTimeDiscretization) = "VerticallyImplicitTimeDiscretization"

@inline time_discretization(::AbstractTurbulenceClosure{TimeDiscretization}) where TimeDiscretization = TimeDiscretization()
@inline time_discretization(::Nothing) = ExplicitTimeDiscretization() # placeholder for closure::Nothing

#####
##### Explicit: move along, nothing to worry about here (use fallbacks).
#####

const ATD = AbstractTimeDiscretization

@inline diffusive_flux_x(i, j, k, grid, ::ATD, args...) = diffusive_flux_x(i, j, k, grid, args...)
@inline diffusive_flux_y(i, j, k, grid, ::ATD, args...) = diffusive_flux_y(i, j, k, grid, args...)
@inline diffusive_flux_z(i, j, k, grid, ::ATD, args...) = diffusive_flux_z(i, j, k, grid, args...)

@inline viscous_flux_ux(i, j, k, grid, ::ATD, args...) = viscous_flux_ux(i, j, k, grid, args...)
@inline viscous_flux_uy(i, j, k, grid, ::ATD, args...) = viscous_flux_uy(i, j, k, grid, args...)
@inline viscous_flux_uz(i, j, k, grid, ::ATD, args...) = viscous_flux_uz(i, j, k, grid, args...)

@inline viscous_flux_vx(i, j, k, grid, ::ATD, args...) = viscous_flux_vx(i, j, k, grid, args...)
@inline viscous_flux_vy(i, j, k, grid, ::ATD, args...) = viscous_flux_vy(i, j, k, grid, args...)
@inline viscous_flux_vz(i, j, k, grid, ::ATD, args...) = viscous_flux_vz(i, j, k, grid, args...)

@inline viscous_flux_wx(i, j, k, grid, ::ATD, args...) = viscous_flux_wx(i, j, k, grid, args...)
@inline viscous_flux_wy(i, j, k, grid, ::ATD, args...) = viscous_flux_wy(i, j, k, grid, args...)
@inline viscous_flux_wz(i, j, k, grid, ::ATD, args...) = viscous_flux_wz(i, j, k, grid, args...)


################################################################################
# Maptables
################################################################################

"""
TODO

API
---
* [`location_type`](@ref)
* [`getlocations`](@ref)
* [`isvalid`](@ref)

Implementations
---------------
* [`MapTable`](@ref)
"""
abstract type AbstractMapTable end

#######
# API #
#######

"""
Return the stored location type for a `MapTable`.

Method List
-----------
$(METHODLIST)
"""
function location_type end

"""
    getlocations(maptable, class::Int)

Return a vector of locations that nodes of type `class` can occupy.

Method List
-----------
$(METHODLIST)
"""
function getlocations end

"""
    isvalid(maptable, class, address :: Address)

Return `true` if nodes of type `class` can occupy `address`. In other words,
there is some component at `adddress` that class can be mapped to.

Method List
-----------
$(METHODLIST)
"""
function isvalid end

################################################################################
# MapTable
################################################################################

"""
Default implementation of [`AbstractMapTable`](@ref)

Parameters:
* `D` - The dimensionality of the `Addresses` in the table.
"""
struct MapTable{D} <: AbstractMapTable
    """
    Bit mask of whether a node class may be mapped to an address.

    Accessing strategy: Look up the node class to index the outer vector. Index the inner
    array with an address.

    A `true` entry means the task class can be mapped.
    A `false` entry means the task calss cannot be mapped.
    """
    mask::Vector{Array{Bool,D}}
end

function MapTable(
    toplevel::TopLevel, ruleset::RuleSet, equivalence_classes, pathtable::PathTable{D}
) where {D}

    # For each node class C, create an array the same size as the
    # pathtable. For each address A, record the indices of the paths at
    # componant_table[A] that C can be mapped to.
    mask = map(equivalence_classes.reps) do node
        map(pathtable) do path
            return path !== nothing && canmap(ruleset, node, toplevel[path])
        end
    end
    return MapTable(mask)
end
location_type(::MapTable{D}) where {D} = Address{D}

# Makes repeated calls to the 3 argument version of getlocations starting at
# all base addresses to generate ALL locations for this class.
function getlocations(maptable::MapTable{D}, class::Int) where {D}
    # Allocation an empty vector of the appropriate location type.
    addresses = location_type(maptable)[]
    table = maptable.mask[class]

    for address in CartesianIndices(table)
        isvalid(maptable, class, address) && push!(addresses, address)
    end
    return addresses
end

function isvalid(maptable::MapTable, class::Integer, address::Address)
    return maptable.mask[class][address]
end


using CxxWrap

module JERI
  using CxxWrap

  @wrapmodule joinpath(@__DIR__,"../libjeri.so") :define_jeri

  function __init__()
    @initcxx
  end

  export initialize, finalize
  export Atom, create_atom
  export Shell, create_shell
  export BasisSet
  export ShellPair, precompute_shell_pair_data

  export OEIEngine
  export compute_overlap_block, compute_overlap_grad_block
  export compute_kinetic_block, compute_kinetic_grad_block
  export compute_nuc_attr_block, compute_nuc_attr_grad_block

  export PropEngine, compute_dipole_block

  export TEIEngine, compute_eri_block
end


module LatBoTests
using LatBo
using FactCheck: facts, context, @fact, not, roughly, exitstatus, exactly, @fact_throws

include("indices/cached.jl")
# include("units.jl")
include("playground.jl")
include("geometry.jl")
include("indices/indices.jl")
include("lb/lattice.jl")
include("lb/thermodynamics.jl")
include("lb/collision.jl")
include("lb/streaming.jl")
include("lb/iolet.jl")
include("lb/initialization.jl")
include("lb/local_kernel.jl")
exitstatus()
end


using Dates

"""
    optimize!

Optimizes a local search index for an specific algorithm to get the desired performance.
Note that optimizing for low-recall will yield to faster searches; the train queries
are specified as part of the `perf` struct.
"""
function optimize!(algosearch::LocalSearchAlgorithm,
                   index::SearchGraph{T},
                   dist::Function,
                   recall::Float64,
                   perf::Performance;
                   bsize::Int=4,
                   tol::Float64=0.01) where T
    n = length(index.db)
    score_function(p) = p.recall < recall ? p.recall : 1.0 + n / p.evaluations
    #score_function(p) = p.recall < recall ? p.recall : 1.0 + 1.0 / (1.0 + sum(p.distances))
    #score_function(p) = 1.0 / (1.0 + sum(p.distances))
    optimize_neighborhood!(index.neighborhood_algo, index, dist, perf, recall)
    p = probe(perf, index, dist)
    best_list = [(score=score_function(p), state=algosearch, perf=p)]
    exploration = Dict(algosearch => 0)  ## -1 unexplored; 0 visited; 1 visited & expanded

    @info "==== BEGIN $(typeof(algosearch)). Starting parameter optimization; expected recall: $recall, n: $n"
    prev_score = -1.0
    iter = 0
    while abs(best_list[1].score - prev_score) > tol
        iter += 1
        prev_score = best_list[1].score
        @info "  == begin iteration $(typeof(algosearch)). Iteration: $iter, expected recall: $recall, n: $n"

        for prev in @view best_list[1:end]  ## the view also fixes the size of best_list even after push!
            S = get(exploration, prev.state, -1)
            if S == 1
                continue  # visited and explored
            elseif S == 0
                exploration[prev.state] = 1
            end

            opt_expand_neighborhood(prev.state, n, iter) do state
                S = get(exploration, state, -1)
                if S == -1
                    exploration[state] = 0
                    index.search_algo = state
                    # p = probe(perf, index, repeat=3, aggregation=:median, field=:seconds)
                    p = probe(perf, index, dist, repeat=1, field=:seconds)
                    score = score_function(p)
                    push!(best_list, (score=score, state=state, perf=p))
                    if score > best_list[1].score
                        @info "  ** $(typeof(algosearch)). A new best conf was found> score: $score, conf: $(JSON.json(state)), perf: $(JSON.json(p)), best_list's length: $(length(best_list)), n: $(n)"
                    end
                end
            end
        end

        sort!(best_list, by=(x) -> -x.score)
        if length(best_list) > bsize
            best_list = best_list[1:bsize]
        end
        @info "  == end $(typeof(algosearch)). Iteration finished; $(JSON.json(best_list[1])), beam: $(length(best_list)), n: $(n)"
    end

    index.search_algo = best_list[1].state
    @info "==== END $(typeof(algosearch)). Finished, best: $(JSON.json(best_list[1])), n: $(n)"
    return index
end


# file for homotopy functions
# TODO: add mechanism for selecting which homotopy function to use
# TODO: better idea: move the homotopy function to NLSolver, because it is
#                    physics agnostic

import PDESolver.evalHomotopy

"""
  This function calls the appropriate homotopy function for the Euler module.
"""
function evalHomotopy(mesh::AbstractMesh, sbp::AbstractOperator,
                      eqn::EulerData{Tsol, Tres},
                      opts::Dict, res::Abstract3DArray, t = 0.0) where {Tsol, Tres}

  fill!(res, 0)
  if opts["homotopy_function"] == "ViscousBO"
    sensor = ShockSensorBO{Tsol, Tres}(mesh, sbp, opts)
    capture = eqn.shock_capturing
    res_orig = eqn.res
    eqn.res = res
    applyShockCapturing(mesh, sbp, eqn, opts, sensor, capture)
    eqn.res = res_orig
  elseif opts["homotopy_function"] == "FirstOrderDissipation"

    calcHomotopyDiss(mesh, sbp, eqn, opts, res)
  else
    error("unrecognized homotop_function: $(opts["homotopy_function"])")
  end

  return nothing
end

"""
  Calculate a first order accurate dissipation to use as a homotopy function

  Inputs:
    mesh: a DG mesh
    sbp: an SBP operator
    eqn: an EulerData object
    opts: options dictionary

  Inputs/Outputs:
    res: 3D array to store the homotopy function in

  Note eqn.res is *not* modified by this function.

  Aliasing restrictions: none
"""
function calcHomotopyDiss(mesh::AbstractDGMesh{Tmsh}, sbp,
        eqn::EulerData{Tsol, Tres}, opts,
        res::Abstract3DArray{Tres}) where {Tsol, Tres, Tmsh}

#  println("\nentered calcHomotopyDiss")

  # some checks for when parallelism is enabled
  @assert opts["parallel_data"] == PARALLEL_DATA_ELEMENT
  for i=1:mesh.npeers
    @assert eqn.shared_data[i].recv_waited
  end

  fill!(res, 0.0)

  #----------------------------------------------------------------------------
  # volume dissipation

  t1 = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerElement)
  t2 = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement)
  fill!(res, 0.0)
  nrm2 = zeros(Tmsh, mesh.dim)

  # compute the D operator in each direction
  D = zeros(mesh.numNodesPerElement, mesh.numNodesPerElement, mesh.dim)
  for d=1:mesh.dim
    D[:, :, d] = inv(diagm(sbp.w))*sbp.Q[:, :, d]
  end

  for el=1:mesh.numEl
    q_el = sview(eqn.q, :, :, el)
    res_el = sview(res, :, :, el)
    for d1=1:mesh.dim
      fill!(t1, 0.0)
      fill!(t2, 0.0)
      differentiateElement!(sbp, d1, q_el, t1)

      for j=1:mesh.numNodesPerElement
        q_j = sview(eqn.q, :, j, el)
        for k=1:mesh.dim
          nrm2[k] = mesh.dxidx[d1, k, j, el]
        end

        lambda_max = getLambdaMax(eqn.params, q_j, nrm2)
        for k=1:mesh.numDofPerNode
          t2[k, j] = lambda_max*t1[k, j]
        end

      end  # end loop j
     weakDifferentiateElement!(sbp, d1, t2, res_el, SummationByParts.Add(), true)
    end  # end loop dim
  end  # end loop el

  #----------------------------------------------------------------------------
  # interface dissipation

  # interpolate to face rather than using eqn.q_face, in case it hasn't been
  # updated since eqn.q was updated

  q_faceL = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerFace)
  q_faceR = zeros(q_faceL)
#  nrm2 = eqn.params.nrm2
  flux = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerFace)
  for i=1:mesh.numInterfaces
    iface_i = mesh.interfaces[i]
    qL = sview(eqn.q, :, :, iface_i.elementL)
    qR = sview(eqn.q, :, :, iface_i.elementR)
    resL = sview(res, :, :, iface_i.elementL)
    resR = sview(res, :, :, iface_i.elementR)
#    fill!(q_faceL, 0.0)
#    fill!(q_faceR, 0.0)

    interiorFaceInterpolate!(mesh.sbpface, iface_i, qL, qR, q_faceL, q_faceR)

    # calculate numerical flux at each face node
    for j=1:mesh.numNodesPerFace
      qL_j = sview(q_faceL, :, j)
      qR_j = sview(q_faceR, :, j)

      # get the face normal
      nrm2 = sview(mesh.nrm_face, :, j, i)

      lambda_max = getLambdaMaxSimple(eqn.params, qL_j, qR_j, nrm2)

      for k=1:mesh.numDofPerNode
        flux[k, j] = 0.5*lambda_max*(qL_j[k] - qR_j[k])
      end
    end  # end loop j

    # integrate over the face
    interiorFaceIntegrate!(mesh.sbpface, iface_i, flux, resL, resR)
  end  # end loop i

# the boundary term makes the predictor-corrector algorithm converge slower
  #----------------------------------------------------------------------------
  # boundary dissipation
  # use q_faceL, nrm2, flux  from interface dissipation
  if opts["homotopy_addBoundaryIntegrals"]
    qg = zeros(Tsol, mesh.numDofPerNode)  # boundary state
    for i=1:mesh.numBoundaryFaces
      bndry_i = mesh.bndryfaces[i]
      qL = sview(eqn.q, :, :, bndry_i.element)
      resL = sview(res, :, :, bndry_i.element)
      fill!(q_faceL, 0.0)

      boundaryFaceInterpolate!(mesh.sbpface, bndry_i.face, qL, q_faceL)

  #    q_faceL = sview(eqn.q_bndry, :, :, i)
      for j=1:mesh.numNodesPerFace
        q_j = sview(q_faceL, :, j)
  #      dxidx_j = sview(mesh.dxidx_bndry, :, :, j, i)

        # calculate boundary state
        coords = sview(mesh.coords_bndry, :, j, i)
        calcFreeStream(eqn.params, coords, qg)
#        calcInvChannelIC(eqn.params, coords, qg)

        # calculate face normal
        nrm2 = sview(mesh.nrm_bndry, :, j, i)

        # calculate lambda_max
        lambda_max = getLambdaMaxSimple(eqn.params, q_j, qg, nrm2)

        # calculate dissipation
        for k=1:mesh.numDofPerNode
          flux[k, j] = 0.5*lambda_max*(q_j[k] - qg[k])
        end

      end  # end loop j

      # integrate over the face
      boundaryFaceIntegrate!(mesh.sbpface, bndry_i.face, flux, resL)
    end  # end loop i
  end  # end if addBoundaryIntegrals

  #----------------------------------------------------------------------------

  # shared face integrals
  # use q_faceL, q_faceR, flux from above

  workarr = zeros(q_faceR)
  for peer=1:mesh.npeers
    # get data for this peer
    bndries_local = mesh.bndries_local[peer]
    interfaces_peer = mesh.shared_interfaces[peer]

    qR_peer = eqn.shared_data[peer].q_recv
    nrm_peer = mesh.nrm_sharedface[peer]
    start_elnum = mesh.shared_element_offsets[peer]

    for i=1:length(bndries_local)
      bndry_i = bndries_local[i]
      iface_i = interfaces_peer[i]
      qL_i = sview(eqn.q, :, :, bndry_i.element)
      qR_i = sview(qR_peer, :, :, iface_i.elementR - start_elnum + 1)
      resL = sview(res, :, :, bndry_i.element)

      # interpolate to face
      interiorFaceInterpolate!(mesh.sbpface, iface_i, qL_i, qR_i, q_faceL, q_faceR)
      # compute flux at every face node
      for j=1:mesh.numNodesPerFace
        qL_j = sview(q_faceL, :, j)
        qR_j = sview(q_faceR, :, j)
        nrm2 = sview(nrm_peer, :, j, i)

        # get max wave speed
        lambda_max = getLambdaMaxSimple(eqn.params, qL_j, qR_j, nrm2)

        # calculate flux
        for k=1:mesh.numDofPerNode
          flux[k, j] = 0.5*lambda_max*(qL_j[k] - qR_j[k])
        end
      end  # end loop j

      # integrate over the face
      boundaryFaceIntegrate!(mesh.sbpface, bndry_i.face, flux, resL)
    end  # end loop i
  end  # end loop peer

  # negate for consistency with the physics module
  for i=1:length(res)
    res[i] = -res[i]
  end

#  println("homotopy residual norm = ", norm(vec(res)))

  return nothing
end


using GeoData, DimensionalData, Test, Statistics, Dates
using GeoData.LookupArrays, GeoData.Dimensions

data1 = cumsum(cumsum(ones(10, 11); dims=1); dims=2)
data2 = 2cumsum(cumsum(ones(10, 11, 1); dims=1); dims=2)
dims1 = X(10.0:10.0:100.0), Y(-50.0:10.0:50.0)
dims2 = (dims1..., Ti([DateTime(2019)]))
refdimz = ()
nme = :test
mval = -9999.0
meta = NoMetadata()

# Formatting only occurs in shorthand constructors
ga1 = GeoArray(data1, dims1; refdims=refdimz, name=nme, metadata=meta, missingval=mval)
ga2 = GeoArray(data2, dims2)

@testset "stack constructors" begin
    st1 = GeoStack((ga1, ga2); name=(:ga1, :ga2))
    st2 = GeoStack((ga1, ga2); keys=(:ga1, :ga2))
    st3 = GeoStack((data1, data2), dims2; keys=(:ga1, :ga2))
    st4 = GeoStack(st3)
    st5 = GeoStack(st3, dims2)
    @test st1 == st2 == st3 == st4 == st5

    # The dimension differences are lost because the table
    # is tidy - every column is the same length
    table_st = GeoStack(DimTable(st3), dims2)
    @test dims(table_st[:ga1]) isa Tuple{<:X,<:Y,<:Ti}
end

st = GeoStack((ga1, ga2); name=(:ga1, :ga2))

@testset "stack layers" begin
    @test length(st) == 2
    @test first(st) == ga1
    @test last(st) == ga2
    @test DimensionalData.layers(st) isa NamedTuple
    @test st[:ga1] == ga1
    @test st[:ga2] == ga2
    @test parent(st[:ga1]) == data1
    @test parent(st[:ga1]) isa Array{Float64,2}
    @test keys(st) == (:ga1, :ga2)
    @test haskey(st, :ga1)
    @test names(st) == (:ga1, :ga2)
    @test collect(values(st)) == [ga1, ga2]
    # We can iterate/splat stacks as GeArrays
    @test st == GeoStack(st...)
end

@testset "st fields " begin
    @test DimensionalData.layerdims(st, :ga1) == DimensionalData.format(dims1, data1)
    @test metadata(st) == NoMetadata()
    @test metadata(st, :ga1) == NoMetadata()
end

@testset "indexing" begin
    # Indexing the st is the same as indexing its child array
    a = st[:ga1][X(2:4), Y(5:6)]
    @test a == st[:ga1, X(2:4), Y(5:6)]

    @inferred st[:ga1][X=2:4, Y=5:6]
    # FIXME: This isn't inferred, the constants don't propagate like they
    # do in the above call. Probably due to the anonymous wrapper function.
    @test_broken @inferred st[:ga1, X(2:4), Y(5:6)] isa GeoArray

    # Getindex for a whole st of new GeoArrays
    a = st[X=2:4, Y=5:6]
    @test a isa GeoStack
    @test a[:ga1] isa GeoArray
    @test parent(a[:ga1]) isa Array
    @test a[:ga1] == data1[2:4, 5:6]
    @test a[:ga2] == data2[2:4, 5:6, 1:1]

    @testset "select new arrays for the whole st" begin
        s = st[Y=Between(-10, 10.0), Ti=At(DateTime(2019))]
        @test s isa GeoStack
        @test s[:ga1] isa GeoArray
        @test parent(s[:ga1]) isa Array
        @test s[:ga1] == data1[:, 5:7]
        @test s[:ga2] == data2[:, 5:7, 1]
        @test dims(s[:ga2]) == (X(Sampled(10.0:10.0:100.0, ForwardOrdered(), Regular(10.0), Points(), NoMetadata())),
                                Y(Sampled(-10.0:10.0:10.0, ForwardOrdered(), Regular(10.0), Points(), NoMetadata())))
        @test refdims(s[:ga2]) == (Ti(Sampled([DateTime(2019)], ForwardOrdered(), Irregular(nothing, nothing), Points(), NoMetadata())),)
        @test ismissing(missingval(s, :ga2)) && ismissing(missingval(s[:ga2]))
    end

    @testset "select views of arrays for the whole st" begin
        sv = view(st, Y=Between(-4.0, 27.0), Ti=At(DateTime(2019)))
        @test sv isa GeoStack
        @test sv[:ga1] isa GeoArray
        @test parent(sv[:ga1]) isa SubArray
        @test sv[:ga1] == data1[:, 6:8]
        @test sv[:ga2] == data2[:, 6:8, 1]
        @test dims(sv[:ga2]) == (X(Sampled(10.0:10:100.0, ForwardOrdered(), Regular(10.0), Points(), NoMetadata())),
                                 Y(Sampled(0.0:10:20.0, ForwardOrdered(), Regular(10.0), Points(), NoMetadata())))
        @test refdims(sv[:ga2])[1] == Ti(Sampled(view([DateTime(2019)], 1:1), ForwardOrdered(), Irregular((nothing, nothing)), Points(), NoMetadata()))
        # Stack of view-based GeoArrays
        v = view(st, X(2:4), Y(5:6))
        # @inferred view(st, X(2:4), Y(5:6))
        @test v isa GeoStack
        @test v[:ga1] isa GeoArray
        @test parent(v[:ga1]) isa SubArray
        @test v[:ga1] == view(data1, 2:4, 5:6)
        @test v[:ga2] == view(data2, 2:4, 5:6, 1:1)
    end

end

@testset "subset st with specific key(s)" begin
    s1 = GeoStack(st; keys=(:ga2,))
    @test keys(s1) == (:ga2,)
    @test length(values(s1)) == 1
    s2 = GeoStack(st; keys=(:ga1, :ga2))
    @test keys(s2) == (:ga1, :ga2)
    @test length(values(s2)) == 2
end

@testset "concatenate stacks" begin
    dims1b = X(110:10:200), Y(-50:10:50)
    dims2b = (dims1b..., Ti([DateTime(2019)]))
    stack_a = GeoStack((ga1=ga1, ga2=ga2))
    stack_b = GeoStack((ga1=GeoArray(data1 .+ 10, dims1b), ga2=GeoArray(data2 .+ 20, dims2b)))
    catstack = cat(stack_a, stack_b; dims=X())
    @test size(first(catstack)) == (20, 11)
    @test val(dims(catstack, X)) ≈ 10.0:10.0:200.0
    #@test step(dims(first(catstack), X())) == 10.0
    @test DimensionalData.bounds(dims(first(catstack), X)) == (10.0, 200.0)
    @test catstack[:ga1][Y(1)] == 1.0:20.0
    @test catstack[:ga2][Y(1), Ti(1)] == 2.0:2.0:40.0
    catstack = cat(stack_a, stack_b; dims=(X(), Y()))
    @test size(first(catstack)) == (20, 22)
end

@testset "copy" begin
    cp = copy(st)
    @test all(st[:ga1] .=== cp[:ga1])
    @test st[:ga1] !== cp[:ga1]
end

@testset "show" begin
    sh = sprint(show, st)
    # Test but don't lock this down too much
    @test occursin("GeoStack", sh)
    @test occursin("Y", sh)
    @test occursin("X", sh)
end


include("panX.jl")

Export = Command(
   "export",
   "pangraph export <options> [arguments]",
   "exports a pangraph to a chosen file format(s)",
   """zero or one pangraph file (native json)
      if no file given, reads from stdin
      stream can be optionally gzipped.""",
   [
    Arg(
        Int,
        "minimum block length for edge connection",
        (short="-ell", long="--edge-minimum-length"),
        "blocks below this length cutoff will be ignored for edges in graph",
        200,
    ),
    Arg(
        Int,
        "maximum block length for edge connection",
        (short="-elu", long="--edge-maximum-length"),
        "blocks above this length cutoff will be ignored for edges in graph",
        typemax(Int),
    ),
    Arg(
        Int,
        "minimum block depth for edge connection",
        (short="-edl", long="--edge-mininum-depth"),
        "blocks below this depth cutoff will be ignored for edges in graph",
        0,
    ),
    Arg(
        Int,
        "maximum block depth for edge connection",
        (short="-edu", long="--edge-maximum-depth"),
        "blocks above this depth cutoff will be ignored for edges in graph",
        typemax(Int),
    ),
    Arg(
        Int,
        "minimum block length for export",
        (short="-ll", long="--minimum-length"),
        "blocks below this length cutoff will not be exported",
        0,
    ),
    Arg(
        Int,
        "maximum block length for export",
        (short="-lu", long="--maximum-length"),
        "blocks above this length cutoff will not be exported",
        typemax(Int),
    ),
    Arg(
        Int,
        "minimum block depth for export",
        (short="-dl", long="--mininum-depth"),
        "blocks below this depth cutoff will not be exported",
        0,
    ),
    Arg(
        Int,
        "maximum block depth for export",
        (short="-du", long="--maximum-depth"),
        "blocks above this depth cutoff will not be exported",
        typemax(Int),
    ),
    Arg(
        String,
        "output directory",
        (short="-o", long="--output-directory"),
        "relative path to directory where output will be stored",
        "export"
    ),
    Arg(
        String,
        "prefix",
        (short="-p", long="--prefix"),
        "basename of files",
        "pangraph"
    ),
    Arg(
        Bool,
        "no GFA export",
        (short="-ng", long="--no-export-gfa"),
        "do not emit GFA file",
        false,
    ),
    Arg(
        Bool,
        "export panX visualization",
        (short="-pX", long="--export-panX"),
        "emit vis directory to input to panX-visualization",
        false,
    ),
    Arg(
        Bool,
        "no duplications in export",
        (short="-nd", long="--no-duplications"),
        "do not export any block that contains at least one strain more than once",
        false
    ),
   ],

   function(args)
       path = parse(Export, args)
       path = if (path === nothing || length(path) == 0)
           nothing
       elseif length(path) == 1
           path
       else
           return 2
       end

       graph = load(path, Export)

       function cutoffs(prefix)
           (
              depth = (
                   min = arg(Export, "-$(prefix)dl"),
                   max = arg(Export, "-$(prefix)du"),
              ),
              length = (
                   min = arg(Export, "-$(prefix)ll"),
                   max = arg(Export, "-$(prefix)lu"),
              )
           )
       end
       noduplications = arg(Export, "-nd")

       node = cutoffs("")
       edge = cutoffs("e")

       filter = (
            connect = function(node)
                return (!(edge.depth.min  ≤ Blocks.depth(node.block)  ≤ edge.depth.max)
                     || !(edge.length.min ≤ Blocks.length(node.block) ≤ edge.length.max))
            end,
            output = function(segment, isolates)
                return (!(node.depth.min  ≤ segment.depth  ≤ node.depth.max)
                     || !(node.length.min ≤ length(segment.sequence) ≤ node.length.max)
                     ||  (noduplications && (length(isolates) != length(unique(isolates)))))
            end
       )

       # Create directory if it doesn't exist
       directory = arg(Export, "-o")
       if !isdir(directory)
           mkpath(directory)
       end
       prefix = arg(Export, "-p")

       # GFA export (default)
       if !arg(Export, "-ng")
           Base.open("$(directory)/$(prefix).gfa", "w") do io
               marshal(io, graph; fmt=:gfa, opt=filter)
           end
           Base.open("$(directory)/$(prefix).fa", "w") do io
               marshal(io, graph; fmt=:fa)
           end
       end

       # panX export (doesn't fit into marshal paradigm)
       if arg(Export, "-pX")
           if !Shell.havecommand("fasttree")
               panic("external command fasttree not found. please install before attempting to export as panX visualization\n")
           end

           vis = "$(directory)/vis"
           if !isdir(vis)
               mkpath(vis)
           end

           PanX.emit(graph, vis)
       end

        # TODO: others?
   end
)


# Constraints parameters
const BENCHED_CONSTRAINTS = deepcopy(usual_constraints)

const DEFAULT_CONCEPTS = [(:all_different, nothing)]
const DEFAULT_LANGUAGES = [:Julia] # [:Julia, :C, :CPP]
const DEFAULT_METRICS = [:hamming, :manhattan]
const DEFAULT_COMPLETE_SEARCH_LIMIT = 5^5# 6^6 # could be improved with multithreading
const DEFAULT_PARTIAL_SEARCH_LIMIT = 4^4
const DEFAULT_SAMPLINGS = [1000, 100]

const MINIMUM_DOMAIN_SIZE = 3
const MAXIMUM_DOMAIN_SIZE = 8 # NOTE - Check if it is too high

function domains_sizes(;
    min_domain_size=MINIMUM_DOMAIN_SIZE, max_domain_size=MAXIMUM_DOMAIN_SIZE
)
    return reverse(collect(min_domain_size:max_domain_size))
end

# Genetic Algorithms default parameters
const MAXIMUM_TOTAL_ITERATIONS = 12 # order => 2^15 = 32768 // lower this param
const MINIMUM_ICN_ITERATIONS = 5 # order => 2^5= 32
const MINIMUM_GENERATIONS = 5 # order => 2^5 = 32

function icn_iterations(;
    min_icn_iterations=MINIMUM_ICN_ITERATIONS,
    min_generations=MINIMUM_GENERATIONS,
    max_total_iterations=MAXIMUM_TOTAL_ITERATIONS,
)
    maximum_icn_iterations = max_total_iterations - min_generations
    return reverse([2^i for i in min_icn_iterations:maximum_icn_iterations])
end

function generations(;
    min_icn_iterations=MINIMUM_ICN_ITERATIONS,
    min_generations=MINIMUM_GENERATIONS,
    max_total_iterations=MAXIMUM_TOTAL_ITERATIONS,
)
    maximum_generations = max_total_iterations - min_icn_iterations
    return reverse([2^i for i in min_generations:maximum_generations])
end

const DEFAULT_POPULATIONS = reverse([2^i for i in 6:8]) # 2^5=32 to 2^10 = 1024

const DEFAULT_LOSS_SAMPLING_THRESHOLD = 2^10

root(n) = x -> round(Int, x^(1 / n))
constant(n) = _ -> n

const DEFAULT_LOSS_SAMPLER = [
    nothing, root(2), root(3), constant(DEFAULT_LOSS_SAMPLING_THRESHOLD)
]

const DEFAULT_MEMOIZATION = [true, false]

# All parameters
const ALL_PARAMETERS = Dict(
    # Search parameters
    :concept => DEFAULT_CONCEPTS,
    :complete_search_limit => DEFAULT_COMPLETE_SEARCH_LIMIT,
    :domains_size => domains_sizes(),
    :partial_search_limit => DEFAULT_PARTIAL_SEARCH_LIMIT,
    :sampling => DEFAULT_SAMPLINGS,
    :search => [:partial, :complete], # :flexible is also an option

    # # Learning parameter
    :metric => DEFAULT_METRICS,
    :generations => [32,64,128], #generations(),
    :icn_iterations => [32,64,128], #icn_iterations(),
    :language => DEFAULT_LANGUAGES,
    :population => DEFAULT_POPULATIONS,
    :loss_sampler => nothing, #DEFAULT_LOSS_SAMPLER,
    :loss_sampling_threshold => Inf, #DEFAULT_LOSS_SAMPLING_THRESHOLD,
    :memoize => DEFAULT_MEMOIZATION,
)

# ALL_PARAMETERS[:search] = [:flexible]
# ALL_PARAMETERS[:domains_size] = 8

"""
to run on khromeleque:
cd ~/.julia/dev/ICNBenchmarks/scripts
../../../../julia -t auto main.jl no_overlap "ALL_PARAMETERS[:generations] = [32,64,128]" "ALL_PARAMETERS[:icn_iterations] = [32,64,128]" "ALL_PARAMETERS[:loss_sampler] = nothing" "ALL_PARAMETERS[:loss_sampling_threshold] = Inf" "ALL_PARAMETERS[:population] = [64,128,256]" "ALL_PARAMETERS[:population] = [64,128,256]" "ALL_PARAMETERS[:complete_search_limit] = 256"
"""
ALL_PARAMETERS[:generations] = [32,64,128]
ALL_PARAMETERS[:icn_iterations] = [32,64,128]
ALL_PARAMETERS[:loss_sampler] = nothing
ALL_PARAMETERS[:loss_sampling_threshold] = Inf
ALL_PARAMETERS[:population] = [64,128,256]


##############################################################################
##
## FixedEffect
##
## The categoricalarray may have pools that are never referred. Note that the pool does not appear in FixedEffect anyway.
##
##############################################################################

struct FixedEffect{R <: AbstractVector{<:Integer}, I <: AbstractVector{<: Real}}
	refs::R                 # refs must be between 0 and n
	interaction::I          # the continuous interaction
	n::Int                  # Number of potential values (= maximum(refs))
	function FixedEffect{R, I}(refs, interaction, n) where {R <: AbstractVector{<:Integer}, I <: AbstractVector{<: Real}}
		length(refs) == length(interaction) || error("refs and interaction don't have the same length")
		new(refs, interaction, n)
	end
end

function FixedEffect(args...; interaction::AbstractVector = uweights(length(args[1])))
	g = group(args...)
	FixedEffect{typeof(g.refs), typeof(interaction)}(g.refs, interaction, g.n)
end

function Base.show(io::IO, fe::FixedEffect)
	println(io, "Refs:        ", length(fe.refs), "-element ", typeof(fe.refs))
	println(io, "             ", Int.(fe.refs[1:min(5, length(fe.refs))]), "...")
	println(io, "Interaction: ", length(fe.interaction), "-element ", typeof(fe.interaction))
	println(io, "             ", fe.interaction[1:min(5, length(fe.interaction))], "...")
end

Base.length(fe::FixedEffect) = length(fe.refs)
Base.eltype(fe::FixedEffect) = eltype(I)

Base.getindex(fe::FixedEffect, esample::Colon) = fe
function Base.getindex(fe::FixedEffect{R, I}, esample::AbstractVector) where {R, I}
	if fe.interaction isa UnitWeights
		FixedEffect{R, I}(fe.refs[esample], uweights(eltype(I), sum(esample)),  fe.n)
	else
		FixedEffect{R, I}(fe.refs[esample], fe.interaction[esample],  fe.n)
	end
end

##############################################################################
##
## group combines multiple refs
## Missings have a ref of 0
##
##############################################################################

mutable struct GroupedArray{N} <: AbstractArray{UInt32, N}
	refs::Array{UInt32, N}   # refs must be between 0 and n. 0 means missing
	n::Int                   # Number of potential values (= maximum(refs))
end
group(xs::GroupedArray) = xs

function group(xs::AbstractArray)
	refs = Array{UInt32}(undef, size(xs))
	invpool = Dict{eltype(xs), UInt32}()
	n = 0
	has_missing = false
	@inbounds for i in eachindex(xs)
		x = xs[i]
		if x === missing
			refs[i] = 0
			has_missing = true
		else
			lbl = get(invpool, x, 0)
			if lbl !== 0
				refs[i] = lbl
			else
				n += 1
				refs[i] = n
				invpool[x] = n
			end
		end
	end
	return GroupedArray{ndims(xs)}(refs, n)
end

function group(args...)
	g1 = deepcopy(group(args[1]))
	for j = 2:length(args)
		gj = group(args[j])
		length(g1.refs) == length(gj.refs) || throw(DimensionError())
		combine!(g1, gj)
	end
	factorize!(g1)
end

function combine!(g1::GroupedArray, g2::GroupedArray)
	@inbounds for i in eachindex(g1.refs, g2.refs)
		# if previous one is missing or this one is missing, set to missing
		g1.refs[i] = (g1.refs[i] == 0 || g2.refs[i] == 0) ? 0 : g1.refs[i] + (g2.refs[i] - 1) * g1.n
	end
	g1.n = g1.n * g2.n
	return g1
end

function factorize!(x::GroupedArray{N}) where {N}
	refs = x.refs
	uu = sort!(unique(refs))
	has_missing = uu[1] == 0
	ngroups = length(uu) - has_missing
	dict = Dict{UInt32, UInt32}(zip(uu, UInt32(1-has_missing):UInt32(ngroups)))
	for i in eachindex(refs)
		refs[i] = dict[refs[i]]
	end
	GroupedArray{N}(refs, ngroups)
end

##############################################################################
##
## Find connected components
##
##############################################################################
# Return a vector of sets that contains the indices of each unique value
function refsrev(fe::FixedEffect)
	out = Vector{Int}[Int[] for _ in 1:fe.n]
	for i in eachindex(fe.refs)
		push!(out[fe.refs[i]], i)
	end
	return out
end

# Returns a vector of all components
# A component is a vector that, for each fixed effect,
# contains all the refs that are included in the component.
function components(fes::AbstractVector{<:FixedEffect})
	refs_vec = Vector{UInt32}[fe.refs for fe in fes]
	refsrev_vec = Vector{Vector{Int}}[refsrev(fe) for fe in fes]
	visited = falses(length(refs_vec[1]))
	out = Vector{Set{Int}}[]
	for i in eachindex(visited)
		if !visited[i]
			# obs not visited yet, so create new component
			component_vec = Set{UInt32}[Set{UInt32}() for _ in 1:length(refsrev_vec)]
			# visit all obs in the same components
			tovisit = Set{Int}(i)
			while !isempty(tovisit)
				for (component, refs, refsrev) in zip(component_vec, refs_vec, refsrev_vec)
					ref = refs[i]
					# if group is not in component yet
					if ref ∉ component
						# add group to the component
						push!(component, ref)
						# visit other observations in same group
						union!(tovisit, refsrev[ref])
					end
				end
				# mark obs as visited
				i = pop!(tovisit)
				visited[i] = true
			end
			push!(out, component_vec)
		end
	end
	return out
end

##############################################################################
##
## normalize! a vector of fixedeffect coefficients using connected components
##
##############################################################################

function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}; kwargs...)
	# The solution is generally not unique. Find connected components and scale accordingly
	idx = findall(fe -> isa(fe.interaction, UnitWeights), fes)
	length(idx) >= 2 && rescale!(view(fecoefs, idx), view(fes, idx))
	return fecoefs
end

function rescale!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect})
	for component_vec in components(fes)
		m = 0.0
		# demean all fixed effects except the first
		for j in length(fecoefs):(-1):2
			fecoef, component = fecoefs[j], component_vec[j]
			mj = 0.0
			for k in component
				mj += fecoef[k]
			end
			mj = mj / length(component)
			for k in component
				fecoef[k] -= mj
			end
			m += mj
		end
		# rescale the first fixed effects
		fecoef, component = fecoefs[1], component_vec[1]
		for k in component
			fecoef[k] += m
		end
	end
end

function full(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect})
	[fecoef[fe.refs] for (fecoef, fe) in zip(fecoefs, fes)]
end


module EMLtray
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...streams.EMLstreams
	import EMLtypes.length
	include("tray/tray.jl")
	include("tray/trayReac.jl")
	include("tray/packedStage.jl")
end


using WindMixing
using JLD2
using Plots

train_files = ["diurnal"]

𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
𝒟train.T.coarse
𝒟train.t

file = jldopen("Data/three_layer_constant_fluxes_linear_hr192_Qu5.0e-04_Qb4.0e-08_f1.0e-04_Nh256_Nz128_diurnal_8days_statistics.jld2")

keys(file["timeseries/wT"])
file["timeseries/wT/1973"][:]
file["timeseries/wT/serialized/location"]
VIDEO_NAME = "u_v_T_3sim_-1e-3_-8e-4_-5e-4_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_1e-2_rate_2e-4_test_-1e-3"
VIDEO_NAME = "test_video"

plot(file["timeseries/wT/292"][:], -128:1:0, label="")

close(file)

animate_training_data_profiles_fluxes(train_files, joinpath(FILE_PATH, VIDEO_NAME))

function a(x, y)
    x-y
end

file = jldopen("D:\\Downloads\\test.jld2", "w") do file
    file["function"] = a(1,2)
end

read_file = jldopen("D:\\Downloads\\test.jld2")
read_file["function"]

read_file["function"](1,2)

close(read_file)


function Rasters.RasterStack(output::DeterministicOutput)
    Rasters.RasterStack(output.name)
end

abstract type AbstractMixingRatio end
struct MixingRatio{T, N} <: AbstractMixingRatio
    raster::AbstractRaster{T, N}
    relpoints::AbstractArray
end
# MixingRatio(path::String) = MixingRatio(Raster(path, name= :spec001_mr))
function MixingRatio(output::AbstractOutputFile)
    rasterstack = RasterStack(output)
    rellat, rellon = rasterstack[:RELLAT1], rasterstack[:RELLNG1]
    [[lon, lat] for (lon, lat) in zip(rellon, rellat)]
    MixingRatio(rasterstack[:spec001_mr], [[lon, lat] for (lon, lat) in zip(rellat, rellon)])
end
Base.show(io::IO, mime::MIME"text/plain", mr::AbstractMixingRatio) = show(io, mime, mr.raster)

struct CubicMixingRatio
    raster::AbstractRaster
    relpoints::AbstractArray
end


### Functions to Generate Pose Graphs ###
### Ted Steiner ###
### October 2014 ###

# Generate pose graph measurements from a route, set of poses, and list of loop-closure sites.
function makePoseGraph( nodes::Dict,
                        anchors::Array{Int,1},
                        route::Array{Int,1},
                        poses::Vector{Pose2D},
                        filename::String="posegraph.txt" )

    anchor_sets = Array{Int,1}[]
    push!(anchor_sets,anchors)

    filenames = String[]
    push!(filenames,filename)

    return makePoseGraph(nodes,anchor_sets,route,poses,filenames)
end

# Generate pose graphs from a single route and set of poses and loop-closures to multiple anchor sets
# Array of filenames should align with array of anchor_sets
function makePoseGraph( nodes::Dict,
                        anchor_sets::Array{Array{Int,1},1},
                        route::Array{Int,1},
                        poses::Vector{Pose2D},
                        filenames::Array{String,1} )

    odom_inf = diagm([5,5,50])
    loop_inf = diagm([50,50,100])
    odom_cov = inv(odom_inf)
    loop_cov = inv(loop_inf)

    odometry = getPoseOdometry(route,poses,odom_cov)
    meas = getMeasurements(route,poses,anchor_sets,odometry,loop_cov)

    for k = 1:length(anchor_sets)
        writeMeasurements(filenames[k],meas[k])
    end

    return nothing
end

function makePoseGraph( nodes::Dict,
                        anchor_sets::Array{Array{Int,1},1},
                        route::Array{Int,1},
                        poses::Vector{Pose2D},
                        odometry::Vector{BetweenFactor},
                        filenames::Array{String,1} )

    loop_inf = diagm([50,50,100])
    loop_cov = inv(loop_inf)

    meas = getMeasurements(route,poses,anchor_sets,odometry,loop_cov)

    for k = 1:length(anchor_sets)
        writeMeasurements(filenames[k],meas[k])
    end

    return nothing
end

function makePoseGraph( anchor_sets::Array{Array{Int,1},1},
                        route::Vector{Int},
                        poses::Vector{Pose2D},
                        odometry_factors::Vector{BetweenFactor},
                        loop_factors::Dict{(Int,Int),BetweenFactor},
                        prior_meas::Dict{Int,Meas2D},
                        prior_inf::Inf2D,
                        filenames::Vector{String} )

    meas = getMeasurements(route,poses,anchor_sets,odometry_factors,loop_factors,prior_meas,prior_inf)

    for k = 1:length(anchor_sets)
        writeMeasurements(filenames[k],meas[k])
    end

    return nothing
end


#
# foundations.jl -
#
# Sub-module exporting types and methods needed to extend or implement
# LazyAlgebra mappings.
#
#-------------------------------------------------------------------------------
#
# This file is part of LazyAlgebra (https://github.com/emmt/LazyAlgebra.jl)
# released under the MIT "Expat" license.
#
# Copyright (c) 2017-2021 Éric Thiébaut.
#

"""
    using LazyAlgebra.Foundations

imports types and methods that may be useful to extend or implement
`LazyAlgebra` mappings.

"""
module Foundations

using ..LazyAlgebra

for sym in (Symbol("@callable"),
            :Adjoint,
            :AdjointInverse,
            :DiagonalMapping,
            :DiagonalType,
            :Direct,
            :Endomorphism,
            :Inverse,
            :InverseAdjoint,
            :Linear,
            :LinearType,
            :Morphism,
            :MorphismType,
            :NonDiagonalMapping,
            :NonLinear,
            :NonSelfAdjoint,
            :Operations,
            :SelfAdjoint,
            :SelfAdjointType,
            :axpby_yields_zero,
            :axpby_yields_y,
            :axpby_yields_my,
            :axpby_yields_by,
            :axpby_yields_x,
            :axpby_yields_xpy,
            :axpby_yields_xmy,
            :axpby_yields_xpby,
            :axpby_yields_mx,
            :axpby_yields_ymx,
            :axpby_yields_mxmy,
            :axpby_yields_bymx,
            :axpby_yields_ax,
            :axpby_yields_axpy,
            :axpby_yields_axmy,
            :axpby_yields_axpby,
            :multiplier_type,
            :multiplier_floatingpoint_type,
            :promote_multiplier)
    @eval begin
        import ..LazyAlgebra: $sym
        export $sym
    end
end

end # module Foundations


using Turing, Test
using Turing: reconstruct, invlink, step
using Turing.VarReplay
using Turing.VarReplay: uid, cuid, getvals, getidcs, set_retained_vns_del_by_spl!, is_flagged, unset_flag!

randr(vi::VarInfo, vn::VarName, dist::Distribution, spl::Turing.Sampler, count::Bool) = begin
  if ~haskey(vi, vn)
    r = rand(dist)
    Turing.push!(vi, vn, r, dist, spl.alg.gid)
    r
  elseif is_flagged(vi, vn, "del")
    unset_flag!(vi, vn, "del")
    r = rand(dist)
    Turing.setval!(vi, Turing.vectorize(dist, r), vn)
    r
  else
    if count Turing.checkindex(vn, vi, spl) end
    Turing.updategid!(vi, vn, spl)
    vi[vn]
  end
end

# Test for uid() (= string())
csym = gensym()
vn1 = VarName(csym, :x, "[1]", 1)
@test string(vn1) == "{$csym,x[1]}:1"
# println(string(vn1))

vn2 = VarName(csym, :x, "[1]", 2)
vn11 = VarName(csym, :x, "[1]", 1)

@test cuid(vn1) == cuid(vn2)
@test vn11 == vn1

vi = VarInfo()
dists = [Normal(0, 1), MvNormal([0; 0], [1.0 0; 0 1.0]), Wishart(7, [1 0.5; 0.5 1])]

alg = PG(PG(5,5),2)
spl2 = Turing.Sampler(alg)
vn_w = VarName(gensym(), :w, "", 1)
randr(vi, vn_w, dists[1], spl2, true)

vn_x = VarName(gensym(), :x, "", 1)
vn_y = VarName(gensym(), :y, "", 1)
vn_z = VarName(gensym(), :z, "", 1)
vns = [vn_x, vn_y, vn_z]

alg = PG(PG(5,5),1)
spl1 = Turing.Sampler(alg)
for i = 1:3
  r = randr(vi, vns[i], dists[i], spl1, false)
  val = vi[vns[i]]
  @test sum(val - r) <= 1e-9
end

# println(vi)

@test length(getvals(vi, spl1)) == 3
@test length(getvals(vi, spl2)) == 1

vn_u = VarName(gensym(), :u, "", 1)
randr(vi, vn_u, dists[1], spl2, true)

# println(vi)
vi.num_produce = 1
set_retained_vns_del_by_spl!(vi, spl2)

# println(vi)

vals_of_1 = collect(getvals(vi, spl1))
# println(vals_of_1)
filter!(v -> ~any(map(x -> isnan.(x), v)), vals_of_1)
@test length(vals_of_1) == 3

vals_of_2 = collect(getvals(vi, spl2))
# println(vals_of_2)
filter!(v -> ~any(map(x -> isnan.(x), v)), vals_of_2)
@test length(vals_of_2) == 1

@model gdemo() = begin
  x ~ InverseGamma(2,3)
  y ~ InverseGamma(2,3)
  z ~ InverseGamma(2,3)
  w ~ InverseGamma(2,3)
  u ~ InverseGamma(2,3)
end

# Test the update of group IDs
g_demo_f = gdemo()
g = Turing.Sampler(Gibbs(1000, PG(10, 2, :x, :y, :z), HMC(1, 0.4, 8, :w, :u)))

pg, hmc = g.info[:samplers]

vi = Turing.VarInfo()
g_demo_f(vi, nothing)
vi = step(g_demo_f, pg, vi)
@test vi.gids == [1,1,1,0,0]

g_demo_f(vi, hmc)
@test vi.gids == [1,1,1,2,2]


for dtype in ("", "Dist")
    for (elty, ext) in ((:Float32, :s),
                        (:Float64, :d),
                        (:Complex64, :c),
                        (:Complex128, :z))
        mat = Symbol(dtype, "Matrix")
        @eval begin
            function inverse!(A::$mat{$elty})
                err = ccall(($(string("ElInverse", dtype, "_", ext)), libEl), Cuint,
                    (Ptr{Void},),
                    A.obj,)
                err == 0 || throw(ElError(err))
                return A
            end
        end
    end
end


function FieldAttitude(arg0::Field, arg1::Attitude)
    return FieldAttitude((Field, Attitude), arg0, arg1)
end

function FieldAttitude(arg0::FieldAbsoluteDate, arg1::Frame, arg2::FieldAngularCoordinates)
    return FieldAttitude((FieldAbsoluteDate, Frame, FieldAngularCoordinates), arg0, arg1, arg2)
end

function FieldAttitude(arg0::FieldAbsoluteDate, arg1::Frame, arg2::FieldRotation, arg3::FieldVector3D, arg4::FieldVector3D)
    return FieldAttitude((FieldAbsoluteDate, Frame, FieldRotation, FieldVector3D, FieldVector3D), arg0, arg1, arg2, arg3, arg4)
end

function FieldAttitude(arg0::FieldAbsoluteDate, arg1::Frame, arg2::Rotation, arg3::Vector3D, arg4::Vector3D, arg5::Field)
    return FieldAttitude((FieldAbsoluteDate, Frame, Rotation, Vector3D, Vector3D, Field), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldAttitude(arg0::Frame, arg1::TimeStampedFieldAngularCoordinates)
    return FieldAttitude((Frame, TimeStampedFieldAngularCoordinates), arg0, arg1)
end

function get_date(obj::FieldAttitude)
    return jcall(obj, "getDate", FieldAbsoluteDate, ())
end

function get_orientation(obj::FieldAttitude)
    return jcall(obj, "getOrientation", TimeStampedFieldAngularCoordinates, ())
end

function get_reference_frame(obj::FieldAttitude)
    return jcall(obj, "getReferenceFrame", Frame, ())
end

function get_rotation(obj::FieldAttitude)
    return jcall(obj, "getRotation", FieldRotation, ())
end

function get_rotation_acceleration(obj::FieldAttitude)
    return jcall(obj, "getRotationAcceleration", FieldVector3D, ())
end

function get_spin(obj::FieldAttitude)
    return jcall(obj, "getSpin", FieldVector3D, ())
end

function interpolate(obj::FieldAttitude, arg0::FieldAbsoluteDate, arg1::Stream)
    return jcall(obj, "interpolate", FieldAttitude, (FieldAbsoluteDate, Stream), arg0, arg1)
end

function interpolate(obj::FieldTimeInterpolable, arg0::FieldAbsoluteDate, arg1::Collection)
    return jcall(obj, "interpolate", FieldTimeInterpolable, (FieldAbsoluteDate, Collection), arg0, arg1)
end

function shifted_by(obj::FieldAttitude, arg0::RealFieldElement)
    return jcall(obj, "shiftedBy", FieldAttitude, (RealFieldElement,), arg0)
end

function shifted_by(obj::FieldAttitude, arg0::jdouble)
    return jcall(obj, "shiftedBy", FieldAttitude, (jdouble,), arg0)
end

function to_attitude(obj::FieldAttitude)
    return jcall(obj, "toAttitude", Attitude, ())
end

function with_reference_frame(obj::FieldAttitude, arg0::Frame)
    return jcall(obj, "withReferenceFrame", FieldAttitude, (Frame,), arg0)
end


@testset "Index: $IndexType" for IndexType in [NaiveIndex, BruteTreeIndex,
                                    KDTreeIndex, HNSWIndex, IVFIndex, NoopIndex]
    T = eltype(1.0)
    data = T[0 0 0 5 5 5; 0 1 2 10 11 12]
    spdata = sparse(data)
    point = T[5.1, 10]
    true_length = size(data, 2)

    if IndexType === IVFIndex
        _indexfunc = d->IVFIndex(d; kc=4, k=2, m=1)
        index = _indexfunc(data)
        spindex = _indexfunc(spdata)
    else
        index = IndexType(data)
        spindex = IndexType(data)
    end
    @test index isa IndexType
    idxs, scores = Garamond.knn_search(index, point, 10; w=4)
    @test idxs isa Vector{Int}
    @test scores isa Vector{T}
    if IndexType != NoopIndex
        @test all(i in idxs for i in 1:true_length)
    else
        @test isempty(idxs)
        @test isempty(scores)
    end
    @test length(index) == length(spindex) == true_length

    # Test push!, pop! interface
    idx = 1
    point = data[:, idx]
    if IndexType in [NaiveIndex, BruteTreeIndex, IVFIndex]
        # push!
        @test try push!(index, point); true; catch; false; end
        @test length(index) == true_length + 1

        # pop!
        if !isa(index, IVFIndex)
            @test pop!(index) == point
        else
            @test pop!(index) isa Vector{T} # ivfadc quantizes the vector, it will not be the same
        end
        @test length(index) == true_length

        # pushfirst!
        @test try pushfirst!(index, point); true; catch; false; end
        @test length(index) == true_length + 1

        # popfirst!
        if !isa(index, IVFIndex)
            @test popfirst!(index) == point
        else
            @test popfirst!(index) isa Vector{T}
        end
        @test length(index) == true_length

        # delete_from_index!
        @test Garamond.delete_from_index!(index, [idx]) === nothing
        @test length(index) == true_length - 1
    elseif IndexType == NoopIndex
        # push!
        @test try push!(index, point); true; catch; false; end
        @test length(index) == true_length  # nothing pushed

        # pop!
        @test pop!(index) === nothing
        @test length(index) == true_length

        # pushfirst!
        @test try pushfirst!(index, point); true; catch; false; end
        @test length(index) == true_length  # nothing pushed

        # popfirst!
        @test popfirst!(index) === nothing
        @test length(index) == true_length

        # delete_from_index!
        @test Garamond.delete_from_index!(index, [idx]) === nothing
        @test length(index) == true_length

    elseif IndexType in [KDTreeIndex, HNSWIndex]
        @test_throws MethodError push!(idx, point)
        @test_throws MethodError pushfirst!(idx, point)
        @test_throws MethodError pop!(idx)
        @test_throws MethodError popfirst!(idx)
        @test_throws MethodError Garamond.delete_from_index!(idx, [1,2])
    end
end


#=

  "TRIBE: Trait and Range Interspecific Biogeographic Evolution" submodule package

=#

module TRIBE

using Random: randexp
using DelimitedFiles: readdlm, writedlm
using ProgressMeter: Progress, next!
using Optim: minimizer, optimize, Options
using LinearAlgebra: BLAS.axpy!, BLAS.gemm!, eigvecs, eigvals, diagm
using RCall: reval, @rput

# other submodules dependencies
using ..Utils

include("tribe/tribe_utils.jl")
include("tribe/cont_DA_prop.jl")
include("tribe/disc_DA_prop.jl")
include("tribe/tribe_data_initializer.jl")
include("tribe/area_lineage_averages.jl")
include("tribe/tribe_loglik_functions.jl")
include("tribe/mcmc_utils.jl")
include("tribe/mcmc.jl")
include("tribe/parameter_updates.jl")
include("tribe/mcmc_burn.jl")
include("tribe/under_prior.jl")
include("tribe/proposal_functions.jl")
include("tribe/tribe_wrapper.jl")
include("tribe/tribe_simulation_utils.jl")

end # module TRIBE


# contains main ASWING calling routines

"""
    setgeom(asw::Configuration)

    Inputs geometry into ASWING
"""
function setgeom(asw::Configuration)

    # clear variable array sizes
    ASWING.NB[:, :] .= 0
    ASWING.NANGJ[:] .= 0
    npylo = 0 #pylon counter

    # NAME
    for i = 1:length(asw.namein)
        ASWING.NAME[i] = asw.namein[i]
    end

    # UNITS
    if asw.si_units
        UNITS.UNITL[1] = 1.0
        UNITS.UNCHL[:] = Vector{Char}(rpad("m",32))
        UNITS.UNITT[1] = 1.0
        UNITS.UNCHT[:] = Vector{Char}(rpad("s",32))
        UNITS.UNITM[1] = 1.0
        UNITS.UNCHM[:] = Vector{Char}(rpad("kg",32))
        UNITS.UNITF[1] = 0.0
        UNITS.UNCHF[:] = Vector{Char}(rpad(" ",32))
    else
        UNITS.UNITL[1] = 1.0
        UNITS.UNCHL[:] = Vector{Char}(rpad("ft",32))
        UNITS.UNITT[1] = 1.0
        UNITS.UNCHT[:] = Vector{Char}(rpad("s ",32))
        UNITS.UNITM[1] = 0.0
        UNITS.UNCHM[:] = Vector{Char}(rpad("  ",32))
        UNITS.UNITF[1] = 1.0
        UNITS.UNCHF[:] = Vector{Char}(rpad("lb",32))
    end

    # CONSTANTS BLOCK
    ASWING.GEEW[1] = asw.constants.gravity
    ASWING.RHOSL[1] = asw.constants.sea_level_air_density
    ASWING.VSOSL[1] = asw.constants.sea_level_speed_of_sound

    # REFERENCE BLOCK
    ASWING.SREF[1] = asw.refvals.area
    ASWING.CREF[1] = asw.refvals.chord
    ASWING.BREF[1] = asw.refvals.span
    ASWING.XYZREF[:,1] = asw.refvals.moments
    ASWING.XYZREF[:,2] = asw.refvals.accelerations
    ASWING.XYZREF[:,3] = asw.refvals.velocities

    # SENSOR BLOCK
    nsen = length(asw.sensors)
    for i = 1:nsen
        ASWING.KPTYPE[npylo+i] = 100 + asw.sensors[i].Ksens
        ASWING.KBPYLO[npylo+i] = asw.sensors[i].Nbeam
        ASWING.QPYLO[0,npylo+i] = asw.sensors[i].t
        ASWING.QPYLO[1:3,npylo+i] = asw.sensors[i].location
        ASWING.QPYLO[4:6,npylo+i] = asw.sensors[i].V
        ASWING.QPYLO[7:9,npylo+i] = asw.sensors[i].A
    end
    npylo += nsen

    # ENGINE BLOCK
    neng = length(asw.engines)
    for i = 1:neng
        keng = asw.engines[i].Keng
        ietyp = asw.engines[i].IEtyp
        ASWING.KBPYLO[npylo+i] = asw.engines[i].Nbeam
        ASWING.QPYLO[0,npylo+i] = asw.engines[i].t
        ASWING.QPYLO[1:3,npylo+i] = asw.engines[i].location
        ASWING.QPYLO[4:6,npylo+i] = asw.engines[i].thrust_direction
        ASWING.QPYLO[7,npylo+i] = asw.engines[i].dFdPe
        ASWING.QPYLO[8,npylo+i] = asw.engines[i].dMdPe
        ASWING.QPYLO[9,npylo+i] = asw.engines[i].Rdisk
        ASWING.QPYLO[10,npylo+i] = asw.engines[i].Omega
        ASWING.QPYLO[11,npylo+i] = asw.engines[i].cdA
        ASWING.QPYLO[12,npylo+i] = asw.engines[i].cl
        ASWING.QPYLO[13,npylo+i] = asw.engines[i].CLa
        ASWING.QPYLO[14,npylo+i] = asw.engines[i].S0
        ASWING.QPYLO[15,npylo+i] = asw.engines[i].C0
        ASWING.QPYLO[16,npylo+i] = asw.engines[i].S1
        ASWING.QPYLO[17,npylo+i] = asw.engines[i].C1
        ASWING.QPYLO[18,npylo+i] = asw.engines[i].S2
        ASWING.QPYLO[19,npylo+i] = asw.engines[i].C2
        ASWING.QPYLO[20,npylo+i] = asw.engines[i].S3
        ASWING.QPYLO[21,npylo+i] = asw.engines[i].C3

        ASWING.KPTYPE[npylo+i] = 10 + keng
        ke = max(1, min(NENGX, keng))
        ASWING.IENGTYP[ke] = ietyp
    end
    npylo += neng

    # WEIGHT BLOCK
    nwei = length(asw.weights)
    for i = 1:nwei
        ASWING.KPTYPE[npylo+i] = 1
        ASWING.KBPYLO[npylo+i] = asw.weights[i].Nbeam
        ASWING.QPYLO[0,npylo+i] = asw.weights[i].t
        ASWING.QPYLO[1:3,npylo+i] = asw.weights[i].location
        ASWING.QPYLO[4,npylo+i] = asw.weights[i].weight
        ASWING.QPYLO[5,npylo+i] = asw.weights[i].CDA
        ASWING.QPYLO[6,npylo+i] = asw.weights[i].vol
        ASWING.QPYLO[7:9,npylo+i] = asw.weights[i].Ho
        ASWING.QPYLO[10,npylo+i] = asw.weights[i].Ixx
        ASWING.QPYLO[11,npylo+i] = asw.weights[i].Iyy
        ASWING.QPYLO[12,npylo+i] = asw.weights[i].Izz
        ASWING.QPYLO[13,npylo+i] = asw.weights[i].Ixy
        ASWING.QPYLO[14,npylo+i] = asw.weights[i].Ixz
        ASWING.QPYLO[15,npylo+i] = asw.weights[i].Iyz
    end
    npylo += nwei

    # STRUT BLOCK
    nstr = length(asw.struts)
    for i = 1:nstr
        ASWING.KPTYPE[npylo+i] = 2
        ASWING.KBPYLO[npylo+i] = asw.struts[i].Nbeam
        ASWING.QPYLO[0,npylo+i] = asw.struts[i].t
        ASWING.QPYLO[1:3,npylo+i] = asw.struts[i].location
        ASWING.QPYLO[4:6,npylo+i] = asw.struts[i].anchor_location
        ASWING.QPYLO[7,npylo+i] = asw.struts[i].dLo
        ASWING.QPYLO[8,npylo+i] = asw.struts[i].EAw
    end
    npylo += nstr
    ASWING.NPYLO[1] = npylo

    # JOINT BLOCK
    njoin = length(asw.joints)
    ASWING.NJOIN[1] = njoin
    for i = 1:njoin
        ASWING.KBJOIN[1,i] = asw.joints[i].Nbeam1
        ASWING.KBJOIN[2,i] = asw.joints[i].Nbeam2
        ASWING.TJOIN[1,i] = asw.joints[i].t1
        ASWING.TJOIN[2,i] = asw.joints[i].t2
        ASWING.KJTYPE[i] = asw.joints[i].KJtype
    end

    # JANGLE BLOCK
    for i = 1:length(asw.jangles)
        kj = asw.jangles[i].Njoint
        nangj = length(asw.jangles[i].hinge_moment)
        ASWING.NANGJ[kj] = nangj
        ASWING.HJAX[1:3,kj] = asw.jangles[i].hinge_axis
        ASWING.MOMJ[1:nangj,kj] = asw.jangles[i].hinge_moment
        ASWING.ANGJ[1:nangj,kj] = asw.jangles[i].hinge_angle
    end

    # GROUNDS BLOCK
    ngrou = length(asw.grounds)
    ASWING.NGROU[1] = ngrou
    for i = 1:ngrou
        ASWING.KBGROU[i] = asw.grounds[i].Nbeam
        ASWING.TGROU[i] = asw.grounds[i].t
        ASWING.KGTYPE[i] = asw.grounds[i].KGtype
    end

    # BEAM DEFINITION BLOCKS
    nbeam = length(asw.beams)
    ASWING.NBEAM[1] = nbeam
    for ibeam = 1:nbeam
        ASWING.BNAME[ibeam][1:length(asw.beams[ibeam].name)] = Vector{Char}(asw.beams[ibeam].name)
        ASWING.KBNUM[ibeam] = asw.beams[ibeam].number
        ASWING.IBEAM[ibeam] = asw.beams[ibeam].physical_index
        for var in keys(beamvarval)
            j = beamvarval[var]
            varval = getfield(asw.beams[ibeam].spanwise_variables,Symbol(var))
            nb = size(varval,1)
            ASWING.NB[j, ibeam] = nb
            ASWING.TB[1:nb, j, ibeam] = varval[:,1]
            ASWING.QB[1:nb, j, ibeam] = varval[:,2]
        end
    end

    ccall((:input_, libaswing), Nothing, (Ref{Cchar},), ' ')
    ccall((:qmset_, libaswing), Nothing, ())
    ccall((:dscheck_, libaswing), Nothing, ())
    return nothing
end

"""
    setpnt(point)
    Inputs operating point(s) into ASWING
"""
function setpnt(point::AbstractArray{<:OperatingPoint,1})

    # set number of operating points
    npoint = length(point)
    ASWING.NPOINT[1] = npoint
    for ipnt = 1:npoint
        # set constraints
        setcons(point[ipnt].constraints, ipnt)
        # set parameter values
        setpars(point[ipnt].parameters, ipnt)
        # set other options and weights
        ASWING.LMACH[ipnt] = point[ipnt].mach_from_airspeed
        ASWING.IGRIM[ipnt] = point[ipnt].ground_image
        ASWING.WFLAP[:,ipnt] = point[ipnt].wflap
        ASWING.WPENG[:,ipnt] = point[ipnt].wpeng
        # solution is unconverged
        ASWING.LCONV[ipnt] = false
    end

    # set Mach number for Prandtl-Glauert correction
    for ipnt = 1:npoint
        for ip = 1:KPFREE
            ccall((:psp2par_, libaswing), Nothing, (Ref{Int32}, Ref{Int32}), ip, ipnt)
        end
        ASWING.PARAM[KPMACH, ipnt] =  point[ipnt].machpg
    end

    return nothing
end

"""
    setcons(pnt::dotpnt,ipnt=1)
    Inputs operating points contraints into ASWING
"""
function setcons(constraints, ipnt::Integer=1)
    ASWING.IPPAR[KPVACX:KPVACZ,ipnt] = constraints.linear_acceleration
    ASWING.IPPAR[KPRACX:KPRACZ,ipnt] = constraints.angular_acceleration
    ASWING.IPPAR[KPVELX:KPVELZ,ipnt] = constraints.velocity
    ASWING.IPPAR[KPROTX:KPROTZ,ipnt] = constraints.rotation_rate
    ASWING.IPPAR[KPPOSX:KPPOSZ,ipnt] = constraints.position
    ASWING.IPPAR[KPBANK,ipnt] = constraints.phi
    ASWING.IPPAR[KPELEV,ipnt] = constraints.theta
    ASWING.IPPAR[KPHEAD,ipnt] = constraints.psi
    ASWING.IPPAR[KPFLP1:KPFLP1+length(constraints.flap_defl_ctrl_var)-1,ipnt] = constraints.flap_defl_ctrl_var
    for ip = KPFLP1+length(constraints.flap_defl_ctrl_var):KPFLP1+NFLPX
        ASWING.IPPAR[ip,ipnt] = ip
    end
    ASWING.IPPAR[KPENG1:KPENG1+length(constraints.eng_pwr_ctrl_var)-1,ipnt] = constraints.eng_pwr_ctrl_var
    for ip = KPENG1+length(constraints.eng_pwr_ctrl_var):KPENG1+NENGX
        ASWING.IPPAR[ip,ipnt] = ip
    end
    for kf=1 : NFLPX
        kpflp1 = KPFLP1[1]
        ipflp1 = IPFLP1[1]
        if ASWING.LSFLAP[kf] == 1
            if ASWING.IPPAR[kpflp1-1+kf,ipnt] == 0
                ASWING.IPPAR[kpflp1-1+kf,ipnt] = ipflp1-1+kf
            end
        else
            ASWING.IPPAR[kpflp1-1+kf,ipnt] = 0
        end
    end
    for ke=1 : NENGX
        kpeng1 = KPENG1[1]
        ipeng1 = IPENG1[1]
        if ASWING.LSPENG[ke] == 1
            if ASWING.IPPAR[kpeng1-1+ke, ipnt] == 0
                ASWING.IPPAR[kpeng1-1+ke, ipnt] = ipeng1-1 + ke
            end
        else
            ASWING.IPPAR[kpeng1[1]-1+ke,ipnt] = 0
        end
    end
    ASWING.IPPAR[KPVIDT,ipnt] = constraints.err_int_Vinf
    ASWING.IPPAR[KPBEDT,ipnt] = constraints.err_int_beta
    ASWING.IPPAR[KPALDT,ipnt] = constraints.err_int_alpha
    ASWING.IPPAR[KPBADT,ipnt] = constraints.err_int_phi
    ASWING.IPPAR[KPELDT,ipnt] = constraints.err_int_theta
    ASWING.IPPAR[KPHEDT,ipnt] = constraints.err_int_psi
    ASWING.IPPAR[KPWDTX:KPWDTZ,ipnt] = constraints.err_int_ROTxyz
    ASWING.IPPAR[KPADTX:KPADTZ,ipnt] = constraints.err_int_VACxyz
    ASWING.LCONV[ipnt] = false
    return nothing
end

"""
    setpars(parameters, ipnt::Integer=1)
    Sets operating point parameters into ASWING
"""
function setpars(parameters, ipnt::Integer=1)

    ASWING.PSPEC[IPVACX:IPVACZ,ipnt] = parameters.linear_acceleration
    ASWING.PSPEC[IPRACX:IPRACZ,ipnt] = parameters.angular_acceleration
    ASWING.PSPEC[IPVIAS,ipnt] = parameters.velocity
    ASWING.PSPEC[IPBETA,ipnt] = parameters.beta
    ASWING.PSPEC[IPALFA,ipnt] = parameters.alpha
    ASWING.PSPEC[IPROTX:IPROTZ,ipnt] = parameters.rotation_rate
    ASWING.PSPEC[IPPOSX:IPPOSZ,ipnt] = parameters.position
    ASWING.PSPEC[IPBANK,ipnt] = parameters.phi
    ASWING.PSPEC[IPELEV,ipnt] = parameters.theta
    ASWING.PSPEC[IPHEAD,ipnt] = parameters.psi
    ASWING.PSPEC[IPFLP1:IPFLP1+length(parameters.flap_defl_ctrl_var)-1,ipnt] = parameters.flap_defl_ctrl_var
    for ip = IPFLP1+length(parameters.flap_defl_ctrl_var):IPFLP1+NFLPX
      ASWING.PSPEC[ip,ipnt] = 0.0
    end
    ASWING.PSPEC[IPENG1:IPENG1+length(parameters.eng_pwr_ctrl_var)-1,ipnt] = parameters.eng_pwr_ctrl_var
    for ip = IPENG1+length(parameters.eng_pwr_ctrl_var):IPENG1+NENGX
      ASWING.PSPEC[ip,ipnt] = 0.0
    end
    ASWING.PSPEC[IPFORX:IPFORZ,ipnt] = parameters.sum_force
    ASWING.PSPEC[IPMOMX:IPMOMZ,ipnt] = parameters.sum_mom
    ASWING.PSPEC[IPLIFT,ipnt] = parameters.lift
    ASWING.PSPEC[IPGAMM,ipnt] = parameters.climb_angle
    ASWING.PSPEC[IPRACC,ipnt] = parameters.radial_acceleration
    ASWING.PSPEC[IPUSR1,ipnt] = parameters.usr1
    ASWING.PSPEC[IPUSR2,ipnt] = parameters.usr2
    ASWING.PSPEC[IPLSQD,ipnt] = parameters.least_squared

    return nothing
end

"""
    solvesteady(;repeat::Integer=0)
    solvesteady(ipnt1::Integer; repeat::Integer=0)
    solvesteady(ipnt1::Integer, ipnt2::Integer; repeat::Integer=0)

    Quasi-steady ASWING solve. If no arguments are specified, all points are
    solved.
"""
function solvesteady(ipnt1::Integer, ipnt2::Integer; repeat::Integer=0)

    ASWING.STEADY[1] = true
    for ipnt in ipnt1:ipnt2
        # set operating point
        ASWING.ARGP1[1:4] .= Vector{Char}(rpad("$ipnt",4)) # hack to avoid string passing
        ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), 0.0, 0)
        for i = 1:repeat+1
            # converge operating point
            ASWING.ARGP1[1:4] .= Vector{Char}(rpad("X",4)) # hack to avoid string passing
            ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), 0.0, 0)
            if Bool(ASWING.LCONV[ipnt])
                break
            end
        end
    end

    # return formatted solution
    return getsolution(ipnt1,ipnt2)
end
solvesteady(; repeat::Integer=0) = solvesteady(1, ASWING.NPOINT[1], repeat=repeat)
solvesteady(ipnt1::Integer; repeat::Integer=0) = solvesteady(ipnt1, ipnt1, repeat=repeat)[ipnt1]

"""
    solveunsteady(deltat::Real, ntimes::Integer)
    Unsteady ASWING time-domain solve
"""
function solveunsteady(deltat::Real, ntimes::Integer, ipnt::Integer=1)

    # set operating point
    ASWING.ARGP1[1:4] .= Vector{Char}(rpad("$ipnt",4)) # hack to avoid string passing
    ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), 0.0, 0)

    # change to unsteady
    ASWING.STEADY[1] = false

    # solve unsteady
    ASWING.ARGP1[1:4] .= Vector{Char}(rpad("X",4)) # hack to avoid string passing
    ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), deltat, ntimes)

    # return formatted solution
    return getsolution(1, ntimes+1)
end

"""
    retainparam()
    retainparam(ipnt1::Integer)
    retainparam(ipnt1::Integer, ipnt2::Integer)
    Sets operating point parameters equivalent to those of the current state
"""
function retainparam(ipnt1::Integer, ipnt2::Integer=ipnt1)

    for ipnt in ipnt1:ipnt2
        # set operating point
        ASWING.ARGP1[1:4] .= Vector{Char}(rpad("$ipnt",4)) #hack to avoid string passing
        ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), 0.0, 0)
        # set parameters on operating point
        ASWING.ARGP1[1:4] .= Vector{Char}(rpad("R",4)) #hack to avoid string passing
        ccall((:operjl_, libaswing), Nothing, (Ref{Float64}, Ref{Int32}), 0.0, 0)
    end

    return nothing
end
retainparam() = retainparam(1, ASWING.NPOINT[1])

"""
    freeze(ipoint = 1)
    Freezes geometry at current state
"""
function freeze(ipoint = 1)

  ccall((:freeze_,libaswing), Nothing, (Ref{Int32},), ipoint)

  for i = 1 : ASWING.NPOINT[1]
    ASWING.LCONV[i] = false
  end

end

"""
    setgust(igust::Integer, params::AbstractArray{<:Real}=[-4.0,0.0,1.0,8.0,2.0,0.0])
    Sets gust type and corresponding parameters
"""
function setgust(igust::Integer, params::AbstractArray{<:Real}=[-4.0,0.0,1.0,8.0,2.0,0.0])

    gustparam = [-4.0,0.0,1.0,8.0,2.0,0.0]
    if length(params) > 6
        error("No more than six gust parameters may be specified")
    end
    gustparam[1:length(params)] = params
    ASWING.IGUST[1] = igust
    ASWING.VGCON[1:6] = gustparam

    return nothing
end

"""
    geteigs(neigen, zshift, ipnt; evmin = 1e-4, evtol = 1e-3)
    Returns eigenvalues of the system
"""
function geteigs(neig, zshift, ipnt; evmin = 1e-4, evtol = 1e-3)

    ASWING.STEADY[1] = false
    ASWING.LFGUST[1] = true

    # number of new modes to compute
    ASWING.NBEIGEN[1] = neig

    # clear current eigenmode counter, and start from new
    EIGEN.NEIGEN[1] = 0
    EIGEN.NROM[1] = 0

    # fix shift if too close to origin
    if abs(zshift) < 0.01
        zshift = -0.01 + 0.0im
    end

    # use shift as new default seed
    EIGEN.ZESEED[1] = zshift

    # compute only eigenvalues for specified operating points
    lrvec = true
    lgerr = false
    ccall((:evexec_, libaswing), Nothing, (Ref{Int32}, Ref{ComplexF64}, Ref{Int32},
        Ref{Int32}, Ref{Int32}, Ref{Float64}, Ref{Float64}, Ref{Int32}),
        ipnt, zshift, ASWING.NBEIGEN[1], lrvec, lgerr, evmin, evtol, 0)

    ccall((:wbcalc_, libaswing), Nothing, (Ref{Int32},), ipnt)

    zlamda = copy(EIGEN.ZLAMDA[1:EIGEN.NEIGEN[1]])

    return zlamda
end

"""
    stabilityderivatives(ipnt = 1, momref=Float64[])
    Calculates stability derivatives
"""
function stabilityderivatives(ipnt::Integer = 1; momref::AbstractArray{<:Real,1}=Float64[])

    if !isempty(momref)
        # assign new moment reference point
        ASWING.XYZREF[:,1] = momref
    end

    # get stability derivatives
    f_vba = zeros(Float64, 3, 3)
    m_vba = zeros(Float64, 3, 3)
    f_rot = zeros(Float64, 3, 3)
    m_rot = zeros(Float64, 3, 3)
    f_flap = zeros(Float64, 3, NFLPX)
    m_flap = zeros(Float64, 3, NFLPX)
    f_peng = zeros(Float64, 3, NENGX)
    m_peng = zeros(Float64, 3, NENGX)
    ccall((:dumpdm_, libaswing), Nothing,
    (Ref{Int32}, Ref{Float64}, Ref{Float64}, Ref{Float64}, Ref{Float64},
    Ref{Float64}, Ref{Float64}, Ref{Float64}, Ref{Float64}, Ref{Int32}),
    6, f_vba, m_vba, f_rot, m_rot,
    f_flap, m_flap, f_peng, m_peng, ipnt)

    # limit output to defined flaps
    idx = findall(x->x!=0, ASWING.LSFLAP)
    f_flap = f_flap[:,idx]
    m_flap = m_flap[:,idx]
    # limit output to defined engines
    idx = findall(x->x!=0, ASWING.LSPENG)
    f_peng = f_peng[:,idx]
    m_peng = m_peng[:,idx]

  return StabilityDerivatives(f_vba, m_vba, f_rot, m_rot, f_flap, m_flap, f_peng, m_peng)
end

"""
    conclr()
    Clears all control laws
"""
function conclr()
    ASWING.CONLAW[1] = false
    ASWING.CONSET[1] = false
    ASWING.LDCON[:] .= false
    ASWING.LUCON[:] .= false
    ASWING.LQCON[:] .= false
    CONLAW.CU[:,:,:,:] .= 0.0
    CONLAW.CQ[:,:,:,:,:] .= 0.0
    CONLAW.DFREQ[:] .= 0.0
    CONLAW.DELTA[:] .= 0.0
    return nothing
end

"""
    conadd(i,j,ls,gfac,c)
    Sets up a proportional control law
"""
function conadd(i::Integer, j::Integer, ls::Integer, gfac::Real, c::Real)

    ASWING.CONLAW[1] = true
    if ls <= 0
        CONLAW.NUPAR1[i,j] = 2
        CONLAW.NUPAR2[i,j] = 2
        CONLAW.UPAR1[i,j,1:2] .= 0.0
        CONLAW.UPAR2[i,j,1:2] .= 0.0
        CONLAW.CU[i,j,1:2,1:2] .= c*gfac
        ASWING.LDCON[i] = true
        ASWING.LUCON[i] = true
    else
        CONLAW.NQPAR1[i,j,ls] = 2
        CONLAW.NQPAR2[i,j,ls] = 2
        CONLAW.QPAR1[i,j,ls,1:2] .= 0.0
        CONLAW.QPAR2[i,j,ls,1:2] .= 0.0
        CONLAW.CQ[i,j,ls,1:2,1:2] .= c*gfac
        ASWING.LDCON[i] = true
        ASWING.LQCON[j,ls] = true
    end
    ASWING.CONSET[1] = true

    return nothing
end

"""
    getintervals(ibeam::Integer)
    Gets number of panels between each break point
"""
function getintervals(ibeam::Integer)
    # get first beam index
    if ASWING.SBRK[1,ibeam] == ASWING.STIPL[ibeam]
        kbfrst = 1
    else
        kbfrst = 0
    end
    # get last beam index
    nbreak = ASWING.NBRK[ibeam]
    if ASWING.SBRK[nbreak,ibeam] == ASWING.STIPR[ibeam]
        kblast = nbreak
    else
        kblast = nbreak+1
    end
    # get spacing
    return ASWING.IBRK[kbfrst+1:kblast,ibeam].-ASWING.IBRK[kbfrst:kblast-1,ibeam].-1
end

"""
    nodeset(intervals::AbstractArray{<:AbstractArray{<:Integer, 1},1})

    Sets new number of nodes for ibeam.  A custom number of panels can be used
    between break points.
"""
function nodeset(intervals::AbstractArray{<:AbstractArray{<:Integer, 1},1})

    if length(intervals) != ASWING.NBEAM[1]
        error("length of intervals must equal ASWING.NBEAM")
    end

    for ibeam = 1:ASWING.NBEAM[1]
        # get first beam index
        if ASWING.SBRK[1,ibeam] == ASWING.STIPL[ibeam]
            kbfrst = 1
        else
            kbfrst = 0
        end

        # get last beam index
        nbreak = ASWING.NBRK[ibeam]
        if ASWING.SBRK[nbreak,ibeam] == ASWING.STIPR[ibeam]
            kblast = nbreak
        else
            kblast = nbreak+1
        end

        # assign new number of nodes to beam
        if !isempty(intervals[ibeam])
            nnodes = sum(intervals[ibeam])+length(intervals[ibeam])
            if nnodes < 2
                error("sum(intervals[$ibeam]) too small. nnodes must be >= 2")
            elseif nnodes > IIX
                error("sum(intervals[$ibeam]) too large. nnodes must be < IIX")
            end
            ASWING.II[ibeam] = nnodes
        end
        ccall((:iiset_, libaswing), Nothing, ())

        # assign new node distribution
        if !isempty(intervals[ibeam])
            THCALC.CUSTOMINTERVALS[1] = true
            THCALC.INTERVALS[kbfrst+1:kblast] = intervals[ibeam]
        else
            THCALC.CUSTOMINTERVALS[1] = false
        end
        ccall((:winit_, libaswing), Nothing, (Ref{Int32}, Ref{Int32}), ibeam, 0)
        ccall((:qinit_, libaswing), Nothing, (Ref{Int32}, Ref{Int32}, Ref{Float64}, Ref{Float64}),
            ASWING.IFRST[ibeam], ASWING.ILAST[ibeam], view(ASWING.Q,1:18,1:ASWING.IITOT[1],ASWING.IPOINT[1]), ASWING.Q0)
    end
    ccall((:spltset_, libaswing), Nothing, ())
    ccall((:micalc_, libaswing), Nothing, (Ref{Int32},), 1)
    ccall((:nnset_, libaswing), Nothing, ())

    # set to unconverged
    ASWING.LQINI[1:NPNTX] .= false
    ASWING.LCONV[1:NPNTX] .= false
    ASWING.LAINI[1:NPNTX] .= false

    # restore default node distribution settings
    THCALC.CUSTOMINTERVALS[1] = false
    return nothing
end

function getstabilityderivatives(ipnt::Integer=1)

    # no convergence failures yet
    fail = false

    # Get quasi-steady solution
    solution = solvesteady(ipnt)

    if solution.converged == false
        fail = true
    end

    # Retain parameters
    retainparam(ipnt)

    # Switch to anchored constraints
    setcons(Constraints("anchored"))

    # Get quasi-steady solution
    solution_qs = solvesteady(ipnt)

    # check if convergence failed
    if solution_qs.converged == false
        fail = true
    end

    # Get stability matrices using center of gravity as moment reference
    stabderivs_qs = stabilityderivatives(ipnt, momref=solution_qs.pnt_param.position_aircraft_cg)

    # rigidize aircraft at current deformed shape
    freeze(ipnt)

    # Get quasi-steady solution
    solution_rigid = solvesteady(ipnt)

    # check if convergence failed
    if solution_rigid.converged == false
        fail = true
    end

    # Get stability matrices using center of gravity as moment reference
    stabderivs_rigid = stabilityderivatives(ipnt, momref=solution_rigid.pnt_param.position_aircraft_cg)

    return stabderivs_qs, stabderivs_rigid, fail
end #getstabilityderivatives

"""
    getstaticmargin(ipnt::Integer=1)

    Returns quasi-steady static margin, rigidized static margin, and a solution
    failure flag.  Assumes constraints are set up for trimmed flight and follows
    the instructions in derivative_calc.txt. This function performs irreversible
    changes so you will need to re-input the geometry afterwards.
"""
function getstaticmargin(ipnt::Integer=1)

    stabderivs_qs, stabderivs_rigid, fail = getstabilityderivatives(ipnt)

    # calculate static margin
    sm_qs = -stabderivs_qs.Cma/stabderivs_qs.CLa*100

    # calculate static margin
    sm_rigid = -stabderivs_rigid.Cma/stabderivs_rigid.CLa*100

    return sm_qs, sm_rigid, fail
end #getstaticmargin


mutable struct Prop2DAcoIsoDenQ_DEO2_FDTD
    p::Ptr{Cvoid}
end

function Prop2DAcoIsoDenQ_DEO2_FDTD(;
        nz=0,
        nx=0,
        nsponge=60,
        nbz=512,
        nbx=8,
        dz=5.0,
        dx=5.0,
        dt=0.001,
        nthreads=Sys.CPU_THREADS,
        freesurface=true,
        freqQ=5.0,
        qMin=0.1,
        qInterior=100.0)
    nz,nx,nsponge,nbz,nbx,nthreads = map(x->round(Int,x), (nz,nx,nsponge,nbz,nbx,nthreads))
    dz,dx,dt = map(x->Float32(x), (dz,dx,dt))

    @assert nx > 0
    @assert nz > 0
    @assert nsponge > 0
    @assert nthreads > 0
    @assert nbx > 0
    @assert nbz > 0

    fs = freesurface ? 1 : 0

    p = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_alloc, libprop2DAcoIsoDenQ_DEO2_FDTD),
        Ptr{Cvoid},
        (Clong, Clong,    Clong, Clong, Clong,   Cfloat, Cfloat, Cfloat, Clong, Clong),
         fs,    nthreads, nx,    nz,    nsponge, dx,     dz,     dt,     nbx,   nbz)

    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_SetupDtOmegaInvQ, libprop2DAcoIsoDenQ_DEO2_FDTD),
        Cvoid,
        (Ptr{Cvoid}, Cfloat, Cfloat, Cfloat),
         p,          freqQ,  qMin,   qInterior)

    Prop2DAcoIsoDenQ_DEO2_FDTD(p)
end

free(prop::Prop2DAcoIsoDenQ_DEO2_FDTD) = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_free, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

function size(prop::Prop2DAcoIsoDenQ_DEO2_FDTD)
    nz = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_getNz, libprop2DAcoIsoDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nx = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_getNx, libprop2DAcoIsoDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    (nz,nx)
end

for _f in (:V, :B, :PSpace, :PCur, :POld, :DtOmegaInvQ)
    symf = "Prop2DAcoIsoDenQ_DEO2_FDTD_get" * string(_f)
    @eval $(_f)(prop::Prop2DAcoIsoDenQ_DEO2_FDTD) = unsafe_wrap(Array, ccall(($symf, libprop2DAcoIsoDenQ_DEO2_FDTD), Ptr{Float32}, (Ptr{Cvoid},), prop.p), size(prop), own=false)
end

abstract type Prop2DAcoIsoDenQ_DEO2_FDTD_Model end
struct Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V <: Prop2DAcoIsoDenQ_DEO2_FDTD_Model end
struct Prop2DAcoIsoDenQ_DEO2_FDTD_Model_B <: Prop2DAcoIsoDenQ_DEO2_FDTD_Model end
struct Prop2DAcoIsoDenQ_DEO2_FDTD_Model_VB <: Prop2DAcoIsoDenQ_DEO2_FDTD_Model end

propagateforward!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD) = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_TimeStep, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)
propagateadjoint!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD) = propagateforward!(prop) # self-adjoint

scale_spatial_derivatives!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD) =
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_ScaleSpatialDerivatives, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

# v in model-space
function forwardBornInjection!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_ForwardBornInjection_V, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"])
end

# v,b in model-space
function forwardBornInjection!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_VB, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_ForwardBornInjection_VB, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},     Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["b"], wavefields["pold"], wavefields["pspace"])
end

# b in model-space
function forwardBornInjection!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_B, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_ForwardBornInjection_B, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},     Ptr{Cfloat}),
         prop.p,     dmodel["b"], wavefields["pold"], wavefields["pspace"])
end

# v in model-space
function adjointBornAccumulation!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V, imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_AdjointBornAccumulation_V, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"])
end

# v,b in model-space
function adjointBornAccumulation!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_VB, imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_AdjointBornAccumulation_VB, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},     Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["b"], wavefields["pold"], wavefields["pspace"])
end

# b in model-space
function adjointBornAccumulation!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_B, imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_AdjointBornAccumulation_B, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},     Ptr{Cfloat}),
         prop.p,     dmodel["b"], wavefields["pold"], wavefields["pspace"])
end

# v in model-space with FWI wavefield separation
function adjointBornAccumulation!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V, imagingcondition::ImagingConditionWaveFieldSeparationFWI, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},      Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], 1)
 end

# v in model-space with RTM wavefield separation
function adjointBornAccumulation!(prop::Prop2DAcoIsoDenQ_DEO2_FDTD, modeltype::Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V, imagingcondition::ImagingConditionWaveFieldSeparationRTM, dmodel, wavefields)
    ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep, libprop2DAcoIsoDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},      Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], 0)
 end

function show(io::IO, prop::Prop2DAcoIsoDenQ_DEO2_FDTD)
    nx = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_getNx, libprop2DAcoIsoDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nz = ccall((:Prop2DAcoIsoDenQ_DEO2_FDTD_getNz, libprop2DAcoIsoDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    write(io, "Prop2DAcoIsoDenQ_DEO2_FDTD -- nx,nz; $nx,$nz")
end


using PsychoJulia
using GeometryTypes
using Random
using Test

@testset "States" begin
    fstate = PsychoJulia.FixationState(Rectangle(-10,-10,20,20),
                                       PsychoJulia.ClockTimer(time(), 5.0, false),Point2f0(0.0, 0.0))
    b =  PsychoJulia.is_inside(fstate, Point2f0(0.0, 0.0))
    @test b

    rstate = PsychoJulia.ResponseState(Rectangle(50.0,50.0, 20.0, 20.0),
                                       PsychoJulia.ClockTimer(time(), 5.0, false), Point2f0(0.0, 0.0))
    b = PsychoJulia.is_inside(rstate, Point2f0(51.0, 51.0))
    @test b

    estate = PsychoJulia.TrialEndState(PsychoJulia.ClockTimer(time(), 5.0, false))
    trans = PsychoJulia.StateTransitions([fstate, rstate, estate], [2 3 1;3 3 3], 1)
    done = false
    record = PsychoJulia.ExperimentRecord()
    ntrials = 0
    PsychoJulia.reset!(fstate.clock)
    while !done
        state = trans.states[trans.current]
        aa = PsychoJulia.get_area(state)
        PsychoJulia.set_position!(state,Point2f0(aa.x + 1.0, aa.y+1.0))
        PsychoJulia.next!(trans, record)
        ntrials = 0
        for ee in record.event
            if occursin("-> 3", ee)
                ntrials += 1
            end
        end
        done = ntrials >= 2
    end
    @test record.event == ["1 -> 2", "2 -> 3", "3 -> 1", "1 -> 2", "2 -> 3"]

    dd = diff(record.timestamp)
    #check that we get the correct timestamps to whitin 1ms precision
    @test dd ≈ [5.0, 5.0, 5.0, 5.0] atol=0.001
end