export CuDiag

const CuDiag{ElT,VecT} = Diag{ElT,VecT} where {VecT<:CuArray{ElT}}
const NonuniformCuDiagTensor{ElT,N,StoreT,IndsT} =
  Tensor{ElT,N,StoreT,IndsT} where {StoreT<:CuDiag}

CuArray(D::NonuniformCuDiagTensor) = CuArray(dense(D))

function NDTensors.dense(T::NonuniformCuDiagTensor{ElT}) where {ElT}
  R_data = CUDA.zeros(ElT, dim(inds(T)))
  diag_inds = diagind(reshape(R_data, dims(inds(T))), 0)
  R_data[diag_inds] .= data(store(T))
  return Tensor(Dense(R_data), inds(T))
end

function Base.promote_rule(
  ::Type{Diag{ElT2,VecT2}}, ::Type{CuDense}
) where {ElT2,VecT2<:Number}
  return promote_type(DenseT1, ElT2)
end

function Base.promote_rule(
  ::Type{<:Tensor{ElT1,N1,StoreT1}}, ::Type{<:Tensor{ElT2,N2,StoreT2}}
) where {ElT1,ElT2,N1,N2,StoreT1<:CuDense,StoreT2<:UniformDiag}
  ElT3 = promote_type(ElT1, ElT2)
  return Tensor{promote_type(ElT1, ElT2),N3,CuDense{ElT3,CuVector{ElT3}}} where {N3}
end
function Base.promote_rule(
  ::Type{<:Tensor{ElT1,N1,StoreT1}}, ::Type{<:Tensor{ElT2,N2,StoreT2}}
) where {ElT1,ElT2,N1,N2,StoreT1<:UniformDiag,StoreT2<:CuDense}
  ElT3 = promote_type(ElT1, ElT2)
  return Tensor{promote_type(ElT1, ElT2),N3,CuDense{ElT3,CuVector{ElT3}}} where {N3}
end

function Base.promote_rule(
  ::Type{<:UniformDiag{ElT1}}, ::Type{<:NonuniformDiag{ElT2,VecT2}}
) where {ElT1,ElT2,VecT2<:CuArray}
  ElT3 = promote_type(ElT1, ElT2)
  return NonuniformDiag{ElT3,CuVector{ElT3}}
end
function Base.promote_rule(
  ::Type{<:NonuniformDiag{ElT2,VecT2}}, ::Type{<:UniformDiag{ElT1}}
) where {ElT1,ElT2,VecT2<:CuArray}
  ElT3 = promote_type(ElT1, ElT2)
  return NonuniformDiag{ElT3,CuVector{ElT3}}
end
function Base.promote_rule(
  ::Type{DenseT1}, ::Type{Diag{ElT2,VecT2}}
) where {DenseT1<:CuDense,ElT2,VecT2<:Number}
  return promote_type(DenseT1, ElT2)
end

function contraction_output_type(
  TensorT1::Type{<:NonuniformCuDiagTensor}, TensorT2::Type{<:CuDenseTensor}, IndsR
)
  return similar_type(promote_type(TensorT1, TensorT2), IndsR)
end
function contraction_output_type(
  TensorT1::Type{<:CuDenseTensor}, TensorT2::Type{<:NonuniformCuDiagTensor}, IndsR
)
  return contraction_output_type(TensorT2, TensorT1, IndsR)
end

function contraction_output_type(
  TensorT1::Type{<:DiagTensor{<:Number,<:CuDiag}},
  TensorT2::Type{<:DiagTensor{<:Number,<:CuDiag}},
  IndsR::Type,
)
  return similar_type(promote_type(TensorT1, TensorT2), IndsR)
end

function contraction_output_type(
  TensorT1::Type{<:UniformDiagTensor},
  TensorT2::Type{<:DiagTensor{<:Number,<:CuDiag}},
  IndsR::Type,
)
  return similar_type(promote_type(TensorT1, TensorT2), IndsR)
end
function contraction_output_type(
  TensorT1::Type{<:DiagTensor{<:Number,<:CuDiag}},
  TensorT2::Type{<:UniformDiagTensor},
  IndsR::Type,
)
  return contraction_output_type(TensorT2, TensorT1, IndsR)
end

function zero_contraction_output(
  T1::UniformDiagTensor{ElT1,N1}, T2::CuDenseTensor{ElT2,N2}, indsR::IndsR
) where {ElT1,N1,ElT2,N2,IndsR}
  ElT3 = promote_type(ElT1, ElT2)
  dat = CUDA.zeros(ElT3, dim(indsR))
  return Tensor(Dense(dat), indsR)
end
function zero_contraction_output(
  T2::CuDenseTensor{ElT2,N2}, T1::UniformDiagTensor{ElT1,N1}, indsR::IndsR
) where {ElT1,N1,ElT2,N2,IndsR}
  return zero_contraction_output(T1, T2, indsR)
end

function zero_contraction_output(
  T1::TensorT1, T2::TensorT2, indsR
) where {TensorT1<:NonuniformDiagTensor,TensorT2<:NonuniformCuDiagTensor}
  ElT3 = promote_type(eltype(TensorT1), eltype(TensorT2))
  dat = CUDA.zeros(ElT3, length(data(store(T2))))
  return Tensor(Diag(dat), indsR)
end

function zero_contraction_output(
  T1::TensorT1, T2::TensorT2, indsR
) where {TensorT1<:UniformDiagTensor,TensorT2<:NonuniformCuDiagTensor}
  ElT3 = promote_type(eltype(TensorT1), eltype(TensorT2))
  dat = CUDA.zeros(ElT3, length(data(store(T2))))
  return Tensor(Diag(dat), indsR)
end
function zero_contraction_output(
  T1::TensorT1, T2::TensorT2, indsR
) where {TensorT2<:UniformDiagTensor,TensorT1<:NonuniformCuDiagTensor}
  return zero_contraction_output(T2, T1, indsR)
end

function zero_contraction_output(
  T1::TensorT1, T2::TensorT2, indsR
) where {TensorT1<:DiagTensor,TensorT2<:CuDenseTensor}
  ElT3 = promote_type(eltype(TensorT1), eltype(TensorT2))
  dat = CUDA.zeros(ElT3, length(data(store(T2))))
  return Tensor(Dense(dat), indsR)
end
function zero_contraction_output(
  T1::TensorT1, T2::TensorT2, indsR
) where {TensorT2<:DiagTensor,TensorT1<:CuDenseTensor}
  return zero_contraction_output(T2, T1, indsR)
end

function contraction_output(
  T1::UniformDiagTensor, T2::DiagTensor{Elt2,<:CuDiag}, indsR
) where {Elt2}
  return zero_contraction_output(T1, T2, indsR)
end
function contraction_output(
  T1::DiagTensor{Elt1,<:CuDiag}, T2::UniformDiagTensor, indsR
) where {Elt1}
  return contraction_output(T2, T1, indsR)
end

function contract!(
  C::CuDenseTensor{<:Number,NC},
  Clabels,
  A::UniformDiagTensor{<:Number,NA},
  Alabels,
  B::CuDenseTensor{<:Number,NB},
  Blabels,
) where {NC,NA,NB}
  Bstore = data(store(B))
  Astore = data(store(A))
  Cstore = data(store(C))
  return copyto!(Cstore, Astore .* Bstore)
end

function contract!(
  C::CuDenseTensor{<:Number,NC},
  Clabels,
  A::CuDenseTensor{<:Number,NA},
  Alabels,
  B::UniformDiagTensor{<:Number,NB},
  Blabels,
) where {NC,NA,NB}
  return contract!(C, Clabels, B, Blabels, A, Alabels)
end

function contract!(
  C::NonuniformCuDiagTensor{EltC,NC,IndsC},
  Clabels,
  A::UniformDiagTensor{EltA,NA,IndsA},
  Alabels,
  B::NonuniformCuDiagTensor{EltB,NB,IndsB},
  Blabels,
) where {EltC<:Number,EltB<:Number,EltA<:Number,NC,NB,NA,IndsA,IndsB,IndsC}
  Bstore = data(store(B))
  Astore = data(store(A))
  Cstore = data(store(C))
  return copyto!(Cstore, Astore .* Bstore)
end

function contract!(
  C::NonuniformCuDiagTensor{EltC,NC,IndsC},
  Clabels,
  B::NonuniformCuDiagTensor{EltB,NB,IndsB},
  Blabels,
  A::UniformDiagTensor{EltA,NA,IndsA},
  Alabels,
) where {EltC<:Number,EltB<:Number,EltA<:Number,NC,NB,NA,IndsA,IndsB,IndsC}
  Bstore = data(store(B))
  Astore = data(store(A))
  Cstore = data(store(C))
  return copyto!(Cstore, Astore .* Bstore)
end

function contract!(
  C::NonuniformCuDiagTensor{EltC,NC,IndsC},
  Clabels,
  B::NonuniformCuDiagTensor{EltB,NB,IndsB},
  Blabels,
  A::NonuniformCuDiagTensor{EltA,NA,IndsA},
  Alabels,
) where {EltC<:Number,EltB<:Number,EltA<:Number,NC,NB,NA,IndsA,IndsB,IndsC}
  Bstore = data(store(B))
  Astore = data(store(A))
  Cstore = data(store(C))
  return copyto!(Cstore, Astore .* Bstore)
end

# Dense * NonuniformCuDiag
function contract!(
  C::CuDenseTensor, Clabels, A::NonuniformCuDiagTensor, Alabels, B::CuDenseTensor, Blabels
)
  Astore = data(store(A))
  newAstore = CUDA.zeros(eltype(A), dims(inds(A))[1], dims(inds(A))[2])
  adi = diagind(newAstore, 0)
  newAstore[adi] .= Astore[:]
  newA = Tensor(Dense(vec(newAstore)), inds(A))
  return contract!(C, Clabels, newA, Alabels, B, Blabels)
end

function contract!(
  C::CuDenseTensor, Clabels, A::CuDenseTensor, Alabels, B::NonuniformCuDiagTensor, Blabels
)
  return contract!(C, Clabels, B, Blabels, A, Alabels)
end


adjacency_mat = [
    false true true true true true true true true false true true true true false false false true false true false true false false false false false false false false false true false false
    true false true true false false false true false false false false false true false false false true false true false true false false false false false false false false true false false false
    true true false true false false false true true true false false false true false false false false false false false false false false false false false true true false false false true false
    true true true false false false false true false false false false true true false false false false false false false false false false false false false false false false false false false false
    true false false false false false true false false false true false false false false false false false false false false false false false false false false false false false false false false false
    true false false false false false true false false false true false false false false false true false false false false false false false false false false false false false false false false false
    true false false false true true false false false false false false false false false false true false false false false false false false false false false false false false false false false false
    true true true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    true false true false false false false false false false false false false false false false false false false false false false false false false false false false false false true false true true
    false false true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true
    true false false false true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    true false false true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    true true true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true true
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true true
    false false false false false true true false false false false false false false false false false false false false false false false false false false false false false false false false false false
    true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true true
    true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true true
    true true false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true true
    false false false false false false false false false false false false false false false false false false false false false false false false false true false true false true false false true true
    false false false false false false false false false false false false false false false false false false false false false false false false false true false true false false false true false false
    false false false false false false false false false false false false false false false false false false false false false false false true true false false false false false false true false false
    false false false false false false false false false false false false false false false false false false false false false false false false false false false false false true false false false true
    false false true false false false false false false false false false false false false false false false false false false false false true true false false false false false false false false true
    false false true false false false false false false false false false false false false false false false false false false false false false false false false false false false false true false true
    false false false false false false false false false false false false false false false false false false false false false false false true false false true false false false false false true true
    false true false false false false false false true false false false false false false false false false false false false false false false false false false false false false false false true true
    true false false false false false false false false false false false false false false false false false false false false false false false true true false false true false false false true true
    false false true false false false false false true false false false false false true true false false true false true false true true false false false false false true true true false true
    false false false false false false false false true true false false false true true true false false true true true false true true false false true true true true true true true false
   ]


# addprocs(3)

push!(LOAD_PATH,"src")
using DalitzPlotAnalysis

v = @parallel for i in 1:50
    sleep(1)
    val = DalitzPlotAnalysis.λ(1,3,4.0)
    writedlm("/tmp/test-$(i).txt", fill(val, 20))
end

fetch(v[1])

BootstrapResults = let Nb = 500
    # path to save
    path_to_tmp_res = "/tmp"
    # path to save
    minpars0 = vcat(readdlm("minpars_compass_$(mass_bin_name).txt")...);
    res = Matrix{Float64}(Nb,length(minpars0))
    llh = Vector{Float64}(Nb)
    # b = 16
    @parallel for b in 1:Nb
        @show b, "progress is ", b/Nb
        writedlm(joinpath(path_to_tmp_res,"BootstrapResults-$(b).txt"), res[b,:])
        writedlm(joinpath(path_to_tmp_res,"llh-$(b).txt"), llh[b])
    end
    res
end


using MuZero

include("../diagnostics/tac/game.jl")

l1 = Learner(Tac)
l1.config.save_interval = 1
l1.config.checkpoint_interval = 1
selfplay!(l1; limit=10)
fit!(l1, steps=1)

batch_tensor = randtrainingbatch(l1) |> tensor

train_net_initial_out1a = initial_inference(l1.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1b = initial_inference(l1.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1a == train_net_initial_out1b

save!(l1)

l2 = Learner(Tac; uuid=l1.uuid)
train_net_initial_out2 = initial_inference(l2.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1a == train_net_initial_out2

test_net = l1.train
import BSON
using Flux
BSON.@save "test_net.bson" test_net

d = BSON.load("test_net.bson")
test_net_same = d[:test_net]

test_out1 = initial_inference(test_net, batch_tensor.images) |> NetworkOutput
test_out2 = initial_inference(test_net_same, batch_tensor.images) |> NetworkOutput

using Luxor
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
learner = Learner(TicTacToe)
selfplay!(learner; limit=10)
game = learner |> randgame
game |> show
game |> svg
batch = learner |> randtrainingbatch
batch |> show

function cell_fn(x,y)
    bg_color = (2,2) == (x,y) ? "#CFE2FC" : nothing
    opacity = (1,1) == (x,y) ? 0.5 : nothing
    tr = (3,3) == (x,y) ? (value="●", color="green", size=18) : "tr"
    (value="O", color="red", bl="bl", tr=tr, bg_color=bg_color, opacity=opacity)
end
svg1 = GridSVG(Grid(3, 3, cell_fn), bg_color="#fff", title="State 0", header="to play: 1", canvas_height=310, canvas_width=240)
#cell_fn_2(x,y) = (value="X", color="#333", bl="bl", tr="tr")
svg2 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="UCB Policy", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
svg3 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="target", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
svg4 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="target", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
mat = [[svg1,svg2,svg3]]
g = GridSVG(Grid(3, 1), cellsize=200, show_border=false, canvas_height=300, canvas_width=600)
#g((x,y)->"X")
#function test_fn(x,y)
#    GridSVG(Grid(2, 2), bg_color="#ccc", canvas_height=20, canvas_width=20)
#end
#g(test_fn)
results_fn(x,y) = mat[y][x]
g(results_fn)

using Base64
using Serialization

base64encode(serialize, rand(UInt128))

# Dataframes Test
using MuZero
import Zygote
using DataFrames
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
config = make_tictactoe_config()
config.batch_size = 10
config.num_simulations = 10
config.num_unroll_steps = 1

replay_buffer = ReplayBuffer(config)
selfplay!(config, storage, replay_buffer; limit=10)
b = batch(replay_buffer, config)
state(b, 1)

using Serialization, Base64
rand(1:2, 3, 3) |> (x)->base64encode(serialize, x)
deserialize(IOBuffer(base64decode(test)))

df = DataFrame(
    "Inf Type"=>["Initial","Recurrent"],
    "State"=>[
        base64encode(serialize, rand(UInt32)),
        base64encode(serialize, rand(UInt32))],
    "Action"=>[nothing, 3],
    ""=>["=>","=>"],
    "Reward"=>Float32.([0,0]),
    "Value"=>Float32.([-1,1]),
    "π[1]"=>Float32.([0.08, 0.12]),
    "π[2]"=>Float32.([0.12, 0.12]),
    "π[3]"=>Float32.([0.20, 0.00]),
    "π[4]"=>Float32.([0.08, 0.20]),
    "π[5]"=>Float32.([0.00, 0.00]),
    "π[6]"=>Float32.([0.08, 0.12]),
    "π[7]"=>Float32.([0.12, 0.20]),
    "π[8]"=>Float32.([0.12, 0.12]),
    "π[9]"=>Float32.([0.20, 0.12]),
)
show(df)

# Full Training Test (needs GPU for reasonable training time)using MuZero
using MuZero
import Zygote
using Grids
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
using Base.Threads: @spawn, @threads
config = make_tictactoe_config()
config.batch_size = 128
config.num_actors = 3
config.weight_decay = 1e-4
config.lr_init=0.002f0
config.save_interval = 10000
config.num_simulations = 100
config.num_unroll_steps = 1

replay_buffer = ReplayBuffer(config)
losses, dynamics_stats = [], []
for i = 1:(config.num_actors)
    @spawn selfplay!(config, storage, replay_buffer)
end
policy_perfs = []; calc_policy_perf() = push!(policy_perfs, policy_perf())
value_perfs = [] ; calc_value_perf() = push!(value_perfs, value_perf())
reward_perfs = [] ; calc_reward_perf() = push!(reward_perfs, reward_perf())
cb = [(500, 1, [calc_reward_perf,calc_value_perf,calc_policy_perf]),]
@spawn fit!(config, storage, replay_buffer, losses=losses, callbacks=cb)

# Half Training Test (needs GPU for reasonable training time)
using MuZero
import Zygote
using Grids
using Plots
using Statistics
using Base.Threads: @spawn, @threads
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
config = make_tictactoe_config()
config.batch_size = 10
config.num_actors = 3
config.training_steps = 20
config.weight_decay = 1e-3
config.lr_init=0.002f0
config.save_interval = 10000
config.num_simulations = 10
config.num_unroll_steps = 1
config.training_steps = 12000

replay_buffer = ReplayBuffer(config)
losses, dynamics_stats = [], []
selfplay!(config, storage, replay_buffer; limit=32)
policy_perfs = []; calc_policy_perf() = push!(policy_perfs, policy_perf())
value_perfs = [] ; calc_value_perf() = push!(value_perfs, value_perf())
reward_perfs = [] ; calc_reward_perf() = push!(reward_perfs, reward_perf())

avgx(x,size) = mean.(collect(Iterators.partition(x, size)))

function plot_losses()
    plot(avgx(losses, 100), fmt = :png)
    savefig("losses-copy.png")
    cp("losses-copy.png", "losses.png", force=true)
    rm("losses-copy.png")
end

cb = [(100, 1, [plot_losses])]
res = @spawn fit!(config, storage, replay_buffer, losses=losses, callbacks=cb)
fetch(res)

using Grids

function cell_fn(x,y)
    "O"
end

s = Grid(3,3) |> Grids.svg
s(GridCellSVG(3, 3, "X"; val_bl="1"))
s(cell_fn)

g |> get_state |> svg

Learner(config)
Learner(config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)
Learner(;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de"
)
Learner(;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)

#
# Pre-existing path: No
# Pre-existing config: No
#
Learner(
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)

learner = Learner(
    Tac;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    config=config,
    schedule=tacschedule!
)


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

_iscontiguous(A::Array) = true
_iscontiguous(A::AbstractArray) = Base.iscontiguous(A)

_car_cdr_impl() = ()
_car_cdr_impl(x, y...) = (x, (y...,))
_car_cdr(tp::Tuple) = _car_cdr_impl(tp...)

function _all_lteq(A::AbstractArray, B::AbstractArray, C::AbstractArray)
    axes(A) == axes(B) == axes(C) || throw(DimensionMismatch("A, B and C must have the same indices"))
    all(x[1] <= x[2] <= x[3] for x in zip(A, B, C))
end

@inline function _all_lteq_impl(a::Real, B::AbstractArray, c::Real)
    result = 0
    @inbounds @simd for b in B
        result += ifelse(a <= b <= c, 1, 0)
    end
    result == length(eachindex(B))
end

_all_lteq(a::Real, B::AbstractArray, c::Real) = _all_lteq_impl(a, B, c)

"""
    @propagate_inbounds sum_first_dim(A::AbstractArray, j::Integer, ks::Integer...)

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

Calculate the equivalent of `sum(A[:, j, ks...])`.
"""
Base.@propagate_inbounds function sum_first_dim(A::AbstractArray, j::Integer, ks::Integer...)
    s = zero(eltype(A))
    @boundscheck if !Base.checkbounds_indices(Bool, Base.tail(axes(A)), (j, ks...))
        throw(BoundsError(A, (:, j)))
    end
    @inbounds for i in axes(A, 1)
        s += A[i, j, ks...]
    end
    s
end

"""
    @propagate_inbounds sum_first_dim(A::AbstractArray)

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

If `A` is a vector, return `sum(A)`, else `sum(A, 1)[:]`.
"""
sum_first_dim(A::AbstractArray) = sum(A, 1)[:]
sum_first_dim(A::AbstractVector) = sum(A)

const SingleArrayIndex = Union{Integer, CartesianIndex}

convert_eltype(::Type{T}, x::AbstractArray) where T = convert.(T, x)

convert_eltype(::Type{T}, x::AbstractArray{T}) where T = x


module SolverInterface

using Base.Meta
using Compat
using Compat.SparseArrays
using Compat.LinearAlgebra

const methods_by_tag = Dict{Symbol,Vector{Symbol}}()

macro define_interface(args)
    @assert args.head == :block
    code = quote end
    for line in args.args
        (isexpr(line, :line) || isa(line, LineNumberNode)) && continue
        (isa(line,Symbol) || isexpr(line,:tuple)) || error("Unexpected code in block")
        if isa(line,Symbol)
            fname = line
            tags = []
        else
            @assert isexpr(line,:tuple)
            fname = line.args[1]
            tags = line.args[2:end]
        end
        func = esc(fname)
        code = quote
            $code
            $(func)() = throw(MethodError($(func),()))
            function $(func)(::Int)
                q = Any[1]
                return q[1] # infer Any return type, workaround for julia issue #9479
            end
        end
        for t in tags
            t = quot(t)
            code = quote
                $code
                if haskey(methods_by_tag,$t)
                    push!(methods_by_tag[$t],$(quot(fname)))
                else
                    methods_by_tag[$t] = [$(quot(fname))]
                end
            end
        end
        push!(code.args, Expr(:export, func))
    end
    return code
end

abstract type AbstractMathProgModel end
export AbstractMathProgModel

# immutable type which we dispatch solvers on
abstract type AbstractMathProgSolver end
export AbstractMathProgSolver

# create dummy method to define function so that we can attach methods in other modules
# these are the common methods for AbstractMathProgModel
@define_interface begin
    getsolution
    getobjval
    optimize!
    status
    getobjbound, rewrap
    getobjgap, rewrap
    getrawsolver, rewrap
    getsolvetime, rewrap
    setsense!
    getsense
    numvar
    numconstr
    freemodel!
    setvartype!
    getvartype
    loadproblem!
    setwarmstart!
end

# solver parameters, may be implemented by AbstractMathProgModel or AbstractMathProgSolver
@define_interface begin
    setparameters!
end

include("LinearQuadratic.jl")
include("callbacks.jl")
include("Nonlinear.jl")
include("Conic.jl")

# Solver conversion routines
include("conic_to_lpqp.jl")
include("lpqp_to_conic.jl")
include("nonlinear_to_lpqp.jl")

end


"""
    bet(
        mag::AbstractArray{<:Real, 3},
        vsz::NTuple{3, Real},
        betargs::AbstractString = "-m -n -f 0.5"
    ) -> Array{Bool, 3}

Interface to FSL's bet.

### Arguments
- `mag::AbstractArray{<:Real, 3}`: magnitude image
- `vsz::NTuple{3, Real}`: voxel size
- `betargs::AbstractString = "-m -n -f 0.5"`: bet options

### Returns
- `Array{Bool, 3}`: binary brain mask
"""
function bet(
    mag::AbstractArray{<:Real, 3},
    vsz::NTuple{3, Real},
    betargs::AbstractString = "-m -n -f 0.5"
)
    if Sys.which("bet") === nothing
        error("bet: command not found")
    end

    magfile = tempname(cleanup=false) * ".nii"
    maskfile = replace(magfile, ".nii" => "_mask.nii.gz")

    nii = NIVolume(mag, voxel_size=vsz)
    niwrite(magfile, nii)

    try
        # bet appends '_mask' to output filename regardless of user choice
        run(`bet $magfile $magfile $(split(betargs))`)
        nii = niread(maskfile)
    finally
        rm(magfile, force=true)
        rm(maskfile, force=true)
    end

    return Array{Bool, 3}(nii.raw)
end


module ArrayFire

using Cxx
using Base.Meta

import Base: rand, show, randn, ones, diag, eltype, size, elsize,
    sizeof, length, showarray, convert, ndims, lu, qr, svd, setindex!
import Cxx: CppEnum
export AFArray, chol!, constant, aftype, @gpu, @icxx_str, __current_compiler__

# If you have a crash, enable this
const AF_DEBUG = true

function init_library()
    if !haskey(ENV, "AFMODE")
        info("Setting Unified Backend by default")
        Libdl.dlopen("libaf", Libdl.RTLD_GLOBAL)
    else
        if ENV["AFMODE"] == "CPU"
            info("Setting CPU Backend")
            Libdl.dlopen("libafcpu", Libdl.RTLD_GLOBAL)
        elseif ENV["AFMODE"] == "OPENCL"
            info("Setting OPENCL Backend")
            Libdl.dlopen("libafopencl", Libdl.RTLD_GLOBAL)
        elseif ENV["AFMODE"] == "CUDA"
            info("Setting GPU Backend")
            Libdl.dlopen("libafcuda", Libdl.RTLD_GLOBAL)
        else
            info("Invalid value for environment variable AFMODE. Setting to default.")
            Libdl.dlopen("libaf", Libdl.RTLD_GLOBAL)
        end
    end
    if !haskey(ENV, "AFPATH")
       addHeaderDir("/usr/local/include/"; kind = C_System)
    else
        addHeaderDir(ENV["APATH"]; kind = C_System)
    end
end
init_library()

cxx"""
#include <arrayfire.h>
"""

setDevice(i) = icxx"af::setDevice($i);"

const f32 = CppEnum{:af_dtype,UInt32}(0) # 32-bit floating point values
const c32 = CppEnum{:af_dtype,UInt32}(1) # 32-bit complex floating point values
const f64 = CppEnum{:af_dtype,UInt32}(2) # 64-bit complex floating point values
const c64 = CppEnum{:af_dtype,UInt32}(3) # 64-bit complex floating point values
const b8  = CppEnum{:af_dtype,UInt32}(4) #  8-bit boolean values
const s32 = CppEnum{:af_dtype,UInt32}(5) # 32-bit signed integral values
const u32 = CppEnum{:af_dtype,UInt32}(6) # 32-bit unsigned integral values
const u8  = CppEnum{:af_dtype,UInt32}(7) #  8-bit unsigned integral values
const s64 = CppEnum{:af_dtype,UInt32}(8) # 64-bit signed integral values
const u64 = CppEnum{:af_dtype,UInt32}(9) # 64-bit unsigned integral values

aftype(::Type{Float32})          = f32
aftype(::Type{Complex{Float32}}) = c32
aftype(::Type{Float64})          = f64
aftype(::Type{Complex{Float64}}) = c64
aftype(::Type{Bool})             = b8
aftype(::Type{Int32})            = s32
aftype(::Type{UInt32})           = u32
aftype(::Type{UInt8})            = u8
aftype(::Type{Int64})            = s64
aftype(::Type{UInt64})           = u64

function jltype(af_dtype)
    if af_dtype == f32
        return Float32
    elseif af_dtype == c32
        return Complex{Float32}
    elseif af_dtype == f64
        return Float64
    elseif af_dtype == c64
        return Complex{Float64}
    elseif af_dtype == b8
        return Bool
    elseif af_dtype == s32
        return Int32
    elseif af_dtype == u32
        return UInt32
    elseif af_dtype == u8
        return UInt8
    elseif af_dtype == s64
        return Int64
    elseif af_dtype == u64
        return UInt64
    end
end

function af_promote{T,S}(::Type{T},::Type{S})
    if T == S
        return T
    elseif T == Complex{Float64} || S == Complex{Float64}
        return Complex{Float64}
    elseif T == Complex{Float32} || S == Complex{Float32}
        (T == Float64 || S == Float64) && return Complex{Float64}
        return Complex{Float32}
    elseif T == Float64 || S == Float64
        return Float64
    elseif T == Float32 || S == Float32
        return Float32
    elseif T == UInt64 || S == UInt64
        return UInt64
    elseif T == Int64 || S == Int64
        return Int64
    elseif T == UInt32 || S == UInt32
        return UInt32
    elseif T == Int32 || S == Int32
        return Int32
    elseif T == UInt8 || S == UInt8
        return UInt8
    elseif T == Bool || S == Bool
        return Bool
    else
        return Float32
    end
end

abstract AFAbstractArray{T,N} <: AbstractArray{T,4}

immutable AFArray{T,N} <: AFAbstractArray{T,N}
    array::vcpp"af::array"
    function AFArray(array::vcpp"af::array")
        ret = new(array)
        if AF_DEBUG && backend_eltype(ret) != T
            error("Tried to create AFArray{$T} with array of eltype $(backend_eltype(ret))")
        end
        ret
    end
end
immutable AFSubArray{T,N} <: AFAbstractArray{T,N}
    array::vcpp"af::array::array_proxy"
    function AFSubArray(array::vcpp"af::array::array_proxy")
        ret = new(array)
        if AF_DEBUG && backend_eltype(ret) != T
            error("Tried to create AFArray{$T} with array of eltype $(backend_eltype(ret))")
        end
        ret
    end
end

AFArray() = icxx"af::array();"

call{T,N}(::Type{AFAbstractArray{T,N}}, proxy::vcpp"af::array::array_proxy") =
    AFSubArray{T,N}(proxy)
call{T,N}(::Type{AFAbstractArray{T,N}}, arr::vcpp"af::array") =
    AFArray{T,N}(arr)
ndims(arr::vcpp"af::array") = icxx"$arr.numdims();"
ndims(arr::vcpp"af::array::array_proxy") = icxx"$arr.numdims();"
(::Type{AFArray{T}}){T}(arr::vcpp"af::array") = AFArray{T,Int(ndims(arr))}(arr)
(::Type{AFSubArray{T}}){T}(arr::vcpp"af::array::array_proxy") = AFSubArray{T,Int(ndims(arr))}(arr)

Cxx.cppconvert{T}(x::AFAbstractArray{T}) = x.array

convert{T,N}(::Type{AFArray{T,N}}, arr::AFSubArray{T,N}) =
    AFArray{T,N}(icxx"(af::array)$arr;")
(::Type{AFArray})(arr::AFSubArray) = AFArray{backend_eltype(arr), ndims(arr)}(arr)
convert(::Type{AFArray}, arr::AFSubArray) = AFArray(arr)

eltype{T}(x::AFAbstractArray{T}) = T
backend_eltype(x) = jltype(icxx"$x.type();")
sizeof{T}(a::AFAbstractArray{T}) = elsize(a) * length(a)

# GPU to Host
function convert{T,N}(::Type{Array{T,N}},x::AFAbstractArray{T,N})
    ret = Array(UInt8,sizeof(x))
    icxx"$x.host($(pointer(ret)));"
    ret = reinterpret(T, ret)
    ret = reshape(ret, size(x)...)
    ret
end
convert{T,N}(::Type{Array},x::AFAbstractArray{T,N}) = convert(Array{T,N},x)

# Host to GPU
function convert{T,N}(::Type{AFArray{T,N}}, x::Array{T,N})
    AFArray{T,N}(icxx"return af::array{$(dims_to_dim4(size(x))),$(pointer(x)),afHost};")
end
convert{T,N}(::Type{AFArray}, x::Array{T,N}) = convert(AFArray{T,N}, x)

# Show this array, by getting a data pointer to it and using julia's printing
# mechanism. ArrayFire copies internally anyway, so there is nothing to be gained
# by using it's printing
function showarray{T,N}(io::IO, X::AFAbstractArray{T,N};
                   header::Bool=true, kwargs...)
    header && print(io, summary(X))
    !isempty(X) && println(io,":")
    showarray(io, convert(Array{T,N},X); header = false, kwargs...)
end

Base.show(io::IO, arr::vcpp"af::array") = print(io, "ArrayFire Array")

af_print(X::AFAbstractArray) = icxx"""af::print("",$X);"""

# Translate dimensions. Note that we need to translate between 0 and 1 based
# indexing
ndims(dim4) = icxx"$dim4.ndims();"
ndims(A::AFAbstractArray) = ndims(icxx"$A.dims();")
function dims_to_dim4(dims)
    if length(dims) == 1
        icxx"af::dim4{$(dims[1])};"
    elseif length(dims) == 2
        icxx"af::dim4{$(dims[1]),$(dims[2])};"
    elseif length(dims) == 3
        icxx"af::dim4{$(dims[1]),$(dims[2]),$(dims[3])};"
    elseif length(dims) == 4
        icxx"af::dim4{$(dims[1]),$(dims[2]),$(dims[3]),$(dims[4])};"
    else
        error("Too many dimensions")
    end
end
function dim4_to_dims(dim4)
    d = ndims(dim4)
    t = (Int(icxx"(int)$dim4[0];"),)
    d == 1 && return t
    t = tuple(t...,Int(icxx"(int)$dim4[1];"))
    d == 2 && return t
    t = tuple(t...,Int(icxx"(int)$dim4[2];"))
    d == 3 && return t
    t = tuple(t...,Int(icxx"(int)$dim4[3];"))
    d == 4 && return t
    error("Bad dim4 object")
end

size(x::AFAbstractArray) = dim4_to_dims(icxx"($x).dims();")

#Specialize vectors and matrices
typealias AFMatrix{T} AFAbstractArray{T,2}
typealias AFVector{T} AFAbstractArray{T,1}

#Functions to generate array
import Base:fill, eye, diag
export iota
constant{T<:Real}(val::T, dims::Integer...) = AFArray{T}(af_constant(val, dims_to_dim4(dims), aftype(typeof(val))))
rand{T}(::Type{AFArray{T}}, dims::Integer...) = AFArray{T}(af_randu(dims_to_dim4(dims), aftype(T)))
randn{T}(::Type{AFArray{T}}, dims::Integer...) = AFArray{T}(af_randn(dims_to_dim4(dims), aftype(T)))
eye{T}(::Type{AFArray{T}}, dims::Integer...) = AFArray{T}(af_identity(dims_to_dim4(dims), aftype(T)))
diag{T}(x::AFArray{T}, k = 0) = AFArray{T}(af_diag(x, k))
getSeed() = af_getSeed()
setSeed(a::Integer) = af_setSeed(a)
range{T}(::Type{AFArray{T}}, dims::Integer...) = AFArray{T}(af_range(dims_to_dim4(dims), length(dims)-1, aftype(T)))
iota{T}(::Type{AFArray{T}}, dims::Integer...) = AFArray{T}(af_iota(dims_to_dim4(dims), dims_to_dim4(1), aftype(T)))
moddims{T}(a::AFAbstractArray{T}, dims::Integer...) = AFArray{T}(af_moddims(a, dims_to_dim4(dims)))

Base.rand{A<:AFArray}(::Type{A},sz::Tuple)=rand(A, sz...)

#TODO : make `tile` compatible with `repeat` in base

export tile

tile{T}(a::AFAbstractArray{T}, inds::Integer...) = AFArray{T}(af_tile(a, dims_to_dim4(inds)))

import Base: getindex

# Note that we need to translate between 0 and 1 based indexing
to_af_idx(x::Real) = convert(Cint, x-1)
to_af_idx(x::Range) = icxx"af::seq($(first(x))-1,$(last(x))-1,$(step(x)));"
to_af_idx(x::Colon) = icxx"af::span;"
to_af_idx(x::rcpp"af::seq") = x
const tis = to_af_idx

IS = Union{Real,Range,Colon,rcpp"af::seq"}

#ArrayFire needs to index with Int32s or Float32s. Also, array_proxy wrap for indexing with arrays crashes currently.
function _getindex(x::AFAbstractArray,y::AFAbstractArray)
    idx = y - 1
    out = AFArray()
    icxx"$out = $x($idx);"
    AFArray{backend_eltype(out)}(out)
end
# Avoid an ambiguity
_getindex{T}(x::AFAbstractArray{T},s0::Real)             = icxx"$x($(tis(s0)));"

_getindex{T}(x::AFAbstractArray{T},s0::IS)               = icxx"$x($(tis(s0)));"

_getindex{T}(x::AFAbstractArray{T},s0::IS,s1::IS)        = icxx"$x($(tis(s0)),
                                                                   $(tis(s1)));"

_getindex{T}(x::AFAbstractArray{T},s0::IS,s1::IS,s2::IS) = icxx"$x($(tis(s0)),
                                                                   $(tis(s1)),
                                                                   $(tis(s2)));"
_getindex{T}(x::AFAbstractArray{T},
                    s0::IS,s1::IS,s2::IS, s3::IS)        = icxx"$x($(tis(s0)),
                                                                   $(tis(s1)),
                                                                   $(tis(s2)),
                                                                   $(tis(s3)));"

#_getindex{T}(x::AFAbstractArray{T}, i) = icxx"$x($i)"

function getindex{T}(x::AFAbstractArray{T}, idxs...)
    proxy = _getindex(x,idxs...)
    if isa(proxy, AFArray)
        return proxy
    end
    num_elems = icxx"$proxy.elements();"
    if num_elems == 1
        return Array(AFSubArray{T}(proxy))[1]
    else
        return AFSubArray{T}(proxy)
    end
end

#TODO : return 0-element array when b is all falses
function getindex(x::AFAbstractArray, b::AFAbstractArray{Bool})
    out = AFArray()
    if any(b)
        icxx"$out = $x($b);"
        AFArray{backend_eltype(out)}(out)
    else
        return 0
    end
end

function setindex!{T}(x::AFAbstractArray{T}, val, idxs...)
    proxy = _getindex(x,idxs...)
    icxx"$proxy = $val;"
end

# Avoid crash when referencing with boolean arrays with all falses
function setindex!(x::AFAbstractArray, val, b::AFAbstractArray{Bool})
    if any(b)
        icxx"$x($b) = $val;"
    else
        return val
    end
end

# Exception handling
import Base: showerror
@exception function showerror(io::IO, e::rcpp"af::exception")
    println(io, bytestring(icxx"$e.what();"))
end

immutable AFFeatures
    feat::vcpp"af::features"
    function AFFeatures(f::vcpp"af::features")
        new(f)
    end
end
AFFeatures() = icxx"af::features();"

Cxx.cppconvert(f::AFFeatures) = f.feat

immutable AFWindow
    win::vcpp"af::Window"
    function AFWindow(w::vcpp"af::Window")
        new(w)
    end
end

Cxx.cppconvert(w::AFWindow) = w.win

#import other files
include("AFWrap.jl")
include("math.jl")
include("image.jl")
include("stats.jl")
include("vector.jl")
include("linalg.jl")
include("signal.jl")
include("graphics.jl")
include("unified.jl")

#Info
export af_info, af_isDoubleAvailable, af_sync, getDevice
getDevice() = af_getDevice()

#Storing and Loading Arrays
save(a::AFAbstractArray, path::AbstractString, key::AbstractString) = af_saveArray(key, a, path)
function load(path::AbstractString, key::AbstractString)
    out = af_readArray(path, key)
    AFArray{backend_eltype(out)}(out)
end

macro gpu(ex...)
    ex = ex[1]
    top = ex.args[1]
    bottom = ex.args[2]
    iterator = top.args[1]
    start = :($(top.args[2]).start)
    step = :(step($(top.args[2])))
    stop = :($(top.args[2]).stop)
    x = bottom.args[2]
    body = Expr(:->,:i,Expr(:block, x, :nothing))
    esc(Expr(:let, quote icxx"""
                    gfor (af::seq i, $start - 1, $stop - 1, $step)
                    {
                        $body(i);
                    }
                    """
                    end , :(body = $body), :(stop = $stop), :(start = $start), :(step = $step)))
end

end # module


"""
Individual representing the solution from the Pareto set.
"""
struct FrontierIndividual{F} <: ArchivedIndividual{F}
    fitness::F
    params::AbstractIndividual
    tag::Int                            # tag of the individual (e.g. gen.op. ID)
    num_fevals::Int                     # number of fitness evaluations so far
    n_restarts::Int                     # the number of method restarts so far
    timestamp::Float64                  # when archived

    FrontierIndividual{F}(fitness::F,
                   params, tag, num_fevals, n_restarts, timestamp=time()) where F =
        new{F}(fitness, params, tag, num_fevals, n_restarts, timestamp)

    FrontierIndividual(fitness::F,
                   params, tag, num_fevals, n_restarts, timestamp=time()) where F =
        new{F}(fitness, params, tag, num_fevals, n_restarts, timestamp)
end

tag(indi::FrontierIndividual) = indi.tag

"""
Individual stored in `EpsBoxArchive`.
"""
const EpsBoxFrontierIndividual{N,F<:Number} = FrontierIndividual{IndexedTupleFitness{N,F}}
const FrontierSpatialIndex{N,F<:Number} = SI.SpatialIndex{Int,N,EpsBoxFrontierIndividual{N,F}}
const FrontierRTree{N,F<:Number} = SI.RTree{Int,N,EpsBoxFrontierIndividual{N,F}}

# traits for spatial indexing
SI.mbrtrait(::Type{EpsBoxFrontierIndividual{N,F}}) where {N,F} = SI.HasMBR{SI.Rect{Int,N}}
SI.mbr(indi::EpsBoxFrontierIndividual{N,F}) where {N,F} = SI.Rect(indi.fitness.index, indi.fitness.index)
SI.idtrait(::Type{<:EpsBoxFrontierIndividual}) = SI.HasID{Int}
SI.id(indi::EpsBoxFrontierIndividual) = indi.num_fevals

EpsBoxFrontierIndividual(fitness::IndexedTupleFitness{N,F},
               params, tag, num_fevals, n_restarts, timestamp=time()) where {N,F} =
    FrontierIndividual(fitness, params, tag, num_fevals, n_restarts, timestamp)

"""
ϵ-box archive saves only the solutions that are not ϵ-box
dominated by any other solutions in the archive.

It also counts the number of candidate solutions that have been added
and how many ϵ-box progresses have been made.
"""
mutable struct EpsBoxArchive{N,F,FS<:EpsBoxDominanceFitnessScheme} <: Archive{IndexedTupleFitness{N,F},FS}
    fit_scheme::FS        # Fitness scheme used
    start_time::Float64   # Time when archive created, we use this to approximate the starting time for the opt...

    num_candidates::Int               # Number of calls to add_candidate!()
    best_front_elem::EpsBoxFrontierIndividual{N,F} # best frontier element: the candidate with the best aggregated fitness
    last_progress::Int                # when (wrt num_candidates) last ϵ-progress has occured
    last_restart::Int                 # when (wrt num_dlast) last restart has occured
    n_restarts::Int                   # the counter of the method restarts
    n_oversize_inserts::Int           # how many times the candidates were inserted into oversized archive

    max_size::Int         # maximal frontier size
    # TODO allow different frontier containers?
    # see e.g. Altwaijry & Menai "Data Structures in Multi-Objective Evolutionary Algorithms", 2012
    frontier::FrontierRTree{N,F}  # candidates along the fitness Pareto frontier

    EpsBoxArchive(fit_scheme::EpsBoxDominanceFitnessScheme{N,F};
                  max_size::Integer = 1_000_000,
                  leaf_capacity::Integer = 10,
                  branch_capacity::Integer = 10) where {N,F} =
        new{N,F,typeof(fit_scheme)}(fit_scheme, time(), 0,
                                    EpsBoxFrontierIndividual{N,F}(nafitness(fit_scheme), Individual(), 0, 0, 0, NaN), #FIXME does not respect the SearchSpace datatype
                                    0, 0, 0, 0, max_size,
                                    FrontierRTree{N,F}(leaf_capacity=leaf_capacity,
                                                       branch_capacity=branch_capacity))
end

EpsBoxArchive(fit_scheme::EpsBoxDominanceFitnessScheme, params::Parameters) =
    EpsBoxArchive(fit_scheme, max_size=params[:MaxArchiveSize],
                  leaf_capacity=params[:LeafCapacity],
                  branch_capacity=params[:BranchCapacity])

const EpsBoxArchive_DefaultParameters = ParamsDict(
    :MaxArchiveSize => 100_000,
    :LeafCapacity => 10,
    :BranchCapacity => 10
)

Base.eltype(::Type{EpsBoxArchive{N,F}}) where {N,F} = EpsBoxFrontierIndividual{N,F}

Base.length(a::EpsBoxArchive) = length(a.frontier)
Base.isempty(a::EpsBoxArchive) = isempty(a.frontier)
capacity(a::EpsBoxArchive) = a.max_size
numdims(a::EpsBoxArchive) = !isempty(a.frontier) ? length(a.frontier[1].params) : 0

"""
    pareto_frontier(a::EpsBoxArchive)

Get the iterator to the individuals on the Pareto frontier.
"""
pareto_frontier(a::EpsBoxArchive) = a.frontier

"""
    rand_front_elem(a::EpsBoxArchive)

Get random Pareto frontier element.

Returns `nothing` if frontier is empty.
"""
function rand_front_elem(a::EpsBoxArchive)
    isempty(a) && return nothing

    node = a.frontier.root
    while SI.level(node) > 1
        # descend to a random child
        node = node[rand(eachindex(SI.children(node)))]
    end
    return node[rand(eachindex(SI.children(node)))]
end

"""
    noprogress_streak(a::EpsBoxArchive, [since_restart])

Get the number of `add_candidate!()` calls since the last ϵ-progress.
If `since_restart` is specified, the number is relative to the last
restart.
"""
noprogress_streak(a::EpsBoxArchive; since_restart::Bool=false) =
    since_restart ?
        a.num_candidates - max(a.last_progress, a.last_restart) :
        a.num_candidates - a.last_progress

has_best_front_elem(a::EpsBoxArchive) = isfinite(a.best_front_elem.timestamp)
best_front_elem(a::EpsBoxArchive) = has_best_front_elem(a) ? a.best_front_elem : nothing
best_candidate(a::EpsBoxArchive) = has_best_front_elem(a) ? params(a.best_front_elem) : nothing
best_fitness(a::EpsBoxArchive) = has_best_front_elem(a) ? fitness(a.best_front_elem) : nafitness(a.fit_scheme)

function rebuild_frontier!(a::EpsBoxArchive)
    # bulk-reload R-tree (leave leaves underfilled to postpone reinserts)
    a.frontier = SI.load!(similar(a.frontier), a.frontier,
                          leaf_fill = floor(Int, a.frontier.fill_factor * SI.capacity(SI.Leaf, a.frontier)))
    return a
end

function notify!(a::EpsBoxArchive, event::Symbol)
    if event == :restart
        a.n_restarts += 1
        a.last_restart = a.num_candidates
        a.n_restarts == 1 && rebuild_frontier!(a)
    end
    return a
end

"""
    tagcounts(a::EpsBoxArchive)

Count the tags of individuals on the ϵ-box frontier.
Each restart the individual remains in the frontier discounts it by `θ`.

Returns the `tag`→`count` dictionary.
"""
function tagcounts(a::EpsBoxArchive, θ::Number = 1.0)
    (0.0 < θ <= 1.0) || throw(ArgumentError("θ ($θ) should be in (0.0, 1.0] range"))
    res = Dict{Int,Float64}()
    for indi in pareto_frontier(a)
        curtag = tag(indi)
        if curtag > 0
            curcounts = get!(res, curtag, 0.0)
            res[curtag] = curcounts+θ^(a.n_restarts-indi.n_restarts)
        end
    end
    return res
end

function add_candidate!(a::EpsBoxArchive{N,F}, cand_fitness::IndexedTupleFitness{N,F},
                        candidate, tag::Int=0, num_fevals::Int=-1) where {N,F}
    a.num_candidates += 1
    if num_fevals == -1
        num_fevals = a.num_candidates
    end
    #@debug "New fitness: $(cand_fitness.orig) agg=$(cand_fitness.agg)"
    #@debug "Params: $candidate"
    # TODO don't traverse the tree twice: some SpatialIndexing API to combine isempty() with the follow-up findfirst() call?
    if !isempty(a.frontier, DominanceCone{is_minimizing(a.fit_scheme)}(cand_fitness.index)) # candidate is dominated by frontier
        if length(a.frontier) <= a.max_size
            a.n_oversize_inserts = 0 # reset the counter since the size is ok
        end
        return a
    end
    # find the Pareto front element with exactly the same indexed fitness as the candidate
    frontix = findfirst(a.frontier, SI.Point(cand_fitness.index))
    if frontix !== nothing # found front element, no eps-progress
        front_leaf, frontel_pos = frontix
        frontel = front_leaf[frontel_pos]
        front_fitness = fitness(frontel)
        hat, index_match = hat_compare(cand_fitness, front_fitness, a.fit_scheme)
        @assert index_match # should have the same indexed fitness since that's how it was found
        if hat < 0 # new fitness dominates (but not eps-dominates) the one in archive, replace the element
            SI.children(front_leaf)[frontel_pos] = frontel =
                EpsBoxFrontierIndividual(cand_fitness, copyto!(frontel.params, candidate), tag, num_fevals, a.n_restarts)
        end
    else # eps-progress: non-dominated solution that has some eps-indices different from all current ones
        a.last_progress = a.num_candidates
        hat = -1
        # remove all dominated frontier elements, if any
        SI.subtract!(a.frontier, DominanceCone{!is_minimizing(a.fit_scheme)}(cand_fitness.index))
        #@debug "Appended non-dominated element to the frontier"
        frontel = EpsBoxFrontierIndividual(cand_fitness, copy(candidate), tag, num_fevals, a.n_restarts)
        insert!(a.frontier, frontel)
        if length(a.frontier) > a.max_size
            a.n_oversize_inserts += 1 # throw(error("Pareto frontier exceeds maximum size"))
        end
        # rebuild the frontier if it has changed significantly since the last rebuild
        if a.n_restarts > 1 && a.frontier.nelem_insertions > 3length(a.frontier)
            rebuild_frontier!(a)
        end
    end
    # check if the new candidate has better aggregate score
    if !has_best_front_elem(a) # initialize best element
        a.best_front_elem = frontel
    elseif hat < 0 # candidate replaced frontier element
        # update best_front_elem if it was replaced or if the new aggscore is better
        # NOTE: when the old best element is replaced, the new aggscore might be
        #       worse due to differences in aggscore and eps-progress metrics
        d = a.best_front_elem.fitness.agg - frontel.fitness.agg
        if  a.best_front_elem.fitness.index == frontel.fitness.index ||
           (d > zero(d) && is_minimizing(a.fit_scheme)) ||
           (d < zero(d) && !is_minimizing(a.fit_scheme))
            #@debug "New best candidate"
            a.best_front_elem = frontel
        end
    end
    if length(a.frontier) <= a.max_size
        a.n_oversize_inserts = 0 # reset the counter since the size is ok
    end
    return a
end

# actually this methods should never be called because the fitness
# is already indexes within the method
add_candidate!(a::EpsBoxArchive{N,F}, cand_fitness::NTuple{N,F},
               candidate::AbstractIndividual, tag::Int=0, num_fevals::Int=-1) where {N,F} =
    add_candidate!(a, archived_fitness(cand_fitness, a), candidate, tag, num_fevals)

# called by check_stop_condition(e::Evaluator, ctrl)
function check_stop_condition(a::EpsBoxArchive, p::OptimizationProblem, ctrl)
    if ctrl.max_steps_without_progress > 0 &&
        noprogress_streak(a, since_restart=false) > ctrl.max_steps_without_progress
        return "No epsilon-progress for more than $(ctrl.max_steps_without_progress) iterations"
    elseif a.n_oversize_inserts >= 10
        # notify that the last 10 inserts were to the oversized archive
        # that means that the ϵ quantization steps are too small
        return "Pareto frontier size ($(length(a.frontier))) exceeded maximum ($(a.max_size))"
    else
        return ""
    end
end


# TODO: clean up

struct PlanarContactSituation{T}
    active_body_poses::Dict{BodyID, Transform3D{T}}
    support_polygon::DConvexHull{T}
end

PlanarContactSituation{T}() where {T} = PlanarContactSituation(Dict{BodyID, Transform3D{T}}(), DConvexHull{T}())

struct ICPWalkingStateMachine{N, I<:ICPTrajectoryGenerator, S<:SE3PDController}
    contacts::Dict{BodyID, Vector{QPControl.ContactPoint{N}}}
    icp_trajectory_generator::I
    end_effector_controllers::Dict{BodyID, S}
    contact_plan::Piecewise{Constant{PlanarContactSituation{Float64}}, Float64, Vector{Constant{PlanarContactSituation{Float64}}}, Vector{Float64}}
    bodies::Vector{BodyID}
    # TODO: footstep generator
end

# TODO: not super nice:
isdone(trajectory::SE3Trajectory, t) = isdone(trajectory.angular, t) || isdone(trajectory.linear, t)
isdone(trajectory::Constant, t) = false
isdone(trajectory::Interpolated, t) = t > trajectory.xf
isdone(trajectory::BasicFootTrajectory, t) = t > trajectory.tf

function ICPWalkingStateMachine(
            mechanism::Mechanism,
            contacts::Dict{BodyID, <:Vector{<:QPControl.ContactPoint}},
            icp_trajectory_generator::ICPTrajectoryGenerator
        )
    T = Float64
    n = num_segments(icp_trajectory_generator)
    contact_situations = [PlanarContactSituation{Float64}() for i = 1 : n]
    times = zeros(T, n + 1)
    contact_plan = Piecewise(Constant.(contact_situations), times, clamp=true)
    bodies = sort!(collect(keys(contacts)), by = x -> x.value)

    # create end effector controllers
    end_effector_controllers = map(bodies) do bodyid
        body = findbody(mechanism, bodyid)
        bodyframe = default_frame(body)
        baseframe = root_frame(mechanism)
        gains = SE3PDGains(FramePDGains(bodyframe, PDGains(0.0, 20.0)), FramePDGains(bodyframe, PDGains(0.0, 0.0)))
        angulartraj = interpolated_orientation_trajectory(0.0, 0.0, one(Quat{T}), one(Quat{T}))
        lineartraj = convert(BasicFootTrajectory{T}, SVector(0.0, 0.0, 0.0), 0.0)
        trajectory = SE3Trajectory(bodyframe, baseframe, angulartraj, lineartraj)
        weight = Diagonal(vcat(fill(10.0, 3), fill(10.0, 3)))
        BodyID(body) => SE3PDController(BodyID(root_body(mechanism)), BodyID(body), trajectory, weight, gains)
    end |> Dict

    statemachine = ICPWalkingStateMachine(
        contacts, icp_trajectory_generator, end_effector_controllers,
        contact_plan, bodies)
end

function enable_contacts!(statemachine::ICPWalkingStateMachine, bodyid::BodyID)
    for contact in statemachine.contacts[bodyid]
        contact.maxnormalforce[] = 1e6 # TODO
    end
end

function disable_contacts!(statemachine::ICPWalkingStateMachine, bodyid::BodyID)
    for contact in statemachine.contacts[bodyid]
        disable!(contact)
    end
end

# TODO: temporary
function init_footstep_plan!(statemachine::ICPWalkingStateMachine, state::MechanismState, foot_polygons)
    T = Float64
    body_poses = Dict{BodyID, Transform3D{T}}()
    for bodyid in keys(foot_polygons)
        body_pose = transform_to_root(state, bodyid)
        body_poses[bodyid] = body_pose
    end
    Δt_first_transfer = 2.0
    Δt_transfer = 0.55
    Δt_swing = 0.75
    step_length = 0.4

    ts = breaks(statemachine.contact_plan)
    ts[1] = 0
    contact_situations = subfunctions(statemachine.contact_plan)
    n = length(contact_situations)
    support_bodies = Set(keys(foot_polygons))
    next_swing_id_iterator = Iterators.cycle(statemachine.bodies)
    next_swing_id, next_swing_id_state = iterate(next_swing_id_iterator)
    first_step = true

    for i = 1 : n
        contact_situation = contact_situations[i].value
        foot_polygons_world = valtype(foot_polygons)[]
        for bodyid in support_bodies
            foot_polygon = foot_polygons[bodyid]
            body = findbody(state.mechanism, bodyid)
            sole_to_world = body_poses[bodyid] * frame_definition(body, foot_polygon.frame)
            push!(foot_polygons_world, transform(foot_polygon, sole_to_world))
        end
        support_points = collect(Iterators.flatten(vertices(unwrap(foot_polygon_world)) for foot_polygon_world in foot_polygons_world))
        jarvis_march!(contact_situation.support_polygon, support_points)
        empty!(contact_situation.active_body_poses)
        for body in support_bodies
            push!(contact_situation.active_body_poses, body => body_poses[body])
        end
        if length(support_bodies) == length(foot_polygons)
            Δt = first_step ? Δt_first_transfer : Δt_transfer
            delete!(support_bodies, next_swing_id)
        else
            Δt = Δt_swing
            foot_displacement = if first_step
                first_step = false
                SVector(step_length, 0.0, 0.0)
            else
                SVector(2 * step_length, 0.0, 0.0)
            end
            foot_pose = body_poses[next_swing_id]
            body_poses[next_swing_id] = Transform3D(foot_pose.from, foot_pose.to, rotation(foot_pose), translation(foot_pose) + foot_displacement)
            next_swing_id, next_swing_id_state = iterate(next_swing_id_iterator, next_swing_id_state)
            for bodyid in keys(foot_polygons)
                push!(support_bodies, bodyid)
            end
        end
        ts[i + 1] = ts[i] + Δt
    end
    ξ0 = horizontal_projection(icp(state).v)
    p0 = ξ0
    ξf = centroid(last(contact_situations).value.support_polygon)
    generator = statemachine.icp_trajectory_generator
    empty!(generator)
    initialize!(generator, p0, ξ0, ξf)
    for i in eachindex(contact_situations)
        Δt = ts[i + 1] - ts[i]
        push_segment!(generator, Δt, contact_situations[i].value.support_polygon)
    end
    solve!(statemachine.icp_trajectory_generator)

    for end_effector_controller in values(statemachine.end_effector_controllers)
        init_support!(end_effector_controller; t0=ts[1], tf=ts[2])
    end
end

function (statemachine::ICPWalkingStateMachine)(t, state::MechanismState)
    active_body_poses = statemachine.contact_plan(t).active_body_poses

    # Enable/disable contacts
    for bodyid in statemachine.bodies
        if bodyid in keys(active_body_poses)
            enable_contacts!(statemachine, bodyid)
        else
            disable_contacts!(statemachine, bodyid)
        end
    end

    # Gain scheduling / trajectory initialization
    contact_plan = statemachine.contact_plan
    for (bodyid, end_effector_controller) in statemachine.end_effector_controllers
        if isdone(end_effector_controller.trajectory[], t)
            t′ = clamp(t, first(breaks(contact_plan)), last(breaks(contact_plan)))
            n = length(subfunctions(contact_plan))
            index = min(searchsortedlast(breaks(contact_plan), t′), n)
            tf = breaks(contact_plan)[min(index + 1, n)]
            if bodyid in keys(active_body_poses)
                init_support!(end_effector_controller; t0=t′, tf=tf)
            else
                previous_pose = subfunctions(contact_plan)[index - 1].value.active_body_poses[bodyid]
                next_pose = subfunctions(contact_plan)[index + 1].value.active_body_poses[bodyid]
                init_swing!(end_effector_controller, previous_pose, next_pose; t0=t′, tf=tf)
            end
        end
    end
end


# -----------------------------------------------------------------------------
# Implementation of Atomistic MolecularDynamicsResult API
# -----------------------------------------------------------------------------

"""
    MollyResult <: MolecularDynamicsResult

The result generated from running an `MollySimulator`.

**Field descriptions**
- `system::System` the `System` resulting from the simulation
- `simulator::MollySimulator` the `MollySimulator` that produced the result
"""
struct MollyResult <: MolecularDynamicsResult
    system::System
    simulator::MollySimulator
end

function get_time_range(result::MollyResult)
    s = result.simulator
    Δt = s.Δt * s.stride
    start = s.t₀ + Δt - s.Δt
    steps = s.steps ÷ s.stride
    start:Δt:start+Δt*steps
end
get_num_bodies(result::MollyResult) = length(result.system)
get_bounding_box(result::MollyResult) = bounding_box(result.system)
get_boundary_conditions(result::MollyResult) = boundary_conditions(result.system)

reference_temperature(result::MollyResult) = coupling_reference_temperature(result.simulator.coupling)
coupling_reference_temperature(coupling) = coupling.temperature
coupling_reference_temperature(::NoCoupling) = missing

get_positions(result::MollyResult) = position(result.system)
get_velocities(result::MollyResult) = velocity(result.system)
get_particles(result::MollyResult) = AtomsBase.Atom.(result.system.atoms_data, get_positions(result), get_velocities(result))
get_system(result::MollyResult) = result.system

get_positions(result::MollyResult, t::Integer) = result.system.loggers["c"].coords[t]
get_velocities(result::MollyResult, t::Integer) = result.system.loggers["v"].velocities[t]
get_particles(result::MollyResult, t::Integer) = AtomsBase.Atom.(result.system.atoms_data, get_positions(result, t), get_velocities(result, t))

temperature(result::MollyResult, t::Integer) = result.system.loggers["t"].temperatures[t]
kinetic_energy(result::MollyResult, t::Integer) = result.system.loggers["k"].energies[t]
potential_energy(result::MollyResult, t::Integer) = result.system.loggers["p"].energies[t]

"""
    visualize(result::MollyResult, filename::String; kwargs...)

Animate a `MollyResult` using Molly's `visualize` function and store the result in a file.
Any kwargs will be passed directly to the `visualize` function.

This function is only available when GLMakie is imported.
"""
Molly.visualize(result::MollyResult, filename::String; kwargs...) = visualize(result.system.loggers["c"], result.system.box_size, filename; kwargs...)


# biowulf.jl
# functions for use on the NIH Biowulf super computer

"""
makeswarm(;G::Int=2,cell="HCT116",swarmfile::String="fit",label="label",inlabel=label,timestamp="",nsets=1,datafolder::String="HCT116_testdata",fish= false,cycle=true,thresholdlow::Float64=0.,thresholdhigh::Float64=1e8,conds::String="MOCK",resultfolder::String= "fit_result",infolder=resultfolder,batchsize=1000,maxtime = 60.,nchains::Int=2,nthreads::Int=1,transient::Bool=false,fittedparam=collect(1:2*G-1),fixedeffects=(),juliafile::String="fitscript",root=".",samplesteps::Int=100000,warmupsteps=20000,annealsteps=0,temp=1.,tempanneal=100.,cv = 0.02,yieldprior=0.05)

makeswarm(genes::Vector;G::Int=2,cell="HCT116",swarmfile::String="fit",label="label",inlabel=label,timestamp="",nsets=1,datafolder::String="HCT116_testdata",fish=false,cycle=true,conds::String="MOCK",resultfolder::String="fit_result",infolder=resultfolder,batchsize=1000,maxtime=60.,nchains::Int=2,nthreads::Int=1,transient::Bool=false,fittedparam=collect(1:2*G-1),fixedeffects=(),juliafile::String="fitscript",root=".",samplesteps::Int=100000,warmupsteps=20000,annealsteps=0,temp=1.,tempanneal=100.,cv=0.02,yieldprior=0.05)

Arguments
    - `G`: number of gene states
    - `cell': cell type for halflives and allele numbers
    - `infolder`: name of folder for initial parameters
    - `swarmfile`: name of swarmfile to be executed by swarm
    - `label`: label of output files produced
    - `inlabel`: label of files used for initial conditions
    - `nsets`: number of histograms to be fit (e.g. one for wild type and one for perturbation)
    - `datafolder`: folder holding histograms, if two folders use `-` (hyphen) to separate, e.g.  "data\folder1-data\folder2"
    - `fish`: Data file type, set to true for FISH and false to scRNA (FISH type assumes different folder structure)
    - `thresholdlow`: lower threshold for halflife for genes to be fit
    - `threhsoldhigh`: upper threshold
    - `conds`: string describing conditions to be fit with `-` to separate if two conditions, e.g. "WT-AUXIN"
    - `result`: folder for results
    - `batchsize`: number of jobs per swarmfile, default = 1000
    - `maxtime`: maximum wall time for run, default = 2 hrs
    - `nchains`: number of MCMC chains = number of processors called by Julia, default = 2
    - 'nthreads`: number of Julia threads per processesor, default = 1
    - `transient::Bool`: true means fit transient model (T0, T30, T120)
    - `fittedparam`: vector of rate indices to be fit, e.g. [1,2,3,5,6,7]
    - `fixedeffects`: tuple of vectors of rates that are fixed between control and treatment where first index is fit and others are fixed to first, e.g. ([3,8],) means  index 8 is fixed to index 3
         (each vector in tuple is a fixed rate set)
    - `juliafile`: name of file to be called by julia in swarmfile
    - `root`: name of root directory for project, e.g. "scRNA\"
    - `samplesteps`: number of MCMC sampling steps
    - `warmupsteps`: number of MCMC warmup steps to find proposal distribution covariance
    - `annealsteps`: number of annealing steps (during annealing temperature is dropped from tempanneal to temp)
    - `temp`: MCMC temperature
    - `tempanneal`: annealing temperature
    - `cv`: coefficient of variation (mean/std) of proposal distribution
    - `yieldprior`: prior for yield, default = .05, (set to 1 for FISH if using scRNA data format for FISH data)

returns swarmfile that calls a julia file that is executed on biowulf

"""

function makeswarm(;G::Int=2,cell="HCT116",swarmfile::String="fit",label="label",inlabel=label,timestamp="",nsets=1,datafolder::String="HCT116_testdata",fish= false,cycle=true,thresholdlow::Float64=0.,thresholdhigh::Float64=1e8,conds::String="MOCK",resultfolder::String= "fit_result",infolder=resultfolder,batchsize=1000,maxtime = 60.,nchains::Int=2,nthreads::Int=1,transient::Bool=false,fittedparam=collect(1:2*G-1),fixedeffects=(),juliafile::String="fitscript",root=".",samplesteps::Int=100000,warmupsteps=20000,annealsteps=0,temp=1.,tempanneal=100.,cv = 0.02,yieldprior=0.05)
    if occursin.("-",conds)
        cond = string.(split(conds,"-"))
    else
        cond = conds
    end
    genes = checkgenes(root,cond,datafolder,cell,thresholdlow,thresholdhigh)
    makeswarm(genes,G=G,cell=cell,infolder=infolder,swarmfile=swarmfile,label=label,inlabel=inlabel,timestamp=timestamp,nsets=nsets,datafolder=datafolder,fish=fish,cycle=cycle,conds=conds,resultfolder=resultfolder,batchsize=batchsize,maxtime=maxtime,nchains=nchains,nthreads=nthreads,transient=transient,fittedparam=fittedparam,fixedeffects=fixedeffects,juliafile=juliafile,root=root,samplesteps=samplesteps,warmupsteps=warmupsteps,annealsteps=annealsteps,temp=temp,tempanneal=tempanneal,cv=cv,yieldprior=yieldprior)
end

function makeswarm(genes::Vector;G::Int=2,cell="HCT116",swarmfile::String="fit",label="label",inlabel=label,timestamp="",nsets=1,datafolder::String="HCT116_testdata",fish=false,cycle=true,conds::String="MOCK",resultfolder::String="fit_result",infolder=resultfolder,batchsize=1000,maxtime=60.,nchains::Int=2,nthreads::Int=1,transient::Bool=false,fittedparam=collect(1:2*G-1),fixedeffects=(),juliafile::String="fitscript",root=".",samplesteps::Int=100000,warmupsteps=20000,annealsteps=0,temp=1.,tempanneal=100.,cv=0.02,yieldprior=0.05)
    if label == "label"
        if fish
            label = "FISH-ss"
        else
            label = "scRNA-ss"
        end
        label = label * "-" * cell
        if ~isempty(timestamp)
            label = label * "-" * timestamp
        end
        label = label * "_" * conds
        inlabel = label
    end
    ngenes = length(genes)
    println("number of genes: ",ngenes)
    juliafile = juliafile * "_" * label * "_" * "$G" * ".jl"
    if ngenes > batchsize
        batches = getbatches(genes,ngenes,batchsize)
        for batch in eachindex(batches)
            sfile = swarmfile * "_" * label * "_" * "$G" * "_" * "$batch" * ".swarm"
            write_swarmfile(sfile,nchains,nthreads,juliafile,batches[batch])
        end
    else
        sfile = swarmfile * "_" * label * "_" * "$G" * ".swarm"
        write_swarmfile(sfile,nchains,nthreads,juliafile,genes)
    end
    write_fitfile(juliafile,nchains,cell,conds,G,maxtime,fittedparam,fixedeffects,infolder,resultfolder,datafolder,fish,cycle,inlabel,label,nsets,transient,samplesteps,warmupsteps,annealsteps,temp,tempanneal,root,cv,yieldprior)
end

"""
    fix(folder)

    Finds jobs that failed and writes a swarmfile for those genes

"""
fix(folder) = writeruns(fixruns(findjobs(folder)))

function fix_filenames(folder,old="scRNA-ss-",new="scRNA-ss_")
    files = readdir(folder)
    for file in files
        if occursin(old,file)
            nfile = replace(file,old => new)
            mv(joinpath(folder,file),joinpath(folder,nfile),force=true)
        end
    end
end

"""
    setup(rootfolder = "scRNA")

    Sets up the folder system prior to use
    Defaults to "scRNA"

"""
function rna_setup(root = ".")

    data = joinpath(root,"data")
    results = joinpath(root,"results")
    alleles = joinpath(data,"alleles")
    halflives = joinpath(data,"halflives")
    testdata = joinpath(data,"HCT116_testdata")

    if ~ispath(data)
        mkpath(data)
    end
    if ~ispath(results)
        mkpath(results)
    end
    if ~ispath(alleles)
        mkpath(alleles)
    end
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/alleles/CH12_alleles.csv","$alleles/CH12_alleles.csv")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/alleles/HCT116_alleles.csv","$alleles/HCT116_alleles.csv")
    if ~ispath(halflives)
        mkpath(halflives)
    end
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/halflives/ESC_halflife.csv","$halflives/ESC_halflife.csv")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/halflives/CH12_halflife.csv","$halflives/CH12_halflife.csv")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/halflives/HCT116_halflife.csv","$halflives/HCT116_halflife.csv")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/halflives/OcaB_halflife.csv","$halflives/OcaB_halflife.csv")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/halflives/aB_halflife.csv","$halflives/aB_halflife.csv")
    if ~ispath(testdata)
        mkpath(testdata)
    end
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/HCT116_testdata/CENPL_MOCK.txt","$testdata/CENPL_MOCK.txt")
    Downloads.download("https://raw.githubusercontent.com/nih-niddk-mbs/StochasticGene.jl/master/data/HCT116_testdata/MYC_MOCK.txt","$testdata/MYC_MOCK.txt")

end

"""
make_fitfile(fitfile,fittedparam,fixedeffects)

make the file the swarm file calls to execute julia code

"""
function write_fitfile(fitfile,nchains,cell,datacond,G,maxtime,fittedparam,fixedeffects,infolder,resultfolder,datafolder,fish,cycle,inlabel,label,nsets,runcycle,samplesteps,warmupsteps,annealsteps,temp,tempanneal,root,cv,yieldprior)
        f = open(fitfile,"w")
        s =   '"'
        write(f,"@everywhere using StochasticGene\n")
        write(f,"@time fit_rna($nchains,ARGS[1],$s$cell$s,$fittedparam,$fixedeffects,$s$datacond$s,$G,$maxtime,$s$infolder$s,$s$resultfolder$s,$s$datafolder$s,$fish,$cycle,$s$inlabel$s,$s$label$s,$nsets,$cv,$runcycle,$samplesteps,$warmupsteps,$annealsteps,$temp,$tempanneal,$s$root$s,$yieldprior)\n")
        close(f)
end

function write_fitfile_include(fitfile,nchains,cell,datacond,G,maxtime,fittedparam,fixedeffects,infolder,resultfolder,datafolder,fish,cycle,inlabel,label,nsets,runcycle,samplesteps,warmupsteps,annealsteps,temp,tempanneal,root,modulepath,cv,yieldprior)
        f = open(fitfile,"w")
        s =   '"'
        write(f,"@everywhere include($s$modulepath$s)\n")
        # write(f,"@everywhere using StochasticGene\n")
        write(f,"@time StochasticGene.fit_rna($nchains,ARGS[1],$s$cell$s,$fittedparam,$fixedeffects,$s$datacond$s,$G,$maxtime,$s$infolder$s,$s$resultfolder$s,$s$datafolder$s,$fish,$cycle,$s$inlabel$s,$s$label$s,$nsets,$cv,$runcycle,$samplesteps,$warmupsteps,$annealsteps,$temp,$tempanneal,$s$root$s,$yieldprior)\n")
        close(f)
end

function getbatches(genes,ngenes,batchsize)
    nbatches = div(ngenes,batchsize)
    batches = Vector{Vector{String}}(undef,nbatches+1)
    println(batchsize," ",nbatches+1)
    for i in 1:nbatches
        batches[i] = genes[batchsize*(i-1)+1:batchsize*(i)]
    end
    batches[end] = genes[batchsize*nbatches+1:end]
    return batches
end

function write_swarmfile(sfile,nchains,nthreads,juliafile,genes::Vector)
    f = open(sfile,"w")
    for gene in genes
        gene = check_genename(gene,"(")
        writedlm(f,["julia -t $nthreads -p" nchains juliafile gene])
        # writedlm(f,["julia -p" nchains juliafile nchains gene cell cond G maxtime infolder resultfolder datafolder fish inlabel label nsets runcycle transient fittedparam fixedeffects])
    end
    close(f)
end

function checkgenes(root,conds::Vector,datafolder,celltype::String,thresholdlow::Float64,thresholdhigh::Float64)
    genes = Vector{Vector}(undef,2)
    if occursin.("-",datafolder)
        datafolder = string.(split(datafolder,"-"))
        for i in 1:2
            genes[i] = checkgenes(root,conds[i],datafolder[i],celltype,thresholdlow,thresholdhigh)
        end
    else
        for i in 1:2
            genes[i] = checkgenes(root,conds[i],datafolder,celltype,thresholdlow,thresholdhigh)
        end
    end
    intersect(genes[1],genes[2])
end

function checkgenes(root,cond::String,datafolder,cell::String,thresholdlow::Float64,thresholdhigh::Float64)
    datafolder = joinpath("data",datafolder)
    genes = intersect(get_halflives(root,cell,thresholdlow,thresholdhigh), get_genes(root,cond,datafolder))
    alleles = get_alleles(root,cell)
    if ~isnothing(alleles)
        return intersect(genes,alleles)
    else
        return genes
    end
end

function checkgenes(cond::String,datafolder,cell::String,thresholdlow::Float64,thresholdhigh::Float64)
    genes = intersect(get_halflives(".",cell,thresholdlow,thresholdhigh), get_genes(cond,datafolder))
    alleles = get_alleles(".",cell)
    if ~isnothing(alleles)
        return intersect(genes,alleles)
    else
        return genes
    end
end

function get_halflives(root,cell,thresholdlow::Float64,thresholdhigh::Float64)
    file = get_file(root,"data/halflives",cell,"csv")
    get_halflives(file,thresholdlow,thresholdhigh)
end

function get_halflives(file,thresholdlow::Float64,thresholdhigh::Float64)
    genes = Vector{String}(undef,0)
    halflives = readdlm(file,',')
    for row in eachrow(halflives)
        if typeof(row[2]) <: Number
            if thresholdlow <= row[2] < thresholdhigh
                push!(genes,string(row[1]))
            end
        end
    end
    return genes
end

function get_alleles(root,cell)
    file = get_file(root,"data/alleles",cell,"csv")
    if ~isnothing(file)
        return readdlm(file,',')[2:end,1]
    else
        return nothing
    end
end

function get_file(root,folder,type,suffix)
    folder = joinpath(root,folder)
    if ispath(folder)
        files = readdir(folder)
        for file in files
            name = split(file,"_")[1]
            if occursin(suffix,file) && name == type
                path = joinpath(folder,file)
                return path
            end
        end
    else
        throw("folder $folder does not exist")
    end
    nothing
end

function findjobs(folder)
    files = readdir(folder)
    files = files[occursin.("swarm_",files)]
    for (i,file) in enumerate(files)
        files[i] = split(file,"_")[2]
    end
    unique(files)
end

function fixruns(jobs,message="FAILED")
    runlist = Vector{String}(undef,0)
    for job in jobs
        if occursin(message,read(`jobhist $job`,String))
            swarmfile = findswarm(job)
            list = readdlm(swarmfile,',')
            runs =  chomp(read(pipeline(`jobhist $job`, `grep $message`),String))
            runs = split(runs,'\n')
            println(job)
            for run in runs
                linenumber = parse(Int,split(split(run," ")[1],"_")[2]) + 1
                while linenumber < length(list)
                    a = String(list[linenumber])
                    linenumber += 1000
                    println(a)
                    push!(runlist,a)
                end
            end
        end
    end
    return runlist
end

function writeruns(runs,outfile="fitfix.swarm")

    f = open(outfile,"w")
    for run in runs
        writedlm(f,[run],quotes=false)
    end
    close(f)

end

function findswarm(job)
    sc = "Swarm Command"
    line = read(pipeline(`jobhist $job`, `grep $sc`),String)
    list = split(line," ")
    list[occursin.(".swarm",list)][1]
end

function collate(folders,type="rates_scRNA_T120")
    genes = Vector{Vector}(undef,0)
    for folder in folders
        files = readdir(folder)
        for file in files
            if occursin(type,file)
                contents=readdlm(file,',')
                for row in eachrow(contents)
                    row[1] .== genes
                end
            end
        end
    end
end

"""
    get_missing_genes(datafolder::String,folder::String,cell,type,label,cond,model)
"""

function get_missing_genes(datafolder::String,folder::String,cell,type,label,cond,model,root=".")
    genes = checkgenes(root,cond,datafolder,cell,0.,1e8)
    genes1=get_genes(folder,type,label,cond,model)
    union(setdiff(genes1,genes),setdiff(genes,genes1))
end

function get_missing_genes(genes::Vector,folder,type,label,cond,model)
    genes1=get_genes(folder,type,label,cond,model)
    union(setdiff(genes1,genes),setdiff(genes,genes1))
end

get_missing_genes(genes,genes1) = union(setdiff(genes1,genes),setdiff(genes,genes1))

function scan_swarmfiles(jobid,folder=".")
    if ~(typeof(jobid) <: String)
        jobid = string(jobid)
    end
    genes = Array{String,1}(undef,0)
    files = readdir(folder)
    files = files[occursin.(jobid,files)]
    for file in files
        genes = vcat(genes,scan_swarmfile(file))
    end
    return genes
end

function scan_swarmfile(file)
    genes = Array{String,1}(undef,0)
    contents = readdlm(file,'\t')
    lines = contents[occursin.("[\"",string.(contents))]
    for line in lines
        push!(genes,split.(string(line)," ")[1])
    end
    return genes
end

function scan_fitfile(file,folder=".")
    genes = Array{String,1}(undef,0)
    joinpath(folder,file)
    file = readdlm(file,'\t')
    for line in eachrow(file)
        push!(genes,line[4])
    end
    return genes
end


abstract type AbstractRenewableFormulation <: AbstractDeviceFormulation end
abstract type AbstractRenewableDispatchFormulation <: AbstractRenewableFormulation end
struct RenewableFullDispatch <: AbstractRenewableDispatchFormulation end
struct RenewableConstantPowerFactor <: AbstractRenewableDispatchFormulation end

########################### renewable generation variables #################################
function AddVariableSpec(
    ::Type{T},
    ::Type{U},
    ::PSIContainer,
) where {T <: ActivePowerVariable, U <: PSY.RenewableGen}
    return AddVariableSpec(;
        variable_name = make_variable_name(T, U),
        binary = false,
        expression_name = :nodal_balance_active,
        lb_value_func = x -> 0.0,
        ub_value_func = x -> PSY.get_max_active_power(x),
    )
end

function AddVariableSpec(
    ::Type{T},
    ::Type{U},
    ::PSIContainer,
) where {T <: ReactivePowerVariable, U <: PSY.RenewableGen}
    return AddVariableSpec(;
        variable_name = make_variable_name(T, U),
        binary = false,
        expression_name = :nodal_balance_reactive,
    )
end

####################################### Reactive Power constraint_infos #########################
function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ReactivePowerVariable},
    ::Type{T},
    ::Type{<:AbstractDeviceFormulation},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
) where {T <: PSY.RenewableGen}
    return DeviceRangeConstraintSpec(;
        range_constraint_spec = RangeConstraintSpec(;
            constraint_name = make_constraint_name(
                RangeConstraint,
                ReactivePowerVariable,
                T,
            ),
            variable_name = make_variable_name(ReactivePowerVariable, T),
            limits_func = x -> PSY.get_reactive_power_limits(x),
            constraint_func = device_range!,
            constraint_struct = DeviceRangeConstraintInfo,
        ),
    )
end

function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ReactivePowerVariable},
    ::Type{T},
    ::Type{<:RenewableConstantPowerFactor},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
) where {T <: PSY.RenewableGen}
    return DeviceRangeConstraintSpec(;
        custom_psi_container_func = custom_reactive_power_constraints!,
    )
end

function custom_reactive_power_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::Type{RenewableConstantPowerFactor},
) where {T <: PSY.RenewableGen}
    names = [PSY.get_name(d) for d in devices]
    time_steps = model_time_steps(psi_container)
    p_var = get_variable(psi_container, ACTIVE_POWER, T)
    q_var = get_variable(psi_container, REACTIVE_POWER, T)
    constraint_val = JuMPConstraintArray(undef, names, time_steps)
    assign_constraint!(psi_container, REACTIVE_RANGE, T, constraint_val)
    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(acos(PSY.get_power_factor(d)))
        constraint_val[name, t] =
            JuMP.@constraint(psi_container.JuMPmodel, q_var[name, t] == p_var[name, t] * pf)
    end
    return
end

function DeviceRangeConstraintSpec(
    ::Type{<:RangeConstraint},
    ::Type{ActivePowerVariable},
    ::Type{T},
    ::Type{<:AbstractRenewableDispatchFormulation},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    use_parameters::Bool,
    use_forecasts::Bool,
) where {T <: PSY.RenewableGen}
    if !use_parameters && !use_forecasts
        return DeviceRangeConstraintSpec(;
            range_constraint_spec = RangeConstraintSpec(;
                constraint_name = make_constraint_name(
                    RangeConstraint,
                    ActivePowerVariable,
                    T,
                ),
                variable_name = make_variable_name(ActivePowerVariable, T),
                limits_func = x -> (min = 0.0, max = PSY.get_active_power(x)),
                constraint_func = device_range!,
                constraint_struct = DeviceRangeConstraintInfo,
            ),
        )
    end

    return DeviceRangeConstraintSpec(;
        timeseries_range_constraint_spec = TimeSeriesConstraintSpec(;
            constraint_name = make_constraint_name(RangeConstraint, ActivePowerVariable, T),
            variable_name = make_variable_name(ActivePowerVariable, T),
            parameter_name = use_parameters ? ACTIVE_POWER : nothing,
            forecast_label = "max_active_power",
            multiplier_func = x -> PSY.get_max_active_power(x),
            constraint_func = use_parameters ? device_timeseries_param_ub! :
                              device_timeseries_ub!,
        ),
    )
end

########################## Addition to the nodal balances ##################################

function NodalExpressionSpec(
    ::Type{T},
    ::Type{<:PM.AbstractPowerModel},
    use_forecasts::Bool,
) where {T <: PSY.RenewableGen}
    return NodalExpressionSpec(
        "max_active_power",
        REACTIVE_POWER,
        use_forecasts ? x -> PSY.get_max_reactive_power(x) : x -> PSY.get_reactive_power(x),
        1.0,
        T,
    )
end

function NodalExpressionSpec(
    ::Type{T},
    ::Type{<:PM.AbstractActivePowerModel},
    use_forecasts::Bool,
) where {T <: PSY.RenewableGen}
    return NodalExpressionSpec(
        "max_active_power",
        ACTIVE_POWER,
        use_forecasts ? x -> PSY.get_max_active_power(x) : x -> PSY.get_active_power(x),
        1.0,
        T,
    )
end

##################################### renewable generation cost ############################
function AddCostSpec(
    ::Type{T},
    ::Type{U},
    ::PSIContainer,
) where {T <: PSY.RenewableDispatch, U <: AbstractRenewableDispatchFormulation}
    # TODO: remove once cost_function is required
    cost_function = x -> isnothing(x) ? 1.0 : PSY.get_variable(x)
    return AddCostSpec(;
        variable_type = ActivePowerVariable,
        component_type = T,
        variable_cost = cost_function,
        multiplier = OBJECTIVE_FUNCTION_NEGATIVE,
    )
end


using Unitful: @u_str
using UnitfulAtomic: austrip, auconvert

using .Calculators: AbstractCalculator, Calculator
using NQCModels: Model

abstract type AbstractSimulation{M} end

Base.broadcastable(sim::AbstractSimulation) = Ref(sim)

struct Simulation{M,Calc,T,Ttype,C} <: AbstractSimulation{M}
    temperature::Ttype
    cell::C
    atoms::Atoms{T}
    calculator::Calc
    method::M
end

"""
    Simulation(atoms::Atoms{T}, model::Model, method::M;
        temperature=0u"K", cell::AbstractCell=InfiniteCell()) where {M,S,T}

Simulation parameters that controls the types of atoms, interactions,
dynamics method, temperature and simulation cell.
"""
function Simulation(atoms::Atoms{T}, model::Model, method::M;
        temperature=0u"K", cell::AbstractCell=InfiniteCell()) where {M,T}
    calc = Calculator(model, length(atoms), T)
    Simulation(temperature, cell, atoms, calc, method)
end

struct RingPolymerSimulation{M,Calc,T,Ttype,C,B} <: AbstractSimulation{M}
    temperature::Ttype
    cell::C
    atoms::Atoms{T}
    calculator::Calc
    method::M
    beads::B
end

function RingPolymerSimulation(temperature, cell::AbstractCell,
        atoms::Atoms{T}, model::Model, method::M,
        n_beads::Integer, quantum_nuclei::Vector{Symbol}) where {M,T}

    if isempty(quantum_nuclei)
        beads = RingPolymers.RingPolymerParameters{T}(n_beads, get_temperature(temperature), length(atoms))
    else
        beads = RingPolymers.RingPolymerParameters{T}(n_beads, get_temperature(temperature), atoms.types, quantum_nuclei)
    end

    calc = Calculator(model, length(atoms), n_beads, T)
    RingPolymerSimulation(temperature, cell, atoms, calc, method, beads)
end

function RingPolymerSimulation(atoms::Atoms, model::Model, method::M, n_beads::Integer;
        temperature=0,
        cell::AbstractCell=InfiniteCell(), quantum_nuclei::Vector{Symbol}=Symbol[]) where {M}
    RingPolymerSimulation(temperature, cell, atoms, model, method, n_beads, quantum_nuclei)
end

NQCModels.nstates(sim::AbstractSimulation) = NQCModels.nstates(sim.calculator)
NQCModels.eachstate(sim::AbstractSimulation) = NQCModels.eachstate(sim.calculator)
NQCModels.nelectrons(sim::AbstractSimulation) = NQCModels.nelectrons(sim.calculator)
NQCModels.eachelectron(sim::AbstractSimulation) = NQCModels.eachelectron(sim.calculator)
NQCModels.mobileatoms(sim::AbstractSimulation) = NQCModels.mobileatoms(sim.calculator)
NQCModels.dofs(sim::AbstractSimulation) = NQCModels.dofs(sim.calculator)

NQCModels.ndofs(sim::AbstractSimulation) = NQCModels.ndofs(sim.calculator.model)
natoms(sim::AbstractSimulation) = length(sim.atoms)
RingPolymers.nbeads(sim::RingPolymerSimulation) = RingPolymers.nbeads(sim.beads)
masses(sim::AbstractSimulation) = sim.atoms.masses
masses(sim::AbstractSimulation, i::Int) = sim.atoms.masses[i]

function masses(sim::AbstractSimulation, I::CartesianIndex)
    masses(sim, I[2])
end

Base.size(sim::Simulation) = (ndofs(sim), natoms(sim))
Base.size(sim::RingPolymerSimulation) = (ndofs(sim), natoms(sim), RingPolymers.nbeads(sim))

function get_temperature(sim::Simulation, t::Real=0)
    t = auconvert(u"fs", t)
    get_temperature(sim.temperature, t)
end
function get_temperature(sim::RingPolymerSimulation, t::Real=0)
    t = auconvert(u"fs", t)
    get_temperature(sim.temperature, t) * length(sim.beads)
end
get_temperature(temperature::Number, t=0) = austrip(temperature)
get_temperature(temperature::Function, t=0) = austrip(temperature(t))

function Base.show(io::IO, sim::Simulation{M}) where {M}
    print(io, "Simulation{$M}:\n  ", sim.atoms, "\n  ", sim.calculator.model)
end

function Base.show(io::IO, sim::RingPolymerSimulation{M}) where {M}
    print(io, "RingPolymerSimulation{$M}:\n\n  ", sim.atoms, "\n\n  ", sim.calculator.model,
          "\n  with ", length(sim.beads), " beads.")
end


using Test, StatsBase, SearchModels, KCenters, KNearestCenters, TextSearch, TextClassification, CategoricalArrays
using Downloads, Random, MLDataUtils, JSON3, GZip

Random.seed!(1)
function train_test_split(corpus, labels, at=0.7)
    n = length(corpus)
    iX = shuffle!(collect(1:n))
    n_ = floor(Int, n * at)
    itrain, itest = iX[1:n_], iX[n_+1:end]
    corpus[itrain], labels[itrain], corpus[itest], labels[itest]
end

function folds_split(corpus, labels, folds=3)
    n = length(corpus)
    indexes = shuffle!(collect(1:n))
    folds = kfolds(indexes, folds)

    [(corpus[itrain], labels[itrain], corpus[itest], labels[itest]) for (itrain, itest) in folds]
end

@testset "microtc" begin
    !isfile("emo50k.json.gz") && Downloads.download("https://github.com/sadit/TextClassificationTutorial/raw/main/data/emo50k.json.gz", "emo50k.json.gz")
    labels = []
    corpus = []
    targets = ("♡", "💔")
    GZip.open("emo50k.json.gz") do f
        for line in eachline(f)
            tweet = JSON3.read(line)
            label = tweet["klass"]
            if label in targets
                push!(labels, tweet["klass"])
                push!(corpus, tweet["text"])
            end
        end
    end
    labels = categorical(labels)
    traincorpus, trainlabels, testcorpus, testlabels = train_test_split(corpus, labels, 0.7)
    folds = folds_split(traincorpus, trainlabels)

    function error_function(config::MicroTC_Config)
        S = Float64[]
        for (_traincorpus, _trainlabels, _testcorpus, _testlabels) in folds
            tc = MicroTC(config, _traincorpus, _trainlabels; verbose=true)
            #valX = vectorize_corpus(tc, _testcorpus)
            #ypred = predict.(tc, valX)
            ypred = predict_corpus(tc, _testcorpus)
            push!(S, recall_score(_testlabels, ypred, weight=:macro))
        end

        1.0 - mean(S)
    end

    for t in traincorpus[1:10]
        @show t
    end

    @show countmap(trainlabels)
    @show countmap(testlabels)

    space = MicroTC_ConfigSpace(
       textconfig=TextConfigSpace(
            qlist=[[4], [3], [5]],
            nlist=[[1], [1, 2], []],
            slist=[]
        ),
        # cls=KncConfigSpace(
        #     centerselection=[TextCentroidSelection()],
        #     kernel=[k_(AngleDistance()) for k_ in [DirectKernel]]
        # )
        #cls = KnnClassifierConfigSpace()
        #cls=LiblinearConfigSpace()
    )

    best_list = search_models(space, error_function, 32;
        maxpopulation=8,
        # search hyper-parameters
        bsize=2, mutbsize=8, crossbsize=8,
        tol=0.0, maxiters=30, verbose=true)

    for (i, b) in enumerate(best_list)
        @info "-- perf best_lists[$i]:", i, b[1], b[2]
    end

    cls = MicroTC(best_list[1][1], traincorpus, trainlabels)
    sc = classification_scores(testlabels, predict_corpus(cls, testcorpus))
    @info "*** Performance on test: " sc
    @test sc.accuracy > 0.6
end

# flush(stdout); flush(stderr)
# sort!(P, :score, rev=true)
# plot(P.score, label="", xlabel="rank", ylabel="score", markershape=:auto, title="RS performance on validation set") |> display


function get_ecef(obj::GLONASSAnalyticalPropagator)
    return jcall(obj, "getECEF", Frame, ())
end

function get_eci(obj::GLONASSAnalyticalPropagator)
    return jcall(obj, "getECI", Frame, ())
end

function get_frame(obj::GLONASSAnalyticalPropagator)
    return jcall(obj, "getFrame", Frame, ())
end

function get_glonass_orbital_elements(obj::GLONASSAnalyticalPropagator)
    return jcall(obj, "getGLONASSOrbitalElements", GLONASSOrbitalElements, ())
end

function get_mu(::Type{GLONASSAnalyticalPropagator})
    return jcall(GLONASSAnalyticalPropagator, "getMU", jdouble, ())
end

function propagate_in_ecef(obj::GLONASSAnalyticalPropagator, arg0::AbsoluteDate)
    return jcall(obj, "propagateInEcef", PVCoordinates, (AbsoluteDate,), arg0)
end

function reset_initial_state(obj::GLONASSAnalyticalPropagator, arg0::SpacecraftState)
    return jcall(obj, "resetInitialState", void, (SpacecraftState,), arg0)
end


# ******************************************************************************************
# Notices:
#
# Copyright © 2022 United States Government as represented by the Administrator of the
# National Aeronautics and Space Administration.  All Rights Reserved.
#
# Disclaimers
#
# No Warranty: THE SUBJECT SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTY OF ANY KIND,
# EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTY
# THAT THE SUBJECT SOFTWARE WILL CONFORM TO SPECIFICATIONS, ANY IMPLIED WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR FREEDOM FROM INFRINGEMENT, ANY
# WARRANTY THAT THE SUBJECT SOFTWARE WILL BE ERROR FREE, OR ANY WARRANTY THAT DOCUMENTATION,
# IF PROVIDED, WILL CONFORM TO THE SUBJECT SOFTWARE. THIS AGREEMENT DOES NOT, IN ANY MANNER,
# CONSTITUTE AN ENDORSEMENT BY GOVERNMENT AGENCY OR ANY PRIOR RECIPIENT OF ANY RESULTS,
# RESULTING DESIGNS, HARDWARE, SOFTWARE PRODUCTS OR ANY OTHER APPLICATIONS RESULTING FROM
# USE OF THE SUBJECT SOFTWARE.  FURTHER, GOVERNMENT AGENCY DISCLAIMS ALL WARRANTIES AND
# LIABILITIES REGARDING THIRD-PARTY SOFTWARE, IF PRESENT IN THE ORIGINAL SOFTWARE, AND
# DISTRIBUTES IT "AS IS."
#
# Waiver and Indemnity:  RECIPIENT AGREES TO WAIVE ANY AND ALL CLAIMS AGAINST THE UNITED
# STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL AS ANY PRIOR RECIPIENT.
# IF RECIPIENT'S USE OF THE SUBJECT SOFTWARE RESULTS IN ANY LIABILITIES, DEMANDS, DAMAGES,
# EXPENSES OR LOSSES ARISING FROM SUCH USE, INCLUDING ANY DAMAGES FROM PRODUCTS BASED ON,
# OR RESULTING FROM, RECIPIENT'S USE OF THE SUBJECT SOFTWARE, RECIPIENT SHALL INDEMNIFY AND
# HOLD HARMLESS THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL
# AS ANY PRIOR RECIPIENT, TO THE EXTENT PERMITTED BY LAW.  RECIPIENT'S SOLE REMEDY FOR ANY
# SUCH MATTER SHALL BE THE IMMEDIATE, UNILATERAL TERMINATION OF THIS AGREEMENT.
# ******************************************************************************************

"""
Sample environment, returning `EnvironmentValue` (default) or array.
"""
Base.rand(env::Environment) = EnvironmentValue(k => rand(dist) for (k, dist) in env)
Base.rand(fenv::Flat) = flatten(fenv.env, rand(fenv.env))

"""
Infer dimension of action space.
"""
act_dim(mdp::ASTMDP{<:State, SampleAction}) = sum(info.n for info in values(mdp.env_info))

"""
Infer dimension of state space.
"""
obs_dim(mdp::ASTMDP{ObservableState, <:Action}) = length(observe(mdp.sim))

"""
Return ordered list dictionary keys.
"""
orderedkeys(dict::Dict) = sort!(collect(keys(dict)))

"""
Flatten `EnvironmentValue` into array.
"""
function flatten(env::Environment, value::EnvironmentValue)
	return reduce(append!, flatten(env[k], value[k]) for k in orderedkeys(env))
end

"""
Reconstruct `EnvironmentValue` from array.
"""
function unflatten(mdp::ASTMDP{<:State, SampleAction}, action::AbstractVector{<:Real})
	value = EnvironmentValue()
	env = environment(mdp.sim)

	i = 0
	for k in orderedkeys(env)
		n = mdp.env_info[k].n
		array = action[i+1:i+n]
		value[k] = unflatten(env[k], array)
		i += n
	end
	return value
end

"""
Calculate log probability of sample from environment variable.
"""
function logprob(distribution::Any, value::Any, marginalize::Bool)
    logp = logpdf(distribution, value)
    if marginalize
        logp -= logpdf(distribution, mode(distribution))
    end
    return logp
end

"""
Calculate total log probability of environment value.
"""
function logprob(env::Environment, value::EnvironmentValue, marginalize::Bool)
    return sum(logprob(env[k], value[k], marginalize) for k in keys(env))
end

"""
Infer information about simulation environment.
"""
function infer_info(env::Environment)
    env_info = EnvironmentInfo()
	for (k, dist) in env
		sample = rand(dist)
		array = flatten(dist, sample)
		env_info[k] = VariableInfo(length(array), typeof(sample))
	end
    return env_info
end

"""
Infer type of state.
"""
function infer_state(sim::AbstractSimulation)
    try
        isempty(observe(sim)) ? UnobservableState : ObservableState
    catch e
        e isa UnimplementedError ? UnobservableState : throw(e)
    end
end

"""
    ASTMDP(sim::AbstractSimulation; kwargs...)

Constructor for ASTMDP object. Infers various properties of MDP. Allows kwargs to be
specified flatly and automatically assigns to correct level.
"""
function ASTMDP(sim::AbstractSimulation; kwargs...)
    reset!(sim)
    s_type = infer_state(sim)
    a_type = sim isa BlackBox ? SeedAction : SampleAction
    env_info = sim isa BlackBox ? EnvironmentInfo() : infer_info(environment(sim))
    mdp = ASTMDP{s_type, a_type}(; sim=sim, env_info=env_info)

    # automatically applies kwarg to reward if match is found
    for k in keys(kwargs)
        obj = k in fieldnames(ASTMDP) || !(k in fieldnames(Reward)) ? mdp : mdp.reward
        setproperty!(obj, k, kwargs[k])
    end

    mdp.reward.heuristic = mdp.episodic ? FinalHeuristic() : mdp.reward.heuristic
    RNG_TEMP[] = mdp.rng == Random.default_rng() ? Random.Xoshiro() : deepcopy(mdp.rng)
    return mdp
end

"""
Wrap action for internal handling.
"""
convert_a(mdp::ASTMDP{<:State, SampleAction}, action::AbstractVector{<:Real}) = SampleAction(unflatten(mdp, action))
convert_a(::ASTMDP{<:State, SampleAction}, action::EnvironmentValue) = SampleAction(action)
convert_a(::ASTMDP{<:State, SeedAction}, action::UInt32) = SeedAction(action)