module Board_Test

using Intan, SpikeSorting, DistributedArrays

if VERSION > v"0.7-"
    using Test
else
    using Base.Test
end

#64 channel amp check

    myamp=RHD2164("PortA1")
    @test myamp == [0,8]
    myamp=RHD2164("PortA2")
    @test myamp == [1,9]
    myamp=RHD2164("PortB1")
    @test myamp == [2,10]
    myamp=RHD2164("PortB2")
    @test myamp == [3,11]
    myamp=RHD2164("PortC1")
    @test myamp == [4,12]
    myamp=RHD2164("PortC2")
    @test myamp == [5,13]
    myamp=RHD2164("PortD1")
    @test myamp == [6,14]
    myamp=RHD2164("PortD2")
    @test myamp == [7,15]

#32 channel amp check

    myamp=RHD2132("PortA1")
    @test myamp == [0]
    myamp=RHD2132("PortA2")
    @test myamp == [1]
    myamp=RHD2132("PortB1")
    @test myamp == [2]
    myamp=RHD2132("PortB2")
    @test myamp == [3]
    myamp=RHD2132("PortC1")
    @test myamp == [4]
    myamp=RHD2132("PortC2")
    @test myamp == [5]
    myamp=RHD2132("PortD1")
    @test myamp == [6]
    myamp=RHD2132("PortD2")
    @test myamp == [7]

myamp=RHD2164("PortA1")
myfpga=FPGA(1,myamp)

    @test myfpga.numDataStreams == 0
    @test myfpga.dataStreamEnabled == zeros(Int64,1,Intan.MAX_NUM_DATA_STREAMS)
    @test myfpga.amps == [0,8]

#=
Single Core Data Structure
=#
myamp=RHD2164("PortA1")
myt=Task_NoTask()
d=Debug(string(dirname(Base.source_path()),"/data/qq.mat"),"qq")

myfpga=FPGA(1,myamp)

(myrhd,ss,myfpgas)=makeRHD([myfpga],debug=d);

    @test myrhd.v == zeros(Int16, Intan.SAMPLES_PER_DATA_BLOCK,64)
    @test length(ss) == 64
    @test typeof(ss) == Array{SpikeSorting.Sorting_1,1}
    @test typeof(myrhd.buf) == Array{SpikeSorting.Spike,2}
    @test size(myrhd.buf,2) == 64
    @test myrhd.nums == zeros(Int64,64)
    @test myrhd.time == zeros(UInt32,Intan.SAMPLES_PER_DATA_BLOCK,1)

Intan.init_board!(myrhd,myfpgas)

    @test myfpgas[1].numDataStreams == 2
    @test myfpgas[1].dataStreamEnabled == [1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    @test myfpgas[1].sampleRate == 30000
    @test myfpgas[1].numWords == Intan.SAMPLES_PER_DATA_BLOCK * (4+2+(2*36) +8+2)
    @test myfpgas[1].numBytesPerBlock == Intan.SAMPLES_PER_DATA_BLOCK * (4+2+(2*36) +8+2) * 2

#=
Multi Core Data Structure
=#
myamp=RHD2132("PortA1")
d=Debug(string(dirname(Base.source_path()),"/data/qq.mat"),"qq")
myt=Task_NoTask()
myfpga1=FPGA(1,myamp)
myfpga2=FPGA(2,myamp)
(myrhd2,ss2,myfpgas)=makeRHD([myfpga1,myfpga2],debug=d,parallel=true);

    @test myrhd2.v == zeros(Int64, Intan.SAMPLES_PER_DATA_BLOCK,64)
    @test length(ss2) == 64
    @test typeof(ss2) == DistributedArrays.DArray{SpikeSorting.Sorting_1,1,Array{SpikeSorting.Sorting_1,1}}
    @test typeof(myrhd2.buf) == SharedArray{SpikeSorting.Spike,2}
    @test size(myrhd2.buf,2) == 64
    @test myrhd2.nums == zeros(Int64,64)

#Intan.init_board!(myrhd2)

#=

    @test myrhd.fpga[1].numDataStreams --> 1
    @test myrhd.fpga[1].dataStreamEnabled --> [1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    @test myrhd.fpga[1].sampleRate --> 30000
    @test myrhd.fpga[1].numWords --> 31200
    @test myrhd.fpga[1].numBytesPerBlock --> 62400

=#
#=
Sample Rate Testing
=#

myamp=RHD2164("PortA1")
myt=Task_NoTask()
d=Debug(string(dirname(Base.source_path()),"/data/qq.mat"),"qq")

myfpga=FPGA(1,myamp)

(myrhd,ss,myfpgas)=makeRHD([myfpga],debug=d);

Intan.init_board!(myrhd,myfpgas)

    for i in [1000,1250,1500,2000,2500,3000,3333,4000,5000,6250,8000,10000,12500,15000,20000,25000,30000]
        Intan.setSampleRate(myfpgas[1],i,myrhd.debug.state)
        @test myfpgas[1].sampleRate == i
    end

end


"""

$(TYPEDEF)

Structure to contain data needed to compute the mddf for a single frame

$(TYPEDFIELDS)

"""
mutable struct FrameData{T<:Trajectory,V}

    trajectory::T
    volume_frame::Volume
    rdf_count_random_frame::Vector{Float64}
    md_count_random_frame::Vector{Float64}

    dc::CutoffDistances
    dmin_mol::Vector{DminMol}
    dref_mol::Vector{Float64}
    x_solvent_random::Vector{V}

    lc_solvent::LinkedCells

end

function FrameData(trajectory::Trajectory, R::Result)
    return FrameData(
        trajectory,                                       # trajectory
        Volume(R.nbins),                                  # volume_frame
        zeros(R.nbins),                                   # rdf_count_random_frame
        zeros(R.nbins),                                   # md_count_random_frame
        CutoffDistances(trajectory.solvent.natoms),       # dc
        [DminMol(+Inf, i, 0, 0) for i = 1:trajectory.solvent.nmols], # dmin_mol
        zeros(trajectory.solvent.nmols),                             # dref_mol
        similar(trajectory.x_solvent),                               # x_solvent_random
        LinkedCells(trajectory.solvent.natoms),
    )                      # lc_solvent
end


module EBIC

export run_ebic

using JSON
using DataStructures: SortedSet
using ProgressMeter: next!, finish!, Progress
using ArgParse: ArgParseSettings, @add_arg_table, parse_args

include("constants.jl")
include("biclusterseval.jl")
include("evolution.jl")
include("scoring.jl")
include("algorithm.jl")

using .evolution: init_population, mutate
using .scoring: score_population
using .algorithm: update_rank_list!, ReverseOrdering
using .biclusterseval: get_biclusters, initialize_input_on_gpus

function run_ebic(;
    input_path::String,
    best_bclrs_stats = true,
    verbose = true,
    max_iterations = MAX_ITERATIONS,
    max_biclusters = MAX_BICLUSTERS_NUMBER,
    overlap_threshold = OVERLAP_THRESHOLD,
    negative_trends = NEGATIVE_TRENDS_ENABLED,
    approx_trends_ratio::Float32 = APPROX_TRENDS_RATIO,
    gpus_num = GPUS_NUMBER,
    output = false,
)
    data_load_time = @elapsed begin
        d_input_data = initialize_input_on_gpus(input_path, gpus_num)
    end
    @debug "Loading data to GPU took: $(data_load_time)s"

    # used to evaluate the iteration and timing of the best bclrs finding
    prev_top_bclrs = Vector()
    last_top_bclrs_change = (0, 0)

    if verbose
        p_bar = Progress(max_iterations, barlen = 20)
    end

    @debug "Starting algorithm..."
    # algorithm initialization steps
    start_time = time_ns()
    tabu_list = Set()
    tabu_hits = 0
    top_rank_list = SortedSet(Vector(), ReverseOrdering())

    cols_number = size(d_input_data[1], 2)

    old_population = init_population(cols_number, tabu_list)

    old_scored_population =
        score_population(d_input_data, old_population, gpus_num = gpus_num)

    update_rank_list!(top_rank_list, old_scored_population, overlap_threshold)

    i = 0
    while i < max_iterations
        i += 1
        new_population = Population()

        # reset penalties
        penalties = fill(1, cols_number)

        # elitism
        for (_, chromo) in top_rank_list
            !(length(new_population) < REPRODUCTION_SIZE) && break

            push!(new_population, chromo)
        end

        # perform mutations to replenish new population
        new_population, tabu_hits =
            mutate(new_population, old_scored_population, tabu_list, penalties, cols_number)

        # check if algorithm has found new solutions
        if tabu_hits >= MAX_NUMBER_OF_TABU_HITS
            verbose && finish!(p_bar)
            break
        end

        # evaluate fitness for new population
        new_scored_population =
            score_population(d_input_data, new_population, gpus_num = gpus_num)

        # save best chromosomes
        update_rank_list!(top_rank_list, old_scored_population, overlap_threshold)

        # proceed to the next generation
        old_scored_population = new_scored_population

        if best_bclrs_stats
            new_top_bclrs = collect(top_rank_list)[1:max_biclusters]
            if !isempty(prev_top_bclrs)
                changed = new_top_bclrs != prev_top_bclrs
                if changed
                    last_top_bclrs_change = (i, time_ns() - start_time)
                    prev_top_bclrs = new_top_bclrs
                end
            else
                prev_top_bclrs = new_top_bclrs
            end
        end

        verbose && next!(p_bar)
    end

    algorithm_time = time_ns() - start_time

    biclusters = get_biclusters(
        d_input_data,
        [last(p) for p in top_rank_list],
        gpus_num,
        negative_trends,
        approx_trends_ratio,
    )

    run_summary = Dict(
        "algorithm_time" => algorithm_time / 1e9,
        "data_load_time" => data_load_time,
        "biclusters" => biclusters[1:max_biclusters],
        "performed_iters" => i,
        "last_iter_tabu_hits" => tabu_hits,
    )

    if best_bclrs_stats
        run_summary["best_bclrs_iter"] = last_top_bclrs_change[1]
        run_summary["best_bclrs_time"] = last_top_bclrs_change[2] / 1e9
    end

    if output
        open("$(basename(input_path))-res.json", "w") do f
            JSON.print(f, run_summary["biclusters"])
        end
    end

    return run_summary
end

function main()
    args = ArgParseSettings("""
    EBIC is a next-generation biclustering algorithm based on artificial intelligence (AI). EBIC is probably the first algorithm capable of discovering the most challenging patterns (i.e. row-constant, column-constant, shift, scale, shift-scale and trend-preserving) in complex and noisy data with average accuracy of over 90%. It is also one of the very few parallel biclustering algorithms that use at least one graphics processing unit (GPU) and is ready for big-data challenges.
    """)

    @add_arg_table args begin
        "--input", "-i"
        help = "The path to the input file."
        arg_type = String
        default = INPUT_PATH
        dest_name = "input_path"

        "--max_iterations", "-n"
        help = "The maximum number of iterations of the algorithm."
        arg_type = Int
        default = MAX_ITERATIONS

        "--biclusters_num", "-b"
        help = "The number of biclusters that will be returned in the end."
        arg_type = Int
        default = MAX_BICLUSTERS_NUMBER
        dest_name = "max_biclusters"

        "--overlap_threshold", "-x"
        help = "The maximum similarity level of each two chromosomes held in top rank list."
        arg_type = Float64
        default = OVERLAP_THRESHOLD

        "--negative_trends", "-t"
        help = "Enable negative trends."
        action = :store_true

        "--gpus_num", "-g"
        help = "The number of gpus the algorithm should run on."
        arg_type = Int
        default = GPUS_NUMBER

        "--approx_trends", "-a"
        help = ""
        arg_type = Float32
        default = APPROX_TRENDS_RATIO
        dest_name = "approx_trends_ratio"

        "--verbose", "-v"
        help = "Turn on the progress bar."
        action = :store_true

        "--best_bclrs_stats", "-s"
        help = "Evaluate resulting biclusters finding iteration and time. Enabled, it slightly worsens overall algorithm performance."
        action = :store_true

        "--output", "-o"
        help = "Save biclusters to a file in the JSON format. The output file name is a concatenation of the input file name and '-res.json' suffix."
        action = :store_true

    end
    args = parse_args(args, as_symbols = true)

    results = run_ebic(; args...)

    JSON.print(results)
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end

end


export saveMagneticFieldAsHDF5, MagneticFieldSweepCurrentsMeas

#TODO: Unit handling should be put into GaussMeter

# define measObj
@compat struct MagneticFieldSweepCurrentsMeas <: MeasObj
  su::SurveillanceUnit
  daq::AbstractDAQ
  gauss::GaussMeter
  positions::Positions
  currents::Matrix{Float64}
  gaussRanges::Vector{Int}
  waitTime::Float64
#  voltToCurrent::Float64
  pos::Matrix{typeof(1.0Unitful.m)}
  magneticField::Array{typeof(1.0Unitful.T),3}
  magneticFieldError::Array{typeof(1.0Unitful.T),3}

  MagneticFieldSweepCurrentsMeas(su, daq, gauss, positions, currents, gausMeterRanges, waitTime) =
                 new(su, daq, gauss, positions, currents, gausMeterRanges, waitTime, #voltToCurrent,
                      zeros(typeof(1.0Unitful.m),3,length(positions)),
                      zeros(typeof(1.0Unitful.T),3,length(positions),size(currents,2)),
                      zeros(typeof(1.0Unitful.T),3,length(positions),size(currents,2)))
end

# define preMoveAction
function preMoveAction(measObj::MagneticFieldSweepCurrentsMeas,
                       pos::Vector{typeof(1.0Unitful.mm)}, index)
  @info "moving to position" pos
end

# define postMoveAction
function postMoveAction(measObj::MagneticFieldSweepCurrentsMeas,
                        pos::Vector{typeof(1.0Unitful.mm)}, index)
  @info "post action" pos index length(measObj.positions)
  #sleep(0.05)
  measObj.pos[:,index] = pos

  for l=1:size(measObj.currents,2)
    # set current at DC sources
    #value(measObj.rp,"AOUT0",measObj.currents[1,l]*measObj.voltToCurrent)
    #value(measObj.rp,"AOUT1",measObj.currents[2,l]*measObj.voltToCurrent)
    setSlowDAC(measObj.daq, measObj.currents[1,l], 0)
    setSlowDAC(measObj.daq, measObj.currents[2,l], 1)

    @debug "Set DC source $(measObj.currents[1,l]*Unitful.A)  $(measObj.currents[2,l]*Unitful.A)"
    # set measurement range of gauss meter
    range = measObj.gaussRanges[l]
    setAllRange(measObj.gauss, Char("$range"[1]))
    # wait until magnet is on field
    sleep(0.6)
    # perform field measurment
    magneticField = getXYZValues(measObj.gauss)
    measObj.magneticField[:,index,l] = magneticField
    # perform error estimation based on gauss meter specification
    magneticFieldError = zeros(typeof(1.0Unitful.T),3,2)
    magneticFieldError[:,1] = abs.(magneticField)*1e-3
    magneticFieldError[:,2] = getFieldError(range).*ones(3)
    measObj.magneticFieldError[:,index,l] = maximum(magneticFieldError,dims=2)

    @info "Field $(uconvert.(Unitful.mT,measObj.magneticField[:,index,l]))"
  end
  setSlowDAC(measObj.daq, 0.0, 0)
  setSlowDAC(measObj.daq, 0.0, 1)

  @showprogress for i=1:measObj.waitTime
    sleep(1)
  end

end

function saveMagneticFieldAsHDF5(measObj::MagneticFieldSweepCurrentsMeas,
       filename::String, params=Dict{String,Any}())
  h5open(filename, "w") do file
    write(file, measObj.positions)
    write(file, "/unitCoords", "m")
    write(file, "/unitFields", "T")
    write(file, "/positions", ustrip.(measObj.pos))
    write(file, "/fields", ustrip.(measObj.magneticField))
    write(file, "/fieldsError", ustrip.(measObj.magneticFieldError))
    write(file, "/currents", measObj.currents)
    for (key,value) in params
      write(file, key, value)
    end
  end
end


MaybeFunction = Nullable{Function}
const nullFunction = MaybeFunction()
typealias AkoFunction Union{Function, MaybeFunction}

Base.isnull(x::MaybeFunction) = x.isnull
Base.isnull(x::Function)      = false

abstract  Multivalued

abstract  Surface
abstract  SurfaceAction        <: Surface
abstract  Boundary
abstract  BoundaryMaterial     <: Boundary

abstract  MateralBounding      <: Boundary
abstract  Wood                 <: MaterialBounding
abstract  Clay                 <: MaterialBounding

abstract  BoundingContext      <: Boundary
abstract  Imprecise            <: BoundingContext
abstract  Precise              <: Imprecise
abstract  Inaccurate           <: BoundingContext
abstract  Accurate             <: Inaccurate

abstract  BoundingCoverness    <: Boundary
abstract  Spread               <: BoundingCoverness
abstract  Shared               <: Spread
abstract  Throughout           <: BoundingCoverness
abstract  Approached           <: Throughout

abstract  AnyBoundaries        <: Boundary

abstract  OneBoundary          <: AnyBoundaries
abstract  TwoBoundaries        <: AnyBoundaries
abstract  ThreeBoundaries      <: AnyBoundaries

abstract
AbsInterval{          Surface,       Tension                  } <: Number

                #   type covering   configures     perceptual
                #   value-context   boundaries     values map

abstract
AbsFlexibleInterval{  Surface,       Tension,        Action   } <: AbsInterval{ Surface, Tension }

immutable
FlexibleInterval{     Surface <: Real,
                      Tension <: AnyBoundaries,
                      Action  <: FunctionMaybe  } <: AbsFlexibleInterval{ Surface, Tension, Action }
    lo::Surface
    hi::Surface

    FlexibleInterval{S}(lo::S,hi::S) = (lo <= hi) ? new(lo,hi) : new(hi,lo)
end

FlexibleInterval{S<:Real,T<:AnyBoundaries,A<:FunctionMaybe}(boundaries::T, action::A, lo::S,hi::S) =
    FlexibleInterval{S,T,A}(lo,hi)

FlexibleInterval{S,T<:AnyBoundaries,A<:FunctionMaybe}(::Type{T}, ::Type{A}, lo::S, hi::S) = FlexibleInterval(lo,hi)
FlexibleInterval{S,T<:AnyBoundaries}(::Type{T}, lo::S, hi::S) = FlexibleInterval(T, nullFunction, lo, hi)
FlexibleInterval{S,A<:FunctionMaybe}(::Type{A}, lo::S, hi::S) = FlexibleInterval(ClosedCloseD, A, lo, hi)
FlexibleInterval{S}(lo::S, hi::S)   = FlexibleInterval(ClosedCloseD, nullFunction, lo, hi)


macro logprob_str(str)
    expr1, expr2 = get_exprs(str)
    return :(logprob($(esc(expr1)), $(esc(expr2))))
end
macro prob_str(str)
    expr1, expr2 = get_exprs(str)
    return :(exp.(logprob($(esc(expr1)), $(esc(expr2)))))
end

function get_exprs(str::String)
    substrings = split(str, '|'; limit = 2)
    length(substrings) == 2 || error("Invalid expression.")
    str1, str2 = substrings

    expr1 = Meta.parse("($str1,)")
    expr1 = Expr(:tuple, expr1.args...)

    expr2 = Meta.parse("($str2,)")
    expr2 = Expr(:tuple, expr2.args...)

    return expr1, expr2
end

function logprob(ex1, ex2)
    ptype, model, vi = probtype(ex1, ex2)
    if ptype isa Val{:prior}
        return logprior(ex1, ex2, model, vi)
    elseif ptype isa Val{:likelihood}
        return loglikelihood(ex1, ex2, model, vi)
    end
end

function probtype(ntl::NamedTuple{namesl}, ntr::NamedTuple{namesr}) where {namesl, namesr}
    if :chain in namesr
        if isdefined(ntr.chain.info, :model)
            model = ntr.chain.info.model
        elseif isdefined(ntr, :model)
            model = ntr.model
        else
            throw("The model is not defined. Please make sure the model is either saved in the chain or passed on the RHS of |.")
        end
        @assert model isa Model
        if isdefined(ntr.chain.info, :vi)
            _vi = ntr.chain.info.vi
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        elseif isdefined(ntr, :varinfo)
            _vi = ntr.varinfo
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        else
            vi = nothing
        end
        defaults = model.defaults
        @assert all(getargnames(model)) do arg
            isdefined(ntl, arg) || isdefined(ntr, arg) ||
                isdefined(defaults, arg) && getfield(defaults, arg) !== missing
        end
        return Val(:likelihood), model, vi
    else
        @assert isdefined(ntr, :model)
        model = ntr.model
        @assert model isa Model
        if isdefined(ntr, :varinfo)
            _vi = ntr.varinfo
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        else
            vi = nothing
        end
        return probtype(ntl, ntr, model), model, vi
    end
end

function probtype(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames}
) where {leftnames,rightnames,argnames,defaultnames,_F}
    defaults = model.defaults
    prior_rhs = all(n -> n in (:model, :varinfo) ||
        n in argnames && getfield(right, n) !== missing, rightnames)
    function get_arg(arg)
        if arg in leftnames
            return getfield(left, arg)
        elseif arg in rightnames
            return getfield(right, arg)
        elseif arg in defaultnames
            return getfield(defaults, arg)
        else
            return nothing
        end
    end
    function valid_arg(arg)
        a = get_arg(arg)
        return a !== nothing && a !== missing
    end
    valid_args = all(valid_arg, argnames)

    # Uses the default values for model arguments not provided.
    # If no default value exists, use `nothing`.
    if prior_rhs
        return Val(:prior)
    # Uses the default values for model arguments not provided.
    # If no default value exists or the default value is missing, then error.
    elseif valid_args
        return Val(:likelihood)
    else
        for argname in argnames
            if !valid_arg(argname)
                throw(ArgumentError(missing_arg_error_msg(argname, get_arg(argname))))
            end
        end
    end
end

missing_arg_error_msg(arg, ::Missing) = """Variable $arg has a value of `missing`, or is not defined and its default value is `missing`. Please make sure all the variables are either defined with a value other than `missing` or have a default value other than `missing`."""
missing_arg_error_msg(arg, ::Nothing) = """Variable $arg is not defined and has no default value. Please make sure all the variables are either defined with a value other than `missing` or have a default value other than `missing`."""

function logprior(
    left::NamedTuple,
    right::NamedTuple,
    _model::Model,
    _vi::Union{Nothing, VarInfo}
)
    # For model args on the LHS of |, use their passed value but add the symbol to
    # model.missings. This will lead to an `assume`/`dot_assume` call for those variables.
    # Let `p::PriorContext`. If `p.vars` is `nothing`, `assume` and `dot_assume` will use
    # the values of the random variables in the `VarInfo`. If `p.vars` is a `NamedTuple`
    # or a `Chain`, the value in `p.vars` is input into the `VarInfo` and used instead.

    # For model args not on the LHS of |, if they have a default value, use that,
    # otherwise use `nothing`. This will lead to an `observe`/`dot_observe`call for
    # those variables.
    # All `observe` and `dot_observe` calls are no-op in the PriorContext

    # When all of model args are on the lhs of |, this is also equal to the logjoint.
    model = make_prior_model(left, right, _model)
    vi = _vi === nothing ? VarInfo(deepcopy(model), PriorContext()) : _vi
    foreach(keys(vi.metadata)) do n
        @assert n in keys(left) "Variable $n is not defined."
    end
    model(vi, SampleFromPrior(), PriorContext(left))
    return getlogp(vi)
end

@generated function make_prior_model(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames}
) where {leftnames,rightnames,argnames,defaultnames,_F}
    argvals = []
    missings = []
    warnings = []

    for argname in argnames
        if argname in leftnames
            push!(argvals, :(deepcopy(left.$argname)))
            push!(missings, argname)
        elseif argname in rightnames
            push!(argvals, :(right.$argname))
        elseif argname in defaultnames
            push!(argvals, :(model.defaults.$argname))
        else
            push!(warnings, :(@warn($(warn_msg(argname)))))
            push!(argvals, :(nothing))
        end
    end

    # `args` is inserted as properly typed NamedTuple expression;
    # `missings` is splatted into a tuple at compile time and inserted as literal
    return quote
        $(warnings...)
        Model{$(Tuple(missings))}(
            model.name, model.f, $(to_namedtuple_expr(argnames, argvals)), model.defaults,
        )
    end
end

warn_msg(arg) = "Argument $arg is not defined. A value of `nothing` is used."

function Distributions.loglikelihood(
    left::NamedTuple,
    right::NamedTuple,
    _model::Model,
    _vi::Union{Nothing, VarInfo},
)
    model = make_likelihood_model(left, right, _model)
    vi = _vi === nothing ? VarInfo(deepcopy(model)) : _vi
    if isdefined(right, :chain)
        # Element-wise likelihood for each value in chain
        chain = right.chain
        ctx = LikelihoodContext(right)
        iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
        logps = map(iters) do (sample_idx, chain_idx)
            setval!(vi, chain, sample_idx, chain_idx)
            model(vi, SampleFromPrior(), ctx)
            return getlogp(vi)
        end
        return reshape(logps, size(chain, 1), size(chain, 3))
    else
        # Likelihood without chain
        # Rhs values are used in the context
        ctx = LikelihoodContext(right)
        model(vi, SampleFromPrior(), ctx)
        return getlogp(vi)
    end
end

@generated function make_likelihood_model(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames},
) where {leftnames,rightnames,argnames,defaultnames,_F}
    argvals = []
    missings = []

    for argname in argnames
        if argname in leftnames
            push!(argvals, :(left.$argname))
        elseif argname in rightnames
            push!(argvals, :(right.$argname))
            push!(missings, argname)
        elseif argname in defaultnames
            push!(argvals, :(model.defaults.$argname))
        else
            throw("This point should not be reached. Please open an issue in the DynamicPPL.jl repository.")
        end
    end

    # `args` is inserted as properly typed NamedTuple expression;
    # `missings` is splatted into a tuple at compile time and inserted as literal
    return :(Model{$(Tuple(missings))}(
        model.name, model.f, $(to_namedtuple_expr(argnames, argvals)), model.defaults,
    ))
end


using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    include("../src/components/MarketDamages.jl")

    marketdamages = addmarketdamages(m)

    set_param!(m, :MarketDamages, :rtl_realizedtemperature, readpagedata(m, "test/validationdata/$valdir/rtl_realizedtemperature.csv"))
    set_param!(m, :MarketDamages, :rcons_per_cap_SLRRemainConsumption, readpagedata(m,"test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv"))
    set_param!(m, :MarketDamages, :rgdp_per_cap_SLRRemainGDP, readpagedata(m,"test/validationdata/$valdir/rgdp_per_cap_SLRRemainGDP.csv"))
    set_param!(m, :MarketDamages, :atl_adjustedtolerableleveloftemprise, readpagedata(m,"test/validationdata/atl_adjustedtolerableleveloftemprise_market.csv"))
    set_param!(m, :MarketDamages, :imp_actualreduction, readpagedata(m,"test/validationdata/$valdir/imp_actualreduction_market.csv"))
    set_param!(m, :MarketDamages, :isatg_impactfxnsaturation, 28.333333333333336)

    p = load_parameters(m)
    p["y_year_0"] = 2015.
    p["y_year"] = Mimi.dim_keys(m.md, :time)
    set_leftover_params!(m, p)

    run(m)

    ## Not a very good test, but damages are based on Burke et al. now anyway
    iref = m[:MarketDamages, :iref_ImpactatReferenceGDPperCap]
    iref_compare = readpagedata(m, "test/validationdata/$valdir/iref_ImpactatReferenceGDPperCap_econ_page.csv")
    @test iref ≈ iref_compare rtol=1e-1

    # rcons_per_cap = m[:MarketDamages, :rcons_per_cap_MarketRemainConsumption]
    # rcons_per_cap_compare = readpagedata(m, "test/validationdata/rcons_per_cap_MarketRemainConsumption.csv")
    # @test rcons_per_cap ≈ rcons_per_cap_compare rtol=1e-1

    # rgdp_per_cap = m[:MarketDamages, :rgdp_per_cap_MarketRemainGDP]
    # rgdp_per_cap_compare = readpagedata(m, "test/validationdata/rgdp_per_cap_MarketRemainGDP.csv")
    # @test rgdp_per_cap ≈ rgdp_per_cap_compare rtol=1e-2
end


module Game

import Util
import Moves
import Search

function pickMoveUCI(board,colour,history)
    Search.search(3,colour,history,board)
    return Search.best_move
end

function play(board,colour,history)
    states = []
    while(true)
        opp_colour = (colour === 'W') ? 'B' : 'W'
        board,states = @time takeTurn(board,history,colour,states)
        colour = opp_colour
        if(checkEndGame(board,states,history,colour)) break end
    end

end

function takeTurn(board,history,colour,states)
    move = pickMoveUCI(board,colour,history)
    push!(history,move)
    println("move picked: ", move)
    board = Moves.makeMoveUCI(move,board,colour)
    state = reduce((x,y) -> x | y, values(board))
    push!(states,state)
    Util.prettyPrintBoard(board)
    return board,states
end

# //=============================  END GAME LOGIC  =============================//
# //============================================================================//

function checkEndGame(board,states,history,colour)
    opp_colour = (colour ==='W') ? 'B' : 'W'
    if (checkForMate(board,colour,history) === 1 || checkForMate(board,colour,history) === 2 ||
        sameBoardStateFiveTimes(states) || fiftyMoves(history) || ~(hasEnoughMaterial(board,colour) || hasEnoughMaterial(board,opp_colour)))
        println("GAME OVER")
        print(board)
        return true
    end
    return false
end

function checkForMate(board,colour,history)
    check_mate = Moves.isCheckMate(board,colour,history)
    if(check_mate == true)
        println("checkmate")
    end
    return check_mate
end

function sameBoardStateFiveTimes(states)
    groups = Dict()
    for i in states
        if i in keys(groups)
            groups[i] = groups[i] + 1
            if(groups[i] > 5)
                println("same state five times")
                return true
            end
        else
            groups[i] = 1
        end
    end
    return false
end

# The fifty-move rule in chess states that a player can claim a draw if no capture has been made and no pawn has been moved in the last fifty moves
# (for this purpose a "move" consists of a player completing their turn followed by the opponent completing their turn)
function fiftyMoves(history)
    if(length(history) < 50)
        return false
    end
    fifty = history[length(history)-49:length(history)]
    pawn_moves = [] #how to get pawn moves from UCI notated?
    if(length(pawn_moves) > 0)
        return false
    end
    captures = ['l'] #how to get caps
    if(length(captures) > 0)
        return false
    end
    print("fifty moves")
    return true
end

function hasEnoughMaterial(board,colour)
    if(board[colour * 'P'] > 0 || (board[colour * 'N'] > 0 && board[colour * 'B'] > 0) || board[colour * 'R'] > 0
        || board[colour * 'Q'] > 0  || count_ones(board[colour * 'B']) === 2)
        return true
    end
    print("not enough material")
    return false
end

end


export monitor_functions, monitor_function, get_active, set_active!

"""
    $(TYPEDEF)

An individual monitor function.
"""
struct MonitorFunction{F}
    f::F
    initial_active::Base.RefValue{Bool}
end

# Slightly annoying special casing of the different combinations of inputs to account for
# the fact that Tuples are unwrapped if they are singletons

function (mfunc::MonitorFunction)(res, u::Tuple, data::Tuple, prob...)
    res[1] =
        mfunc.f(Base.tail(u), Base.tail(data), prob...) -
        (isempty(u[1]) ? data[1][] : u[1][1])
    return nothing
end

function (mfunc::MonitorFunction)(res, u::Tuple{<:Any,<:Any}, data::Tuple, prob...)
    res[1] = mfunc.f(u[2], Base.tail(data), prob...) - (isempty(u[1]) ? data[1][] : u[1][1])
    return nothing
end

function (mfunc::MonitorFunction)(res, u::Tuple, data::Tuple{<:Any,<:Any}, prob...)
    res[1] = mfunc.f(Base.tail(u), data[2], prob...) - (isempty(u[1]) ? data[1][] : u[1][1])
    return nothing
end

function (mfunc::MonitorFunction)(
    res, u::Tuple{<:Any,<:Any}, data::Tuple{<:Any,<:Any}, prob...
)
    res[1] = mfunc.f(u[2], data[2], prob...) - (isempty(u[1]) ? data[1][] : u[1][1])
    return nothing
end

function (mfunc::MonitorFunction)(res, u::Tuple, data::Tuple{<:Any}, prob...)
    res[1] = mfunc.f(Base.tail(u), prob...) - (isempty(u[1]) ? data[] : u[1][1])
    return nothing
end

function (mfunc::MonitorFunction)(res, u::Tuple{<:Any,<:Any}, data::Tuple{<:Any}, prob...)
    res[1] = mfunc.f(u[2], prob...) - (isempty(u[1]) ? data[] : u[1][1])
    return nothing
end

function monitor_function(
    name,
    f,
    var=(),
    data=();
    initial_value=nothing,
    active=false,
    group=:embedded,
    pass_problem=false,
    top_level=true,
)
    mvar = Var(name; initial_dim=1, initial_u=initial_value, top_level=top_level)
    mdata = Data("mfunc_data", Ref(initial_value))
    fullgroup = (group isa Symbol ? push! : append!)([:mfunc], group)
    return Func(
        name,
        MonitorFunction(f, Ref(active)),
        (mvar, var...),
        (mdata, data...);
        initial_dim=1,
        group=fullgroup,
        pass_problem=pass_problem,
    )
end

function parameter(name, var; active=false, top_level=true, index=1)
    local mfunc
    let index = index
        mfunc = monitor_function(
            name, u -> u[index], (var,); active=active, top_level=top_level
        )
    end
end

function parameters(names, var; kwargs...)
    return [parameter(name, var; index=i, kwargs...) for (i, name) in enumerate(names)]
end

"""
    $(SIGNATURES)

Set a monitor function active or inactive. Only used during the creation of the problem
structure; once the continuation problem is closed, changing this has no effect.
"""
function set_active!(mfunc::Func, active::Bool)
    if mfunc.func isa MonitorFunction
        mfunc.func.initial_active[] = active
    else
        throw(ArgumentError("Func provided is not a MonitorFunction"))
    end
end

"""
    $(SIGNATURES)

Return whether a monitor function is active (able to change value) or not.
"""
function get_active(mfunc::Func)
    if mfunc.func isa MonitorFunction
        return mfunc.func.initial_active[]
    else
        throw(ArgumentError("Func provided is not a MonitorFunction"))
    end
end

"""
    $(TYPEDEF)

A problem structure to contain monitor functions.
"""
struct MonitorFunctions <: ProblemOwner
    mfunc::Vector{MonitorFunction}  # MonitorFunction is an abstract type but that's fine since it's not used in a hot loop
    idx_var::Vector{Int64}
    idx_data::Vector{Int64}
    idx_func::Vector{Int64}
    name::Dict{String,Int64}
end
function MonitorFunctions()
    return MonitorFunctions(
        MonitorFunction[], Int64[], Int64[], Int64[], Dict{String,Int64}()
    )
end

function (mfuncs::MonitorFunctions)(::Signal{:post_correct}, problem)
    # Store the var value in data (if var is non-empty)
    chart = get_current_chart(problem)
    u = get_uview(chart)
    data = get_data(chart)
    for (idx_var, idx_data) in zip(mfuncs.idx_var, mfuncs.idx_data)
        if !isempty(u[idx_var])
            data[idx_data][] = u[idx_var]
        end
    end
    return nothing
end

function (mfuncs::MonitorFunctions)(::Signal{:initial_state}, problem, u, t, data)
    T = eltype(eltype(u))
    for (mfunc, idx_var) in zip(mfuncs.mfunc, mfuncs.idx_var)
        if !mfunc.initial_active[]
            u[idx_var] = T[]
            t[idx_var] = T[]
        end
    end
    return nothing
end

function init!(mfuncs::MonitorFunctions, problem)
    flat = get_flatproblem(problem)
    # Find all monitor functions and store their locations
    for (i, func) in pairs(flat.func)
        if func.func isa MonitorFunction
            idx_var = findfirst(==(func.var[1]), flat.var)
            idx_data = findfirst(==(func.data[1]), flat.data)
            push!(mfuncs.mfunc, func.func)
            push!(mfuncs.idx_var, idx_var)
            push!(mfuncs.idx_data, idx_data)
            push!(mfuncs.idx_func, i)
            mfuncs.name[first(k for (k, v) in flat.var_names if v == idx_var)] = lastindex(mfuncs.mfunc)
        end
    end
end

function exchange_pars(flat::FlatProblem) end

"""
    $(SIGNATURES)

Create a problem structure to contain monitor functions
"""
function monitor_functions(name="mfuncs")
    problem = Problem(name, MonitorFunctions())
    return problem
end


"""
    ShapeColl{S<:AbsShapeDesc}

This is the type of a homogeneous collection of finite element shapes.

- `S` = shape descriptor: subtype of AbsShapeDesc.
"""
struct ShapeColl{S<:AbsShapeDesc}
    shapedesc::S # Shape descriptor
    nshapes::Int64 # Number of shapes in the collection
    attributes::Dict{String,AbsAttrib} # dictionary of attributes
    name::String # name of the shape collection
end

"""
    ShapeColl(shapedesc::S, nshapes::Int64) where {S<:AbsShapeDesc}

Set up new shape collection with an empty dictionary of attributes and a default name.
"""
function ShapeColl(shapedesc::S, nshapes::Int64) where {S<:AbsShapeDesc}
    ShapeColl(shapedesc::S, nshapes, Dict{String,AbsAttrib}(), "shapes")
end

"""
    ShapeColl(shapedesc::S, nshapes::Int64, s::String) where {S<:AbsShapeDesc}

Set up new shape collection with an empty dictionary of attributes.
"""
function ShapeColl(shapedesc::S, nshapes::Int64, s::String) where {S<:AbsShapeDesc}
    ShapeColl(shapedesc::S, nshapes, Dict{String,AbsAttrib}(), s)
end

"""
    shapedesc(shapes::ShapeColl)

Retrieve the shape descriptor.
"""
shapedesc(shapes::ShapeColl) = shapes.shapedesc

"""
    nshapes(shapes::ShapeColl)

Number of shapes in the collection.
"""
nshapes(shapes::ShapeColl) = shapes.nshapes

"""
    manifdim(shapes::ShapeColl)

Retrieve the manifold dimension of the collection.
"""
manifdim(shapes::ShapeColl) = manifdim(shapes.shapedesc)

"""
    nvertices(shapes::ShapeColl)

Retrieve the number of vertices per shape.
"""
nvertices(shapes::ShapeColl) = nvertices(shapes.shapedesc)

"""
    facetdesc(shapes::ShapeColl)

Retrieve the shape type of the boundary facet.
"""
facetdesc(shapes::ShapeColl) = shapes.shapedesc.facetdesc

"""
    nfacets(shapes::ShapeColl)

Retrieve the number of boundary facets per shape.
"""
nfacets(shapes::ShapeColl) = nfacets(shapes.shapedesc)

"""
    facetconnectivity(shapes::ShapeColl, j::Int64)

Retrieve connectivity of the `j`-th facet shape of the `i`-th shape from the collection.
"""
function facetconnectivity(shapes::ShapeColl, j::Int64)
    return shapes.shapedesc.facets[j, :]
end

"""
    ridgedesc(shapes::ShapeColl)

Retrieve the shape type of the boundary ridge.
"""
ridgedesc(shapes::ShapeColl) = shapes.shapedesc.ridgedesc

"""
    nridges(shapes::ShapeColl)

Retrieve the number of boundary ridges per shape.
"""
nridges(shapes::ShapeColl) = nridges(shapes.shapedesc)

"""
    ridgeconnectivity(shapes::ShapeColl, j::Int64)

Retrieve connectivity of the `j`-th ridge shape of the `i`-th shape from the collection.
"""
function ridgeconnectivity(shapes::ShapeColl, j::Int64)
    return shapes.shapedesc.ridges[j, :]
end

"""
    attribute(shapes::ShapeColl, s::Symbol)

Retrieve a named attribute.
"""
attribute(shapes::ShapeColl, s::String) = shapes.attributes[s]

function _sense(fc, oc, nshifts) # is the facet used in the positive or in the negative sense?
    if fc == oc
        return +1 # facet used in the positive sense
    end
    for i in 1:(nshifts - 1)
        fc = circshift(fc, 1) # try a circular shift
        if fc == oc
            return +1 # facet used in the positive sense
        end
    end
    return -1 # facet used in the positive sense
end

"""
    Base.summary(sc::S) where {S<:ShapeColl}

Form a brief summary of the shape collection.
"""
function Base.summary(sc::S) where {S<:ShapeColl}
    s = "$(sc.name) = $(nshapes(sc)) x $(shapedesc(sc).name)"
    if !isempty(values(sc.attributes))
        s = s * " {"
        for k in keys(sc.attributes)
            s = s * k * ","
        end
        s = s * "}"
    end

    return s
end


# This file was generated, do not modify it. # hide
model1 = linear
model2 = knn


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using DifferentialEquations, DiffEqCallbacks, Zygote

"""Handle events and return the values and nominals of the changed continuous states"""
function handleEvents(c::fmi2Component, initialEventMode)

    if initialEventMode == false
        fmi2EnterEventMode(c)
    end

    eventInfo = fmi2NewDiscreteStates(c)

    eventInfo.newDiscreteStatesNeeded = fmi2True
    valuesOfContinuousStatesChanged = fmi2False
    nominalsOfContinuousStatesChanged = fmi2False

    #set inputs here
    #fmiSetReal(myFMU, InputRef, Value)

    while eventInfo.newDiscreteStatesNeeded == fmi2True
        # update discrete states
        eventInfo = fmi2NewDiscreteStates(c)
        valuesOfContinuousStatesChanged = eventInfo.valuesOfContinuousStatesChanged
        nominalsOfContinuousStatesChanged = eventInfo.nominalsOfContinuousStatesChanged

        if eventInfo.terminateSimulation == fmi2True
            error("Event info returned an error")
        end
    end

    fmi2EnterContinuousTimeMode(c)

    return valuesOfContinuousStatesChanged, nominalsOfContinuousStatesChanged

end

"""Returns the event indicators for an FMU"""
function condition(c::fmi2Component, out, x, t, integrator) # Event when event_f(u,t) == 0

    fmi2SetTime(c, t)
    fmi2SetContinuousStates(c, x)
    indicators = fmi2GetEventIndicators(c)

    copy!(out, indicators)

end

"""Handles the upcoming events"""
function affect!(c::fmi2Component, integrator, idx)

    # Event found - handle it
    continuousStatesChanged, nominalsChanged = handleEvents(c, false)

    if continuousStatesChanged == fmi2True
        integrator.u = fmi2GetContinuousStates(c)
    end

    if nominalsChanged == fmi2True
        x_nom = fmi2GetNominalsOfContinuousStates(c)
    end
end

"""Does one step in the simulation"""
function stepCompleted(c::fmi2Component, x, t, integrator)

     (status, enterEventMode, terminateSimulation) = fmi2CompletedIntegratorStep(c, fmi2True)
     if enterEventMode == fmi2True
        affect!(c, integrator, 0)
     end

end

"""Returns the derivatives of the FMU"""
function fx(c::fmi2Component, x, p, t)

    fmi2SetTime(c, t)
    fmi2SetContinuousStates(c, x)
    dx = fmi2GetDerivatives(c)

end

""" Source: FMISpec2.0.2[p.90 ff]: 3.2.4 Pseudocode Example
    simulates an fmu instance for the set simulation time """
function fmi2SimulateME(c::fmi2Component, dt::Real, t_start::Real = 0.0, t_stop::Real = 1.0, solver = Tsit5(), customFx = nothing)

    if customFx == nothing
        customFx = (x, p, t) -> fx(c, x, p, t)
    end

    eventCb = VectorContinuousCallback((out, x, t, integrator) -> condition(c, out, x, t, integrator),
                                       (integrator, idx) -> affect!(c, integrator, idx),
                                       Int64(c.fmu.modelDescription.numberOfEventIndicators))

    stepCb = FunctionCallingCallback((x, t, integrator) -> stepCompleted(c, x, t, integrator);
                                     func_everystep=true,
                                     func_start=true)

     # First evaluation of the FMU
     x0 = fmi2GetContinuousStates(c)
     x0_nom = fmi2GetNominalsOfContinuousStates(c)

     fmi2SetContinuousStates(c, x0)
     handleEvents(c, true)

     # Get states of handling initial Events
     x0 = fmi2GetContinuousStates(c)
     x0_nom = fmi2GetNominalsOfContinuousStates(c)

     p = []
     problem = ODEProblem(customFx, x0, (t_start, t_stop), p,)
     solution = solve(problem, solver, callback=CallbackSet(eventCb, stepCb))
     #data = solution(ts)
end

function fmi2SimulateME(fmu::FMU2, dt::Real, t_start::Real = 0.0, t_stop::Real = 1.0, solver = Tsit5(), customFx = nothing)
    fmi2SimulateME(fmu.components[end], dt, t_start, t_stop, solver, customFx)
end


type FlatBreeder <: Breeder
  selection::Selection
  mutation::Mutation
  crossover::Crossover

  FlatBreeder(s::Selection, m::Mutation, c::Crossover) =
    new(s, m, c)
end

type FlatBreederDefinition <: BreederDefinition
  selection::SelectionDefinition
  mutation::MutationDefinition
  crossover::CrossoverDefinition

  FlatBreederDefinition() = new()
  FlatBreederDefinition(s::SelectionDefinition, m::MutationDefinition, c::CrossoverDefinition) =
    new(s, m, c)
end

"""
Composes a flat breeder from its definition.
"""
compose!(d::FlatBreederDefinition, sp::Species) =
  FlatBreeder(compose!(d.selection),
              compose!(d.mutation, sp),
              compose!(d.crossover, sp))

"""
TODO: Document flat breeders.
"""
function flat(f::Function)
  def = FlatBreederDefinition()
  f(def)
  def
end

"""
Performs the breeding process for each of the demes within the population
contained by a given state.
"""
function breed!{F}(b::FlatBreeder, sp::Species, buffer::IndividualCollection{F}, members::IndividualCollection{F}, n::Int)
  n_crossover = num_required(b.crossover, n)
  select(b.selection, buffer, sp, members, n_crossover)
  operate!(b.crossover, buffer, n_crossover)
  operate!(b.mutation, buffer, n)
  return
end


module Throttles

abstract type Throttle end

"""
    IntervalThrottle(interval)

Instruct [`with_preview`](@ref) etc. to call `on_preview` only if the
previous invocation of `on_preview` is at least `interval` seconds ago.
"""
struct IntervalThrottle <: Throttle
    interval::Float64
end

init(::IntervalThrottle) = time_ns()
step(th::IntervalThrottle, prev) =
    if (time_ns() - prev) / 1e9 > th.interval
        (true, time_ns())
    else
        (false, prev)
    end

end # module Throttles


__precompile__(true)
module Reversi

export
    # types
    Board, State, Player, RandomPlayer, Minimax, StochasticMinimax,

    # board
    encode, legal_actions_bits, next_state_bits, score, decode, legal_actions,
    game_over,

    # state
    do_action, do_pass, opponent,

    # player
    select_action

const MARKERS = Dict(0 => "   ", 2 => " \u25cf ", 1 => " \u25cb ")

include("board.jl")
include("state.jl")
include("players.jl")
include("gameplay.jl")

include("edax.jl")
# package code goes here

end # module


function define_packed(ty::DataType)
    packed_offsets = cumsum([sizeof(x) for x in ty.types])
    sz = pop!(packed_offsets)
    pushfirst!(packed_offsets, 0)

    @eval begin
        function Base.read(io::IO, ::Type{$ty})
            bytes = read!(io, Array{UInt8}(undef, $sz...))
            hdr = $(Expr(:new, ty, [:(unsafe_load(convert(Ptr{$(ty.types[i])}, pointer(bytes)+$(packed_offsets[i])))) for i = 1:length(packed_offsets)]...,))
            if hdr.sizeof_hdr == ntoh(Int32(348))
                return byteswap(hdr), true
            end
            hdr, false
        end
        function Base.write(io::IO, x::$ty)
            bytes = UInt8[]
            for name in fieldnames($ty)
                append!(bytes, reinterpret(UInt8, [getfield(x,name)]))
            end
            write(io, bytes)
            $sz
        end
    end
    nothing
end

mutable struct NIfTI1Header
    sizeof_hdr::Int32

    data_type::NTuple{10,UInt8}
    db_name::NTuple{18,UInt8}
    extents::Int32
    session_error::Int16
    regular::Int8

    dim_info::Int8
    dim::NTuple{8,Int16}
    intent_p1::Float32
    intent_p2::Float32
    intent_p3::Float32
    intent_code::Int16
    datatype::Int16
    bitpix::Int16
    slice_start::Int16
    pixdim::NTuple{8,Float32}
    vox_offset::Float32
    scl_slope::Float32
    scl_inter::Float32
    slice_end::Int16
    slice_code::Int8
    xyzt_units::Int8
    cal_max::Float32
    cal_min::Float32
    slice_duration::Float32
    toffset::Float32

    glmax::Int32
    glmin::Int32

    descrip::NTuple{80,UInt8}
    aux_file::NTuple{24,UInt8}

    qform_code::Int16
    sform_code::Int16
    quatern_b::Float32
    quatern_c::Float32
    quatern_d::Float32
    qoffset_x::Float32
    qoffset_y::Float32
    qoffset_z::Float32

    srow_x::NTuple{4,Float32}
    srow_y::NTuple{4,Float32}
    srow_z::NTuple{4,Float32}

    intent_name::NTuple{16,UInt8}

    magic::NTuple{4,UInt8}
end
define_packed(NIfTI1Header)

# byteswapping

function byteswap(hdr::NIfTI1Header)
    for fn in fieldnames(typeof(hdr))
        val = getfield(hdr, fn)
        if isa(val, Number) && sizeof(val) > 1
            setfield!(hdr, fn, ntoh(val))
        elseif isa(val, NTuple) && sizeof(eltype(val)) > 1
            setfield!(hdr, fn, map(ntoh, val))
        end
    end
    hdr
end

"""
    freqdim(x)::Int

Returns the frequency dimension.
"""
freqdim(x::NIfTI1Header) = Int(getfield(x, :dim_info) & 0x03)
freqdim(x) = freqdim(header(x))

"""
    phasedim(x)::Int

Returns the phase dimension.
"""
phasedim(x::NIfTI1Header) = Int((getfield(x, :dim_info) >> 2) & 0x03)
phasedim(x) = phasedim(header(x))

"""
    slicedim(x)::Int

Returns the slice dimension.
"""
slicedim(x::NIfTI1Header) = Int((getfield(x, :dim_info) >> 4) & 0x03)
slicedim(x) = slicedim(header(x))

"""
    slice_start(x)::Int

Which slice corresponds to the first slice acquired during MRI acquisition (i.e. not padded slices).
"""
slice_start(x::NIfTI1Header) = Int(getfield(x, :slice_start)) + 1
slice_start(x) = slice_start(header(x))

"""
    slice_end(x)::Int

Which slice corresponds to the last slice acquired during MRI acquisition (i.e. not padded slices).
"""
slice_end(x::NIfTI1Header) = Int(getfield(x, :slice_end)) + 1
slice_end(x) = slice_end(header(x))

"""
    slice_duration(x)

Time to acquire one slice
"""
slice_duration(x::NIfTI1Header) = getfield(x, :slice_duration)
slice_duration(x) = slice_duration(header(x))


module Darwin

export AbstractEvolutionaryModel,
    FitnessValue

abstract type AbstractPopulationModel{G} end

const FitnessValue = Float64

include("utils.jl")
include("individual.jl")
include("rates.jl")
include("fitness.jl")
include("population_model.jl")

include("selection.jl")
include("mutation.jl")
include("crossover.jl")

include("GA.jl")

end # module


struct RelaxAllInOne
    max_width::Int
    merge
    sort_by  # TODO: add default implementation (distance)
end

function Diderot.process(
    relax::RelaxAllInOne, layer::Layer{S,D,V}
) where {S,D,V}
    if length(layer) <= relax.max_width
        return layer
    end

    # Sort states (best ones first).
    candidates = collect(layer)
    sort!(candidates, by=relax.sort_by)

    # Keep first (width - 1) unchanged.
    new_layer = Layer{S,D,V}(Dict(candidates[1:(relax.max_width - 1)]), false)

    # Merge the remaining states into a single relaxed state.
    rest = @view candidates[relax.max_width:end]
    merged_state = foldl(relax.merge, first.(rest))

    # Use predecessor for longest distance of any merged state.
    index = argmax(map(pair -> pair.second.distance, rest))
    merged_node = rest[index].second
    new_node = Node{S,D,V}(merged_node.distance, merged_node.inarc, false)
    new_layer[merged_state] = new_node

    return new_layer
end


"""
Module for providing MPI functionality for `Rimu`. This module is unexported. To use it, run

```
using Rimu.RMPI
```
"""
module RMPI

import MPI
using Rimu
using LinearAlgebra
using Rimu.ConsistentRNG
using Random
using StaticArrays

import Rimu: sort_into_targets!

export MPIData
export mpi_rank, is_mpi_root, @mpi_root, mpi_barrier
export mpi_comm, mpi_root, mpi_size, mpi_seed_CRNGs!, mpi_allprintln

function __init__()
    # Initialise the MPI library once at runtime.
    MPI.Initialized() || MPI.Init()
    # make sure that MPI ranks have independent random numbers
    mpi_seed_CRNGs!()
end

const mpi_registry = Dict{Int,Any}()
# The registry keeps references to distributed structures that must not be
# garbage collected at random
abstract type DistributeStrategy end

include("mpidata.jl")
include("helpers.jl")
include("noexchange.jl")
include("pointtopoint.jl")
include("onesided.jl")
include("alltoall.jl")

end # module RMPI


# Test basic STRIPS functionality in a gripper domain
@testset "strips" begin

path = joinpath(dirname(pathof(PDDL)), "..", "test", "strips")

domain = load_domain(joinpath(path, "domain.pddl"))
@test domain.name == :gripper
@test convert(Term, domain.predicates[:room]) == pddl"(room ?r)"

problem = load_problem(joinpath(path, "problem.pddl"))
@test problem.name == Symbol("gripper-problem")
@test problem.objects == @pddl("rooma", "roomb", "ball1", "ball2", "left", "right")

state = initstate(domain, problem)
implementations = [
    "concrete interpreter" => (domain, state),
    "abstract interpreter" => abstracted(domain, state),
    "ground interpreter" => (ground(domain, state), state),
    "concrete compiler" => compiled(domain, state),
    "abstract compiler" => compiled(abstracted(domain), state)
]

@testset "strips ($name)" for (name, (domain, _)) in implementations

    # Test forward execution of plans
    state = initstate(domain, problem)
    state = execute(domain, state, pddl"(pick ball1 rooma left)")
    @test domain[state => pddl"(carry ball1 left)"] ≃ true
    state = execute(domain, state, pddl"(move rooma roomb)")
    @test domain[state => pddl"(robbyat roomb)"] ≃ true
    state = execute(domain, state, pddl"(drop ball1 roomb left)")
    @test domain[state => pddl"(at ball1 roomb)"] ≃ true

    @test satisfy(domain, state, problem.goal) ≃ true

    # Test action availability
    state = initstate(domain, problem)
    @test Set{Term}(available(domain, state)) == Set{Term}(@pddl(
        "(pick ball1 rooma right)", "(pick ball1 rooma left)",
        "(pick ball2 rooma right)", "(pick ball2 rooma left)",
        "(move rooma roomb)", "(move rooma rooma)"
    ))

end

# Test backward regression of plans
state = goalstate(domain, problem)
state = regress(domain, state, pddl"(drop ball1 roomb left)")
@test domain[state => pddl"(carry ball1 left)"] == true
state = regress(domain, state, pddl"(move rooma roomb)")
@test domain[state => pddl"(robbyat rooma)"] == true
state = regress(domain, state, pddl"(pick ball1 rooma left)")
@test domain[state => pddl"(at ball1 rooma)"] == true
@test issubset(state, initstate(domain, problem))

# Test action relevance
state = goalstate(domain, problem)
@test Set{Term}(relevant(domain, state)) == Set{Term}(@pddl(
    "(drop ball1 roomb left)", "(drop ball1 roomb right)",
    "(drop ball1 roomb ball1)", "(drop ball1 roomb ball2)",
    "(drop ball1 roomb rooma)", "(drop ball1 roomb roomb)"
))

end # strips


# This file was generated, do not modify it. # hide
rfr = @load RandomForestRegressor pkg=DecisionTree

rfr_m = machine(rfr, X, y);


function Base.show(io::IO, diagram::Diagram)
    println(io, "already fixed: ", diagram.partial_sol)
    println(io, "root: ", only(diagram.layers[1]))
    for (l, variable) in enumerate(diagram.variables)
        println(io, "variable: ", variable)
        for pair in diagram.layers[l + 1]
            println(io, " ", pair)
        end
    end
end

function fixed_variables(diagram)
    return vcat(diagram.partial_sol, diagram.variables)
end

function add_transition(layer, new_state, new_node)
    if haskey(layer, new_state)
        if new_node.distance > layer[new_state].distance
            layer[new_state] = new_node
        end
    else
        layer[new_state] = new_node
    end
end

function build_layer(instance, diagram::Diagram{S,D,V}, variable) where {S,D,V}
    layer = Layer{S,D,V}()

    # Collect new states, keep only "best" arcs.
    for (state, node) in diagram.layers[end]
        for (arc, new_state) in transitions(instance, state, variable)
            new_node = Node{S,D,V}(node.distance + arc.value, arc)
            add_transition(layer, new_state, new_node)
        end
    end

    return layer
end

struct InOrder end

function next_variable(instance, diagram, variable_order::InOrder)
    num_variables = length(instance)
    fixed = fixed_variables(diagram)
    for i in 1:num_variables
        if !(i in fixed)  # TODO: efficient!
            return i
        end
    end
    return nothing
end

struct KeepAllNodes end

process(::KeepAllNodes, layer) = layer

"""
    top_down!(diagram::Diagram{S,D,V}, instance; variable_order, processing)

Build a decision diagram for given instance. It is developed layer by layer,
from the top, by following the state transitions. The first layer with a single
root node should already be present in the diagram.

Without any layer processing, the diagram will be exact, that is will contain a
path representing an optimal solution.

Restrictions or relaxations are achieved by dropping or merging the nodes and
states in each layer.
"""
function top_down!(
    diagram::Diagram{S,D,V}, instance;
    variable_order=InOrder(), processing=KeepAllNodes()
) where {S,D,V}
    @assert length(diagram.layers) == 1   # root layer

    # Intermediate layers
    while true
        variable = next_variable(instance, diagram, variable_order)
        if variable === nothing
            break
        end

        layer = build_layer(instance, diagram, variable)
        layer = process(processing, layer)   # restrict/relax

        push!(diagram.layers, layer)
        push!(diagram.variables, variable)
    end

    # Terminal node (last layer reduced to best)
    max_state, max_node = nothing, Node{S,D,V}(typemin(V))
    for (state, node) in diagram.layers[end]
        if node.distance > max_node.distance
            max_state = state
            max_node = node
        end
    end
    diagram.layers[end] = Layer{S,D,V}(Dict(max_state => max_node))
end

"""
    longest_path(diagram::Diagram{S,D,V})

Extract optimal solution by following a longest path from root to terminal.
"""
function longest_path(diagram::Diagram{S,D,V}) where {S,D,V}
    # Collect path in reverse, from terminal to root.
    terminal = only(values(diagram.layers[end].nodes))
    num_variables = length(diagram.partial_sol) + length(diagram.variables)
    decisions = Vector{D}(undef, num_variables)
    node, depth = terminal, length(diagram.layers) - 1
    while depth != 0
        decisions[diagram.variables[depth]] = node.inarc.decision
        state = node.inarc.tail
        node = diagram.layers[depth][state]
        depth -= 1
    end

    return Solution(decisions, terminal.distance)
end

function last_exact_layer(diagram)
    for (l, layer) in enumerate(diagram.layers)
        if !layer.exact
            # Current layer has at least one relaxed node.
            @assert l > 1

            # Return previous layer (all exact)
            return l - 1
        end
    end
    # If we reached the end then even the terminal layer is exact.
    return len(diagram.layers)
end

# TODO: pass solver object to store options, statistics, results?
"""
    branch_and_bound(instance; restrict, relax, variable_order)

Solve given instance using branch-and-bound. For each subproblem, a restriction
is computed to update the incumbent. Then a relaxation is computed. If the
problem can not be pruned from the relaxation bound, new subproblems are created
from an exact cutset in the relaxation's diagram.
"""
function branch_and_bound(instance; restrict, relax, variable_order=InOrder())
    state = initial_state(instance)
    S = typeof(state)
    D = domain_type(instance)
    V = value_type(instance)
    original_problem = Subproblem(Int[], D[], zero(V), state)

    # Assume maximization problems, so relaxation provides upper bound.
    # Following the book, we will use a node's distance (from root) as priority
    # and pick the smallest. This corresponds to picking nodes that are on early
    # layers?!
    #
    # TODO: turn this into a configurable strategy.
    problems = PriorityQueue(original_problem => zero(V))
    incumbent = Solution(D[], typemin(V))

    # Solve subproblems, one at a time.
    while !isempty(problems)
        current = dequeue!(problems)

        root_layer = Layer{S,D,V}(
            Dict(current.state => Node{S,D,V}(current.distance)))

        # solve restriction
        diagram = Diagram{S,D,V}(current.variables, [root_layer], [])
        top_down!(diagram, instance,
                  variable_order=variable_order, processing=restrict)
        solution = longest_path(diagram)

        # update incumbent
        if solution.objective > incumbent.objective
            for (variable, decision) in zip(current.variables, current.decisions)
                solution.decisions[variable] = decision
            end
            incumbent = solution
        end

        # have we solved the subproblem already?
        if all(l -> l.exact, diagram.layers)
            continue
        end

        # solve relaxation
        diagram = Diagram{S,D,V}(current.variables, [root_layer], [])
        top_down!(diagram, instance;
                  variable_order=variable_order, processing=relax)
        solution = longest_path(diagram)

        # create subproblems if not pruned
        if solution.objective > incumbent.objective
            cutset = last_exact_layer(diagram)
            @assert length(diagram.layers[cutset]) > 1
            for (sub_state, sub_node) in diagram.layers[cutset]
                depth = cutset - 1
                new_decisions = Vector{D}(undef, depth)
                node = sub_node
                while depth != 0
                    new_decisions[depth] = node.inarc.decision
                    state = node.inarc.tail
                    node = diagram.layers[depth][state]
                    depth -= 1
                end

                variables = vcat(current.variables,
                                 diagram.variables[1:cutset - 1])
                decisions = vcat(current.decisions, new_decisions)

                subproblem = Subproblem(variables, decisions,
                                        sub_node.distance, sub_state)
                problems[subproblem] = subproblem.distance
            end
        end
    end

    return incumbent
end


using Revise
using HybridFunnelGraphs

graph = create_funnel_graph("test/pddl/domain.pddl", "test/pddl/problem.pddl"; max_levels=60)
graph.num_levels


abstract Layer{A<:AbstractArray, B<:AbstractArray}

# Pointwise layers always use the same dimensionality for input and output
abstract Pointwise{A<:AbstractArray} <: Layer{A, A}

abstract Container{A<:AbstractArray, B<:AbstractArray} <: Layer{A, B}

# By default, accGradParameters! does nothing
function accGradParameters!{A, B}(self :: Layer{A, B}, input :: A,
                                  gradOutput :: B, scale :: eltype(A))
end

# Call the scale version if scale unspecified
function accGradParameters!{A, B}(self :: Layer{A, B}, input :: A,
                                  gradOutput :: B)
    return accGradParameters!(self, input, one(eltype(A)))
end

function forward!{A, B}(self :: Layer{A, B}, input :: A)
    return updateOutput!(self, input)
end
export forward!

function backward!{A, B}(self :: Layer{A, B}, input :: A, gradOutput :: B,
                         scale :: eltype(A))
    updateGradInput!(self, input, gradOutput)
    accGradParameters!(self, input, gradOutput, scale)
    return self.gradInput
end

# Call the scale version if scale unspecified
function backward!{A, B}(self :: Layer{A, B}, input :: A, gradOutput :: B)
    return backward!(self, input, gradOutput, one(eltype(A)))
end
export backward!

# Add a layer to the container
function add_layer!{A, B}(self :: Container{A, B}, layer :: Layer)
    push!(self.modules, layer)
end
export add_layer!

function validate{A, B}(self :: Layer{A, B})
end

function validate{A, B}(self :: Container{A, B})
    for layer in self.modules
        validate(layer)
    end
end
export validate

# Iteration; layers return themselves; containers return themselves and
# children

import Base.start
import Base.done
import Base.next

function start(self :: Layer)
    return false
end

function done(self :: Layer, state)
    return state
end

function next(self :: Layer, state)
    return self, true
end

# state is (0, false) for initial setup (must return self next)
#          (idx, state_for_modules[idx]) afterwards
function start(self :: Container)
    return 0, false
end

function done(self :: Container, state)
    local idx, o_state
    idx, o_state = state
    if idx == 0 then
        return false
    end
    if idx > length(self.modules)
        return true
    end

    # Note that this check is simpler than in the general case, as we know
    # each layer will return at least one element (self). So, as long
    # as idx < length(self.modules), we know we're not done; this would no
    # longer be true if we allowed empty layers.
    return done(self.modules[idx], o_state) && idx == length(self.modules)
end

function next(self :: Container, state)
    local idx, o_state
    idx, o_state = state
    if idx == 0
        return self, (1, isempty(self.modules) ? true : start(self.modules[1]))
    end

    if done(self.modules[idx], o_state)
        idx += 1
        o_state = start(self.modules[idx])
    end

    item, o_state = next(self.modules[idx], o_state)
    return item, (idx, o_state)
end

function zeroGradParameters!(self :: Layer)
    for layer in self
        for field in (:gradWeight, :gradBias)
            if field in names(layer)
                local a = getfield(layer, field)
                fill!(a, zero(eltype(a)))
            end
        end
    end
end
export zeroGradParameters!

using Distributions
function initWeights!(self :: Layer)
    for layer in self
        if :weight in names(layer)
            rand!(Distributions.Uniform(-0.1, 0.1), layer.weight)
        end
        if :bias in names(layer)
            rand!(Distributions.Uniform(0, 0.1), layer.bias)
        end
    end
end
export initWeights!

function getParameters(self :: Layer)
    local params = Array((AbstractArray, AbstractArray), 0)
    for layer in self
        if :weight in names(layer) && :gradWeight in names(layer)
            push!(params, (layer.weight, layer.gradWeight))
        end
        if :bias in names(layer) && :gradBias in names(layer)
            push!(params, (layer.bias, layer.gradBias))
        end
    end
    return params
end
export getParameters

abstract Criterion{A<:AbstractArray, B<:AbstractArray}

function forward!{A,B}(c :: Criterion{A,B}, x :: A, t :: A)
    return updateOutput!(c, x, t)
end

function backward!{A,B}(c :: Criterion{A,B}, x :: A, t :: A)
    return updateGradInput!(c, x, t)
end

#######

type Parallel <: Container
    # TODO: Define my fields here...
end

type SoftMax <: Layer
    # TODO: Define my fields here...
end

type LogSoftMax <: Layer
    # TODO: Define my fields here...
end

type Affine <: Layer
    # TODO: Define my fields here...
end

type SpatialMaxPool <: Layer
    # TODO: Define my fields here...
end

type SpatialConv{T, I <: Integer} <: Layer
    nInputPlane::I
    nOutputPlane::I
    kH::I
    kW::I
    dH::I
    dW::I
    weight::Array{T}
    bias::Array{T}
    grad_weight::Array{T}
    grad_bias::Array{T}
    output::Array{T}
    gradInput::Array{T}
end


"""
    constant_extrapolant(t, integrator::DEIntegrator, idxs, deriv)

Calculate constant extrapolation of derivative `deriv` at time `t` and indices `idxs`
for `integrator`.
"""
function constant_extrapolant(t, integrator::DEIntegrator, idxs, deriv)
    [constant_extrapolant(τ, integrator, idxs, deriv) for τ in t]
end

function constant_extrapolant(t::Number, integrator::DEIntegrator, idxs, T::Type{Val{0}})
    idxs === nothing ? integrator.u : integrator.u[idxs]
end

function constant_extrapolant(t::Number, integrator::DEIntegrator, idxs, T::Type{Val{1}})
    idxs === nothing ? zero(integrator.u)./oneunit(t) : zero(integrator.u[idxs])./oneunit(t)
end

"""
    constant_extrapolant!(val, t, integrator::DEIntegrator, idxs, deriv)

Calculate constant extrapolation of derivative `deriv` at time `t` and indices `idxs`
for `integrator`, and save result in `val` if `val` is not `nothing`.
"""
function constant_extrapolant!(val, t, integrator::DEIntegrator, idxs, deriv)
    [constant_extrapolant!(val, τ, integrator, idxs, deriv) for τ in t]
end

function constant_extrapolant!(val, t::Number, integrator::DEIntegrator, idxs, T::Type{Val{0}})
    if val === nothing
        if idxs === nothing
            return integrator.u
        else
            return integrator.u[idxs]
        end
    elseif idxs === nothing
        @. val = integrator.u
    else
        @views @. val = integrator.u[idxs]
    end
end

function constant_extrapolant!(val, t::Number, integrator::DEIntegrator, idxs, T::Type{Val{1}})
    if val === nothing
        if idxs === nothing
            return zero(integrator.u)./oneunit(t)
        else
            return zero(integrator.u[idxs])./oneunit(t)
        end
    elseif idxs === nothing
        val .= zero(integrator.u)./oneunit(t)
    else
        @views val .= zero(integrator.u[idxs])./oneunit(t)
    end
end


using Abaco
using JSON3
using BenchmarkTools

function nop(s, ne, name, value, inputs)
end

function onresult(ts, ne, name, value, inputs)
    @info "age [$ts]: scope: [$ne] $name = $value"
end

interval = 5
ages = 4
abaco = Abaco.init(nop, interval, ages)
formula(abaco, "y = x1 + x2")

function main(sns)
    for ts in 0:5:80
        for i in 1:100
            ingest(abaco, ts, sns[i].first, :x1, sns[i].second)
            ingest(abaco, ts, sns[i].first, :x2, sns[i].second)
        end
    end
end

sns = Dict(i => "SN$i" => float(i) for i in 1:100)

main(sns)
ingest(abaco, 85, "SN", "x1", 1.0)
abaco.oncomplete = onresult

# expect (0 allocations: 0 bytes)
@btime ingest(abaco, 85, "SN", "x2", 1.0)


############################
### google challenge    ####
### Author F. Clautiaux ####
############################

using DelimitedFiles

# data structure for data used to compute the balance cost
struct BalanceData
    resource1::Int64        # first resource
    resource2::Int64        # second resource
    target::Int64           # target aimed for the recource consumption
    weight::Int64           # weight of the balance cost
end

# data struture for an instance of Google challenge problem
struct DataGoogle
    nbResources::Int64               # number of resources in the instance
    transientStatus::Array{Int64}    # table indicating if each resource is transient (O/1)
    weightLoadCost::Array{Int64}     # table indicating the weight of load cost for each resource

    nbMachines::Int64                # number of machines
    neighborhoods::Array{Int64}      # table indicating the neighborhood of each machine
    nbNeighborhoods::Int64           # total number of neighbordhoods
    locations::Array{Int64}          # table indicating the location of each machine
    nbLocations::Int64               # total number of locations
    softCapacities::Array{Int64,2}   # softcapacities[m,r] indicates the soft cap of machine m for resource r
    hardCapacities::Array{Int64,2}   # hardcapacities[m,r] indicates the hard cap of machine m for resource r
    machineMoveCosts::Array{Int64,2} # machineMoveCosts[m1,m2] indicates the cost for moving from m1 to m2

    nbServices::Int64                # number of services
    spreadMin::Array{Int64}          # for each service, its spread min
    dependences::Array{Array{Int64}} # for each service, a list of dependences (service numbers)

    nbProcess::Int64                 # number of processes
    servicesProcess::Array{Int64}    # for each process  its service
    processReq::Array{Int64,2}       # processReq[p,r] indicates the consumption of resource r by p
    processMoveCost::Array{Int64}    # cost of moving each process p

    nbBalanceCostData::Int64               # number of balance constraints
    balanceCostDataList::Array{BalanceData} # one for each balance target
    processMoveWeight::Int64        # weight for the process move cost
    serviceMoveWeight::Int64        # weight for the service move cost
    machineMoveWeight::Int64        # weight for the machine move cost

    initialAssignment # for each process, id of the machine in the initial solution
end

struct SolutionGoogle
    assignment::Array{Int64}
    cost::Int64
end

function printData(data::DataGoogle)
    println("### Resources part ###")
    println("nb resources = ", data.nbResources)
    println("transient status = ", data.transientStatus)
    println("weight of load cost = ", data.weightLoadCost)

    println("### Machine part ###")
    println("nbMachines = ", data.nbMachines)
    println("Neighborhoods = ", data.neighborhoods)
    println("Locations = ", data.locations)
    println("Hard cap = ", data.hardCapacities)
    println("Soft cap = ", data.softCapacities)
    println("Machine move costs = ", data.machineMoveCosts)

    println("### Service part ###")
    println("Nb services = ", data.nbServices)
    println("Spread min = ", data.spreadMin)
    println("Dependences = ", data.dependences)

    println("### Process part ###")
    println("Nb process = ", data.nbProcess)
    println("Process Services = ", data.servicesProcess)
    println("Proces requirement = ", data.processReq)
    println("Proces move cost = ", data.processMoveCost)

    println("###Weight part ###")
    println("Nb balance constraints = ", data.nbBalanceCostData)
    println("Balance constraints = ", data.balanceCostDataList)
    println("Process move weight : ", data.processMoveWeight)
    println("Service move weight : ", data.serviceMoveWeight)
    println("Machine move weight : ", data.machineMoveWeight)

    println("###Assignment part ###")
    println("Initial assignment : ", data.initialAssignment)
end

function parserGoogle(datafilename::String, assignmentFileName::String, verbose)
    datafile = open(datafilename)
    data = readdlm(datafile)
    close(datafile)

    ### RESOURCES ###
    nbResources = data[1,1]

    transientStatus = data[2:nbResources+1,1]
    weightLoadCost = data[2:nbResources+1,2]

    offset = nbResources + 2 # offset is used to record the current line number in the file

    ### MACHINES ###
    nbMachines = data[offset,1]

    neighborhoods = data[offset+1:offset+nbMachines,1]

    for i in 1:size(neighborhoods,1)
        neighborhoods[i] += 1
    end

    locations = data[offset+1:offset+nbMachines,2]
    for i in 1:size(locations,1)
        locations[i] += 1
    end

     hardcapacities = data[offset+1:offset+nbMachines,3:2+nbResources]
     softcapacities = data[offset+1:offset+nbMachines,3+nbResources:2+2*nbResources]

    machineMoveCosts = data[offset+1:offset+nbMachines,3+2*nbResources:2+2*nbResources+nbMachines]
    offset += nbMachines+1

    ### SERVICES ###
    nbServices = data[offset,1]

    spreadMin = data[offset+1:offset+nbServices,1]

    dependences = [[] for s=1:nbServices]
    for s in 1:nbServices
        nbDep = data[offset+s,2]
        if nbDep != 0
            dependences[s] = data[offset+s,3:2+nbDep]
            for i in 1:nbDep
                dependences[s][i] += 1
            end
        end
    end

    offset += nbServices + 1

    ### PROCESSES ###
    nbProcess = data[offset,1]

    servicesProcess = data[offset+1:offset+nbProcess,1]
    for i in 1:size(servicesProcess,1)
        servicesProcess[i] += 1
    end

    processRequirement = data[offset+1:offset+nbProcess,2:1+nbResources]
    processMoveCost = data[offset+1:offset+nbProcess,2+nbResources]

    offset += nbProcess + 1

    ### BALANCE ###
    nbBalanceCostData = data[offset,1]

    if(nbBalanceCostData > 1)
        println("Only 0 or 1 balance constraint in the files")
        return
    end

    if(nbBalanceCostData == 0)
        balanceConstraints = []
        balanceCostWeight = []
        balanceDataList = []
    else
        balanceConstraints = data[offset+1:offset+nbBalanceCostData,1:3]
        balanceConstraints[1]+=1
        balanceConstraints[2]+=1
        balanceCostWeight = [data[offset+2,1]]
        balanceDataList = [BalanceData(balanceConstraints[1],balanceConstraints[2],balanceConstraints[3],balanceCostWeight[1])]
    end

    offset += 2*nbBalanceCostData+1

    ### MOVE WEIGHTS
    processMoveWeight = data[offset,1]
    serviceMoveWeight = data[offset,2]
    machineMoveWeight = data[offset,3]

    assignmentfile = open(assignmentFileName)
    rawAssignment = readdlm(assignmentfile)
    close(assignmentfile)

    assignment = Array{Int64}(undef,nbProcess)
    for i in 1:nbProcess
        assignment[i] = rawAssignment[i] + 1
    end

    return DataGoogle(nbResources, transientStatus, weightLoadCost,
                      nbMachines, neighborhoods, size(neighborhoods,1), locations, size(locations,1),
                      softcapacities,hardcapacities, machineMoveCosts,
                      nbServices, spreadMin, dependences,
                      nbProcess, servicesProcess,processRequirement, processMoveCost,
                      nbBalanceCostData,
                      balanceDataList,
                      processMoveWeight, serviceMoveWeight, machineMoveWeight,
                      assignment)
end


@coeff function solve!(reactor::Reactor, R_0_guess=nothing)
  @assert !reactor.solved

  @assert isa(R_0, SymEngine.Basic)
  @assert !isa(B_0, SymEngine.Basic)

  if reactor.steady
    solve_steady!(reactor)
  else
    solve_pulsed!(reactor, R_0_guess)
  end

  if isnan(reactor.R_0)
    reactor.valid = false
  else
    update!(reactor)
  end

  reactor.solved = true

  reactor
end


## tape compilation

## CONVERSION TO EXGRAPH

make_name(op::AbstractOp) = Symbol("%$(getid(getvar(op)))")
make_name(var::TAny) = Symbol("%$(getid(var))")

to_exnode(op::Input) = ExNode{:input}(make_name(op), make_name(op); val=getvalue(op))
to_exnode(op::Constant) = ExNode{:constant}(make_name(op), getvalue(op); val=getvalue(op))
to_exnode(op::Assign) = ExNode{:(=)}(make_name(op), make_name(op.src); val=getvalue(op))

function to_exnode(op::Call)
    arg_names = map(make_name, op.args)
    # fns = maybe_to_symbol(op.fn)
    if isempty(op.kwargs)
        ex = Expr(:call, op.fn, arg_names...)
    else
        ex = Expr(:call, op.fn, Espresso.make_kw_params(op.kwargs), arg_names...)
    end
    return ExNode{:call}(make_name(op), ex; val=getvalue(op))
end

function to_exnode(op::Bcast)
    arg_names = map(make_name, op.args)
    #  fns = maybe_to_symbol(op.fn)
    ex = Expr(:., op.fn, Expr(:tuple, arg_names...))
    ExNode{:bcast}(make_name(op), ex; val=getvalue(op))
end

function to_exgraph(tape::Tape)
    g = ExGraph()
    for op in tape
        push!(g, to_exnode(op))
    end
    return g
end

## CODE GENERATION

"""
Rewrite :(Z = X * Y) into :(mul!(Z, X, Y)),  but only if X and Y are arrays
"""
function rewrite_mul(tape::Tape, ex::Expr)
    @assert ex.head == :block
    new_ex = Expr(:block)
    for subex in ex.args
        if (matchingex(:(_R = $*(_X, _Y)), subex)
            && tape[subex.args[1]].val isa AbstractArray
            && all(tape[subex.args[2].args[i]].val isa AbstractArray for i=2:3))
            new_subex = rewrite(subex, :(_Z = $*(_X, _Y)), :($mul!(_Z, _X, _Y)))
            push!(new_ex.args, new_subex)
        else
            push!(new_ex.args, subex)
        end
    end
    return new_ex
end

"""
Generate initialization block of function code
"""
function generate_prologue(tape::Tape)
    ex = Expr(:body)
    for op in tape
        name = make_name(op)
        push!(ex.args, :($name = $(op.var).val))
    end
    return ex
end

"""
Main part of generated function code
"""
function generate_body(tape::Tape)
    ret_var_ids = vcat(values(tape.derivs) |> collect, [tape.resultid])
    ret_var_names = [Symbol("%$id") for id in ret_var_ids]
    exg = to_exgraph(tape)
    # exg = Espresso.fuse_broadcasting(exg)  # TODO: this breaks test_examples.jl
    exg = Espresso.fuse_assigned(exg; outvars=ret_var_names)
    exg = Espresso.eliminate_common(exg)
    ex = Espresso.to_expr_kw(exg)
    # ex = rewrite_all(ex, INPLACE_RULES; phs=Set([:X, :Y, :Z]))
    ex = rewrite_mul(tape, ex)
    return ex
end

"""
Generate last block of function code
"""
function generate_epilogue(tape::Tape)
    ex = Expr(:body)
    for op in tape
        name = make_name(op)
        push!(ex.args, :($(op.var).val = $name))
    end
    return ex
end

function generate_function_expr(tape::Tape)
    fn_ex = :(function $(gensym("tape_fn"))() end)
    fn_ex_body = fn_ex.args[2]
    prologue = generate_prologue(tape)
    body = generate_body(tape)
    epilogue = generate_epilogue(tape)
    # abusing LineNumberNode to insert comments
    push!(fn_ex_body.args, LineNumberNode(0, "prologue"))
    push!(fn_ex_body.args, prologue.args...)
    push!(fn_ex_body.args, LineNumberNode(0, "body"))
    push!(fn_ex_body.args, body.args...)
    push!(fn_ex_body.args, LineNumberNode(0, "epilogue"))
    push!(fn_ex_body.args, epilogue.args...)
    return fn_ex
end

## COMPILATION

function compile(tape::Tape)
    fn_ex = generate_function_expr(tape)
    return Core.eval(@__MODULE__, fn_ex)
end

function compile!(tape::Tape)
    tape.compiled = compile(tape)
end

function rerecord_inputs!(tape::Tape, args...)
    minitape = Tape()
    targs = make_tracked_args(minitape, args...)
    for i=1:length(minitape)
        val = getvalue(minitape[i])
        setvalue!(tape[i], val)
    end
end

function play!(tape::Tape, args...; use_compiled=true)
    rerecord_inputs!(tape, args...)
    if use_compiled && tape.compiled != nothing
        Base.invokelatest(tape.compiled)
    else
        for op in tape
            exec!(tape, op)
        end
    end
end


"""
Lookup the nearest neighbors and teh estimate matrix to get a rough estimate.
"""
function get_travel_time_estimate(halton_nn_tree::BallTree, loc1::LatLonCoords, loc2::LatLonCoords,
                                  estimate_matrix::Matrix{Float64}, time_scale::Float64)

    # Get the nearest neighbor for each location
    # IMP - If they are the same, return 0.0
    loc1_vect = convert_to_vector(loc1)
    loc1_idxs, _ = knn(halton_nn_tree, loc1_vect, 1)
    loc1_idx = loc1_idxs[1]

    loc2_vect = convert_to_vector(loc2)
    loc2_idxs, _ = knn(halton_nn_tree, loc2_vect, 1)
    loc2_idx = loc2_idxs[1]

    # Returns 0.0 if same!
    tt_est = estimate_matrix[loc1_idx, loc2_idx]

    # Use direct distance divided by speed if 0.0
    if tt_est == 0.0
        tt_est = Distances.evaluate(EuclideanLatLongMetric(), convert_to_vector(loc1), convert_to_vector(loc2)) / 0.00777 # Avg. drone speed
    end

    return tt_est/time_scale
end

function generate_package_request(lat_dist::Distribution, lon_dist::Distribution, current_time::Float64,
                                  tw_duration::Float64, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    delivery = (lat = rand(rng, lat_dist), lon = rand(rng, lon_dist))

    # Randomness in walking time and in duration
    start_time = current_time + rand(rng, Distributions.Uniform(tw_duration/2.0, tw_duration))
    duration = rand(rng, Distributions.Uniform(tw_duration/2.0, tw_duration))
    time_window = (start_time, start_time + duration)

    return Package(delivery, time_window)
end

function sample_true_delivery_return_time(drone_pkg_window::SVector{3,Float64}, current_time::Float64,
                                         std_scale::Float64, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    travel_time = drone_pkg_window[3] - drone_pkg_window[2]

    meanval = travel_time
    stdval = travel_time/std_scale
    tt_dist = Distributions.Epanechnikov(meanval, stdval)

    true_delivery_time = current_time + rand(rng, tt_dist)
    true_delivery_time = (true_delivery_time <= drone_pkg_window[1]) ? drone_pkg_window[1] : true_delivery_time

    true_return_time = ceil(true_delivery_time) + rand(rng, tt_dist)

    return (true_delivery_time, true_return_time)
end # end function

# For simulator - Set downtime based on distance between pickup and dropoff - but no uncertainty there
function setup_routing_sim(params_fn::String, halton_nn_tree, estimate_matrix::Matrix{Float64};
                           num_init_requests::Int64=5, new_request_prob::Float64=0.75, delivery_reward::Float64=1000.0,
                           time_window_duration::Float64=30.0, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    # Parse the parameters for the relevant city
    city_params = parse_city_params(params_fn)
    lat_dist = Distributions.Uniform(city_params.lat_start, city_params.lat_end)
    lon_dist = Distributions.Uniform(city_params.lon_start, city_params.lon_end)

    # Update the number of active packages etc.
    active_packages = Dict{String,Package}()

    for n = 1:num_init_requests
        new_package_name = string("pkg",n)

        new_package = generate_package_request(lat_dist, lon_dist, 0.0, time_window_duration, rng)
        active_packages[new_package_name] = new_package
    end

    return RoutingSimulator(current_time=0.0, city_params=city_params, new_request_prob=new_request_prob, delivery_reward=delivery_reward,
                            time_window_duration=time_window_duration, halton_nn_tree=halton_nn_tree, estimate_matrix=estimate_matrix,
                            active_packages=active_packages, num_total_packages=num_init_requests, num_active_packages=num_init_requests)
end

## Need something that maintains the TRUE pickup and dropoff times but only uses them when called upon
# 1. Generate new request based on probability
# 2. If new time is > TRUE pickup, simulate pickup (same for dropoff)
# 3. Assign drones as free or active; Update current pos of drone if picked / dropped off
function update_routing_sim!(sim::RoutingSimulator, server::RoutingAllocation,
                            rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    # First update the simulator and server time
    old_time = sim.current_time
    sim.current_time += 1
    server.current_time += 1

    curr_drone_locs_cols = Tuple{LatLonCoords,Symbol}[]
    curr_sites_locs_cols = Tuple{LatLonCoords,Symbol}[]

    # Simulate any true pickups and dropoffs
    keys_to_del = Set{Tuple{String,String}}()
    busy_packages_to_del = Set{String}()
    done_packages_to_del = Set{String}()
    for (k,v) in server.agent_task_allocation
        drone_nm = k
        package_nm, _ = v

        # Lookup true delivery and return times
        delivery, return_time = sim.true_delivery_return[(drone_nm, package_nm)]

        # For pickup, assign package to drone and mark package inactive
        if delivery > old_time && delivery <= sim.current_time

            server.agent_prop_set[drone_nm].current_package = ""
            server.agent_prop_set[drone_nm].at_depot = false # For good measure

            # Need to check if it reached in time!
            if delivery <= sim.busy_packages[package_nm].time_window[2]
                sim.delivered_packages += 1
                @info "$(drone_nm) has delivered $(package_nm)!"

                # Drone is green and at package location
                push!(curr_drone_locs_cols, (sim.busy_packages[package_nm].delivery, :green))

            else
                sim.late_packages += 1
                @info "$(package_nm) was not delivered in time!"

                # Drone is green and at package location
                push!(curr_drone_locs_cols, (sim.busy_packages[package_nm].delivery, :red))

            end

            push!(busy_packages_to_del, package_nm)
            sim.done_packages[package_nm] = sim.busy_packages[package_nm]

            # Increment loss with difference from start of window
            sim.sum_of_delivery_time += delivery - server.agent_task_windows[(drone_nm, package_nm)][1]

        # For dropoff, free up drone and setup deletion of (drone,package) keys
        elseif return_time > old_time && return_time <= sim.current_time

            # Drone at depot again
            server.agent_prop_set[drone_nm].at_depot = true
            push!(keys_to_del, (drone_nm, package_nm))
            @info "$(drone_nm) back at depot!"
            push!(done_packages_to_del, package_nm)

            # Drone is blue and at depot
            push!(curr_drone_locs_cols, (server.agent_set[drone_nm].depot_loc, :blue))

        else
            ## NOTE: Super hacky, only for plotting
            if haskey(sim.busy_packages, package_nm) == false && haskey(sim.done_packages, package_nm) == false
                continue
            end
            depot_loc = server.agent_set[drone_nm].depot_loc

            if sim.current_time < delivery
                pp_loc = sim.busy_packages[package_nm].delivery
                maxdiff = max(delivery, return_time - delivery)
                interp_factor = (delivery - sim.current_time)/maxdiff
                new_lat = depot_loc.lat + (1.0 - interp_factor)*(pp_loc.lat - depot_loc.lat)
                new_lon = depot_loc.lon + (1.0 - interp_factor)*(pp_loc.lon - depot_loc.lon)
            else
                pp_loc = sim.done_packages[package_nm].delivery
                interp_factor = (return_time - sim.current_time)/(return_time - delivery)
                new_lat = depot_loc.lat + (interp_factor)*(pp_loc.lat - depot_loc.lat)
                new_lon = depot_loc.lon + (interp_factor)*(pp_loc.lon - depot_loc.lon)
                # @infiltrate
            end
            # @infiltrate

            push!(curr_drone_locs_cols, ((lat=new_lat, lon=new_lon), :blue))

        end # if delivery >
    end # for (k,v)

    # Delete bookkeeping keys for dropped off package
    for k in keys_to_del
        delete!(sim.true_delivery_return, k)
        delete!(server.agent_task_windows, k)
        delete!(server.agent_task_allocation, k[1]) # Only delete the agent key
    end

    for r in busy_packages_to_del
        delete!(sim.busy_packages, r)
    end # for r

    for r in done_packages_to_del
        delete!(sim.done_packages, r)
    end

    # Now deal with deliveries that were never attempted?
    packages_to_del = Set{String}()
    for (package_nm, rp) in sim.active_packages
        if rp.time_window[2] <= sim.current_time
            @info "$(package_nm) was not even attempted!"
            sim.late_packages += 1
            push!(packages_to_del, package_nm)

            # Add package location with red
            push!(curr_sites_locs_cols, (rp.delivery, :red))
        end
    end

    for r in packages_to_del
        delete!(sim.active_packages, r)
        sim.num_active_packages -= 1
    end # for r

    # Finally, generate new packages
    if rand(rng) <= sim.new_request_prob
        lat_dist = Distributions.Uniform(sim.city_params.lat_start, sim.city_params.lat_end)
        lon_dist = Distributions.Uniform(sim.city_params.lon_start, sim.city_params.lon_end)
        new_package = generate_package_request(lat_dist, lon_dist, sim.current_time,
                                                sim.time_window_duration, rng)
        sim.num_total_packages += 1
        sim.num_active_packages += 1
        new_package_nm = string("pkg",sim.num_total_packages)
        @info "$(new_package_nm) added!"
        sim.active_packages[new_package_nm] = new_package
    end # if rand(rng)

    # Grey packages
    for (pkg_nm, pp) in sim.active_packages
        push!(curr_sites_locs_cols, (pp.delivery, :grey))
    end
    for (pkg_nm, pp) in sim.busy_packages
        push!(curr_sites_locs_cols, (pp.delivery, :grey))
    end

    sim.curr_drone_site_locs = CurrDroneSiteLocs(curr_drone_locs_cols, curr_sites_locs_cols)
end # function

function update_time_windows!(sim::RoutingSimulator, server::RoutingAllocation)

    @assert sim.current_time == server.current_time

    # We look at all active packages and update their time windows w.r.t the drones
    for (drone_nm,drone_props) in server.agent_prop_set

        interaction_events = server.agent_prop_set[drone_nm].interaction_events
        drone = server.agent_set[drone_nm]

        for (package_nm,rp) in sim.active_packages
            if ~(haskey(server.agent_task_windows, (drone_nm, package_nm)))
                # Compute downtime in terms of dropoff time
                return_time = get_travel_time_estimate(sim.halton_nn_tree, rp.delivery,
                                                        drone.depot_loc, sim.estimate_matrix, sim.time_scale)

                # Insert interaction event timestamps
                ie_timestamps = Dict{MODE,Float64}(START=>rp.time_window[1], FINISH=>rp.time_window[2],
                                                   SUCCESS=>ceil(rp.time_window[2])+return_time)
                push!(interaction_events, InteractionEvent(drone_nm, package_nm, ie_timestamps))

                # Update time windows on the assignment server side
                server.agent_task_windows[(drone_nm, package_nm)] = [rp.time_window[1], rp.time_window[2], ceil(rp.time_window[2])+return_time]
            end # if ~(haskey)
        end # for (package,rp)

        # Sort the temporal event list
        sort!(interaction_events, by=x->x.timestamps[SUCCESS])
        server.agent_prop_set[drone_nm].interaction_events = interaction_events
    end

end # function update_time_windows

function parse_routing_commandline()

    s = ArgParseSettings()

    @add_arg_table s begin
        "--trials", "-n"
            arg_type = Int64
            default = 100
        "--timesteps", "-t"
            arg_type = Int64
            default = 500
        "--n_drones", "-d"
            arg_type = Int64
        "--n_depots", "-e"
            arg_type = Int64
        "--new_request_prob", "-p"
            arg_type = Float64
        "--time_window", "-w"
            arg_type = Float64
        "baseline"
            help = "Which method to evaluate (dispatch/hungarian/css)"
            required = true
        "out_file_name"
            required = true
    end

    return parse_args(s)
end # end function parse routing


Base.@kwdef mutable struct AlgorithmDesc
    name::String
    params = []
    method = nothing
end

function algorithm_gui(algorithm_desc)
    function ()
        algorithms = ["NSGA-II", "MOEA/D"]
        @cstatic item_current = "MOEA/D" begin
            if CImGui.BeginCombo("Algorithm", item_current)
                for n = 0:length(algorithms)-1
                    is_selected = item_current == algorithms[n+1]
                    if CImGui.Selectable(algorithms[n+1], is_selected)
                        item_current = algorithms[n+1]
                        update_algorithm_desc(algorithm_desc, algorithms[n+1])
                    end
                    # set the initial focus when opening the combo (scrolling + for keyboard navigation support in the upcoming navigation branch)
                    is_selected && CImGui.SetItemDefaultFocus()
                end
                CImGui.EndCombo()
            end
        end
        for param_desc in algorithm_desc.params
            eval(param_gui(algorithm_desc, param_desc))
        end
    end
end

function get_algorithm_desc(algorithm_name)
    if algorithm_name == "MOEA/D"
        AlgorithmDesc(name = algorithm_name,
                      params = [ParamDesc(name="N", type="Int",
                                          value=100, default_value=100,
                                          desc="Population size",
                                          min=2, max=10000),
                                ParamDesc(name="T", type="Int",
                                          value=10, default_value=10,
                                          desc="Number of solutions",
                                          min=2, max=100)])
    elseif algorithm_name == "NSGA-II"
        AlgorithmDesc(name = algorithm_name,
                      params = [ParamDesc(name = "N", type = "Int",
                                          value = 100, default_value = 100,
                                          desc = "Population size",
                                          min = 2, max = 100)])
    end
end

function update_algorithm_desc(algorithm_desc, algorithm_name, param_name=nothing, param_value=nothing)
    should_update = algorithm_desc.method === nothing ? true : false
    if algorithm_name !== nothing && algorithm_desc.name != algorithm_name
        new_algorithm_desc = get_algorithm_desc(algorithm_name)
        algorithm_desc.name = new_algorithm_desc.name
        algorithm_desc.params = new_algorithm_desc.params
        should_update = true
    else
        for param in algorithm_desc.params
            if param.name == param_name
                if param.max !== nothing && param_value > param.max
                        param_value = param.max
                    end
                    if param.min !== nothing && param_value < param.min
                        param_value = param.min
                    end
                if param.value != param_value
                    param.value = param_value
                    should_update = true
                end
            end
        end
    end
    if should_update
        if algorithm_desc.name == "MOEA/D"
            N = get_algorithm_param(algorithm_desc, "N")
            T = get_algorithm_param(algorithm_desc, "T")
            algorithm_desc.method = MOEAD(N=N, T=T)
        elseif algorithm_desc.name == "NSGA-II"
            N = get_algorithm_param(algorithm_desc, "N")
            algorithm_desc.method = NSGAII(populationSize=N)
        end
    end
end

function get_algorithm_param(algorithm_desc, param_name)
    for param in algorithm_desc.params
        if param.name == param_name
            return param.value === nothing ? param.default_value : param.value
        end
    end
end


using Test

using CiphertextStealing
using Random: randstring

@testset "Smoke" begin
    for text_len in 1:100
        plaintext = Vector{UInt8}(randstring(text_len))
        iv = rand(UInt8, 16)
        key = rand(UInt8, 16)

        iv_backup = copy(iv)
        ciphertext = aes_128_cbc_encode_cts!(plaintext, key, iv)
        @test length(ciphertext) == text_len
        @test (text_len > 16) == (iv == iv_backup)
        iv_backup = copy(iv)
        deciphered = aes_128_cbc_decode_cts!(ciphertext, key, iv)
        @test (text_len > 16) == (iv == iv_backup)
        @test deciphered == plaintext
    end
end


using Logging
using ArgParse
using MPI

using QXContexts

"""
    main(ARGS)

QXContexts entry point
"""
function main(args)
    s = ArgParseSettings("QXContexts")
    @add_arg_table! s begin
        "--sub-comm-size", "-s"
            help = "The number of ranks to assign to each sub-communicator for partitions"
            default = 1
            arg_type = Int
        "--mpi", "-m"
            help = "Use MPI"
            action = :store_true
        "--gpu", "-g"
            help = "Use GPU if available"
            action = :store_true
        "--verbose", "-v"
            help = "Enable logging"
            action = :store_true
    end
    parsed_args = parse_args(args, s)

    if parsed_args["verbose"]
        if parsed_args["mpi"]
            if !MPI.Initialized() MPI.Init() end
            global_logger(QXContexts.Logger.QXLoggerMPIPerRank())
        else
            global_logger(QXContexts.Logger.QXLogger())
        end
    end

    file_path      = @__DIR__
    dsl_file       = joinpath(file_path, "ghz/ghz_5.qx")
    param_file     = joinpath(file_path, "ghz/ghz_5.yml")
    input_file     = joinpath(file_path, "ghz/ghz_5.jld2")
    output_file    = joinpath(file_path, "ghz/out.jld2")

    results = execute(dsl_file, input_file, param_file, output_file;
                      use_mpi=parsed_args["mpi"],
                      sub_comm_size=parsed_args["sub-comm-size"],
                      use_gpu=parsed_args["gpu"])
    if parsed_args["verbose"]
        @info results
    end
end

main(ARGS)


module MDTrajAnalyse

export AbstractTrajectory,
       Trajectory,
       natoms,
       getatoms,
       dihedral,
       distances,
       distances!,
       cellvolume,
       angletoframe,
       sphericalview,
       subtrajectory,
       compute_rdf,
       get_close_atoms,

       AbstractAtomNames,
       AtomNames,
       NoAtomName,

       read_pdb,
       read_trajectory,
       read_dcd,
       read_xyz,
       write_xyz,

       AbstractUnitCell,
       AbtractPeriodicCell,
       AbstractOrthorombicCell,
       CubicCell,
       NonPeriodic,
       OrthorombicCell,
       TriclinicCell,
       VariableCell,
       celldiag,
       cellmatrix,
       parsecell,
       volume

include("cell.jl")
include("trajectory.jl")
include("fileaccess.jl")
include("properties.jl")
include("SubModules/visualize.jl")

end # module


"""
abstract type Geometry end
"""

struct NoGeometry{C} <: Geometry
    capacitance::C
end
capacitance(g::NoGeometry) = g.capacitance

######### Geometries ###########
# is_geometric(::Compartment{T}) = T
# is_geometric(::Component) = false
"""
Every geometry needs methods for:
    capacitance()
    maybe reversals as dictionary?
    initial_conditions(f::Species)
"""