module PiecewiseIncreasingRanges
export PiecewiseIncreasingRange, NoNearestSampleError, findnearest, resample

using Compat, Base.Order

function combine_ranges{R<:Range}(ranges::Vector{R}, firstrg::R, firstrgidx::Int)
    newranges = R[]
    offsets = Int[1]

    step(firstrg) < 0 && throw(ArgumentError("ranges must be strictly monotonically increasing"))
    currg = firstrg

    for i = firstrgidx:length(ranges)
        newrg = ranges[i]
        isempty(newrg) && continue

        if step(newrg) == step(currg) && first(newrg) == last(currg)+step(currg)
            # Can extend current range
            currg = first(currg):step(currg):last(newrg)
        else
            if first(newrg) <= last(currg) || step(newrg) < 0
                throw(ArgumentError("ranges must be strictly monotonically increasing"))
            end

            # Need to make a new range
            push!(newranges, currg)
            push!(offsets, offsets[end]+length(currg))
            currg = newrg
        end
    end

    push!(newranges, currg)
    (newranges, offsets)
end

function combine_ranges{R<:FloatRange}(ranges::Vector{R}, firstrg::R, firstrgidx::Int)
    newranges = R[]
    offsets = Int[1]

    if signbit(firstrg.step) != signbit(firstrg.divisor)
        throw(ArgumentError("ranges must be strictly monotonically increasing"))
    end
    currg = firstrg

    for i = firstrgidx:length(ranges)
        newrg = ranges[i]
        isempty(newrg) && continue

        if step(newrg) == step(currg) && newrg.start*currg.divisor == (currg.start+currg.step*currg.len)*newrg.divisor
            currg = FloatRange(currg.start, currg.step, currg.len + newrg.len, currg.divisor)
        else
            if first(newrg) <= last(currg) || signbit(newrg.step) != signbit(newrg.divisor)
                throw(ArgumentError("ranges must be strictly monotonically increasing"))
            end

            # Need to make a new range
            push!(newranges, currg)
            push!(offsets, offsets[end]+length(currg))
            currg = newrg
        end
    end

    push!(newranges, currg)
    (newranges, offsets)
end

immutable PiecewiseIncreasingRange{T,R<:Range,S} <: AbstractVector{T}
    ranges::Vector{R}
    offsets::Vector{Int}
    divisor::S

    function PiecewiseIncreasingRange(ranges::Vector{R}, divisor)
        isempty(ranges) && return new(ranges, Int[])

        # Find first non-empty range
        firstrg = ranges[1]
        j = 0
        for j = 1:length(ranges)
            firstrg = ranges[j]
            !isempty(firstrg) && break
        end
        isempty(firstrg) && return new(R[], Int[], divisor)

        newranges, offsets = combine_ranges(ranges, firstrg, j+1)
        new(newranges, offsets, divisor)
    end
end
PiecewiseIncreasingRange{R<:Range}(ranges::Vector{R}, divisor) = PiecewiseIncreasingRange{typeof(inv(one(eltype(R)))),R,typeof(divisor)}(ranges, divisor)
PiecewiseIncreasingRange{R<:Range}(ranges::Vector{R}) = PiecewiseIncreasingRange{eltype(R),R,@compat(Void)}(ranges, nothing)

# Avoid applying the divisor if it is one, to get types right
divide_divisor{T,R}(r::PiecewiseIncreasingRange{T,R,@compat(Void)}, x) = x
multiply_divisor{T,R}(r::PiecewiseIncreasingRange{T,R,@compat(Void)}, x) = x
divide_divisor{T,R,S}(r::PiecewiseIncreasingRange{T,R,S}, x) = x/r.divisor
multiply_divisor{T,R,S}(r::PiecewiseIncreasingRange{T,R,S}, x) = x*r.divisor

function Base.size(r::PiecewiseIncreasingRange)
    isempty(r.ranges) && return 0
    return (r.offsets[end]+length(r.ranges[end])-1,)
end

function Base.getindex(r::PiecewiseIncreasingRange, i::Integer)
    rgidx = searchsortedlast(r.offsets, i, Forward)
    divide_divisor(r, r.ranges[rgidx][i-r.offsets[rgidx]+1])
end

function Base.getindex{T,R,S}(r::PiecewiseIncreasingRange{T,R,S}, x::Range{Int})
    isempty(x) && return PiecewiseIncreasingRange{T,R,S}(R[], Int[], r.divisor)
    (first(x) >= 1 && last(x) <= length(r)) || throw(BoundsError())

    # Not type-stable!
    step(x) < 0 && return invoke(getindex, (typeof(r), AbstractVector{Int}), r, x)

    firstrgidx = searchsortedlast(r.offsets, first(x), Forward)
    lastrgidx = searchsortedlast(r.offsets, last(x), Forward)
    newrgs = Array(R, lastrgidx-firstrgidx+1)
    for irange = firstrgidx:lastrgidx
        elmax = min(last(x)-r.offsets[irange]+1, length(r.ranges[irange]))
        newrg = newrgs[irange-firstrgidx+1] = r.ranges[irange][first(x)-r.offsets[irange]+1:elmax]
        x = x[length(newrg)+1:end]
    end
    PiecewiseIncreasingRange(newrgs, r.divisor)
end

function resample{T,R,S}(r::PiecewiseIncreasingRange{T,R,S}, ratio::Rational{Int})
    excess = zero(eltype(R))
    curpt = 0
    newrgs = R[]
    for i = 1:length(r.ranges)
        endpt = last(r.ranges[i])*num(ratio)
        newrg = first(r.ranges[i])*num(ratio)+excess*step(r.ranges[i]):step(r.ranges[i])*den(ratio):endpt
        curpt += length(newrg)
        push!(newrgs, newrg)
        if i != length(r.ranges)
            # Need to linearly interpolate between endpoint and next range
            # nextind is the next index in r (= 1 + curpt/ratio) times num(ratio)
            nextind = num(ratio) + curpt*den(ratio)
            if nextind < r.offsets[i+1]*num(ratio)
                # Need to interpolate between this point and the next
                weight = mod(nextind, num(ratio))
                interpt = last(r.ranges[i])*(num(ratio)-weight) + first(r.ranges[i+1])*weight
                interpstep = (first(r.ranges[i+1]) - last(r.ranges[i]))*den(ratio)
                newrg = interpt:interpstep:first(r.ranges[i+1])*num(ratio)
                push!(newrgs, newrg)
                curpt += length(newrg)
                nextind += length(newrg)*den(ratio)
            end
            excess = nextind - r.offsets[i+1]*num(ratio)
        end
    end
    PiecewiseIncreasingRange(newrgs, multiply_divisor(r, num(ratio)))
end

# searchsortedfirst, searchsortedlast
immutable PiecewiseIncreasingRangeFirstOrdering <: Ordering end
Base.Order.lt(o::PiecewiseIncreasingRangeFirstOrdering, a, b) = isless(first(a), first(b))
immutable PiecewiseIncreasingRangeLastOrdering <: Ordering end
Base.Order.lt(o::PiecewiseIncreasingRangeLastOrdering, a, b) = isless(last(a), last(b))

function Base.searchsortedfirst(r::PiecewiseIncreasingRange, x)
    isempty(r.ranges) && return 1
    xd = multiply_divisor(r, x)

    rgidx = searchsortedfirst(r.ranges, xd, PiecewiseIncreasingRangeLastOrdering())
    rgidx > length(r.ranges) && return length(r) + 1
    searchsortedfirst(r.ranges[rgidx], xd, Forward) + r.offsets[rgidx] - 1
end

function Base.searchsortedlast(r::PiecewiseIncreasingRange, x)
    isempty(r.ranges) && return 1
    xd = multiply_divisor(r, x)

    rgidx = searchsortedlast(r.ranges, xd, PiecewiseIncreasingRangeFirstOrdering())
    rgidx == 0 && return 0
    searchsortedlast(r.ranges[rgidx], xd, Forward) + r.offsets[rgidx] - 1
end

immutable NoNearestSampleError <: Exception end

function findnearest(r::PiecewiseIncreasingRange, x, within_half_step::Bool=false)
    isempty(r.ranges) && throw(NoNearestSampleError())
    xd = multiply_divisor(r, x)

    rgidx = searchsortedfirst(r.ranges, xd, PiecewiseIncreasingRangeLastOrdering())
    if rgidx > length(r.ranges)
        rgend = r.ranges[end]
        within_half_step && xd > rgend[end]+step(rgend)/2 && throw(NoNearestSampleError())
        return length(r)
    end

    rg = r.ranges[rgidx]
    idxinrg = searchsortedfirst(rg, xd, Forward)
    idx = idxinrg + r.offsets[rgidx] - 1
    d = rg[idxinrg] - xd

    if idxinrg == 1
        if rgidx == 1
            # First element of all elements
            within_half_step && d > step(rg)/2 && throw(NoNearestSampleError())
            return 1
        end

        # Could be closer to last element of preceding range
        rgprev = r.ranges[rgidx-1]
        dprev = xd - rgprev[end]
        within_half_step && d > step(rg)/2 && dprev > step(rgprev)/2 && throw(NoNearestSampleError())
        return idx - (dprev <= d)
    end

    ifelse(d >= step(rg)/2, idx-1, idx)
end

end # module


using Pkg: @pkg_str
pkg"activate ."

#---

using Cassette
using JuliaInterpreter: JuliaInterpreter, BreakpointSignature, BreakpointRef
using MagneticReadHead
using Debugger
using Dates

using JuliaInterpreter: Frame, enter_call, get_return
@inline function run_interpretted(f, args...; kwargs...)
    fr = enter_call(f, args...)
    cmd = :c
    while true
        ret = JuliaInterpreter.debug_command(JuliaInterpreter.finish_and_return!, fr, cmd)
        ret == nothing && break
        fr, pc = ret
    end
    return get_return(fr)
end

Cassette.@context MixModeDebugCtx

function setup_overdubs(breakpoints=JuliaInterpreter._breakpoints)
    for bp in breakpoints
        create_overdub(bp)
    end
end

function create_overdub(bp)
    ft = MagneticReadHead.functiontypeof(bp.f)
    #TODO: select based on all the method sig info, not just the function
    @eval @inline function Cassette.overdub(::MixModeDebugCtx, f::$ft, args...)
        run_interpretted(f, args...)
    end
end

function run_mixedmode(f, args)
    Cassette.recurse(MixModeDebugCtx(), f, args)
end
##############

using Base: gc_num, time_ns, GC_Diff, gc_alloc_count, time_print, show_unquoted
macro time_trial(expr)
    run = quote
        GC.gc()
        local stats = gc_num()
        local elapsedtime = time_ns()
        local val = $(esc(expr))
        elapsedtime = time_ns() - elapsedtime
        local diff = GC_Diff(gc_num(), stats)
        time_print(elapsedtime, diff.allocd, diff.total_time,
                   gc_alloc_count(diff))
        println()
    end

    show_op = :(println($(sprint(show_unquoted, expr))))

    return quote
        $show_op
        print(" - 1st Run: ")
        $run
        print(" - 2nd Run: ")
        $run
        print(" - 3nd Run: ")
        $run
        println()
    end
end

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

# Setup:
function resetup()
    # Reinclude so that it recompiles
    include("../test/demo_funs.jl")
    JuliaInterpreter.remove()
    @breakpoint exp(1.0)
    setup_overdubs()  # for MixedModeDebugger
end

const x = rand(1_000, 500)
# Timing
resetup()
@time_trial winter(x)
resetup()
@time_trial run_interpretted(winter, x)
resetup()
@time_trial run_mixedmode(winter, x)


# matlab estimate of isi = 0.0177
# julia/glfw estimate of isi = 0.0166 (monitor is at 60Hz. 1/60 = 0.016666...)
#
# a basic draw test was:
#	- draw 100 randomly placed ellipses, with random widths and heights
#	- make sure each ellipse has a randomly assigned fill color and stroke
#	  color
# for this test:
#	fastest matlab roundtrip time = ~0.007860 secs
#	fastest julia roundtrip time = ~0.001945 secs
#
# julia is faster (~4x faster) and its timing fluctuates much less than
# matlab's. plus, it appears that julia accurately syncs to vertical refresh,
# even on OS X.
#
# (however, i'm not sure if this is an illusion or a mistake on my part...)
#
# anyway, "test.tsv" contains timing data for Processing.jl that was generated
# by the script below. the times are how long it takes from the top of the
# animation loop to the screen being updated. if we have proper vsync, then
# this should always be 0.01666... secs. "test.tsv" shows that this is the
# case, at an average frame time of 0.016649 secs with a fluctuation error of
# ~0.0012 secs.
#
# note that, all of this here is without any of the special psychtoolbox tricks
# that try to make stuff sync with the monitor via tweaks. its just standard
# opengl and glfw commands, although we are using the commands that are some
# of the fastest for rendering.
#
# the julia code is also easier to write, understand, and maintain.

using Processing

window = Processing.screen()
ftbf = zeros(600)
ft = zeros(600)
x = 1
t = 0

image = Processing.loadImage(Pkg.dir()*"/Processing/test/image.png")
tex, w, h = Processing.texture(image)

@time while x < 600
	st = time_ns()
	Processing.background(0.94, 0.92, 0.9, 1.0)
	Processing.strokecol(0, 0, 0, 1)
	Processing.fillcol(0.7, 0.7, 0.7, 1.0)
	Processing.ellipse(0, 0, 0.3, 0.5)
	Processing.ellipse(-.5, .5, 0.2, 0.2)
	Processing.triangle(.3, .75, .58, .20, .86, .75)
	Processing.triangle(.3, .6, .58, .40, .86, .7)
	Processing.strokeWeight(5)
	Processing.point(randn(10), randn(10))
	Processing.text("Processing.jl", 0.25, -0.85)
	Processing.strokeWeight(1)
	Processing.fillcol(0, 0, 0.9, 1)
	Processing.quad(-.3, -.75, -.58, -.20, -.86, -.75, -.2, -.4)
	Processing.line(.5, -.4, .7, -.5)
	Processing.noFill()
	Processing.rect(-.6, -.4, .2, .5)
	Processing.line(.7, -.4, .5, -.5)
	Processing.strokecol(0.9, 0, 0, 1)
	Processing.arc(-.6, .6, .3, .3, 0.4pi, pi)
	Processing.strokecol(0, 0.9, 0, 1)
	Processing.fillcol(1.0, 1.0, 1.0, 1.0)
	Processing.ellipse(0, 0, 0.2, 0.2)
	Processing.fillcol(sin(t)/2+0.5, cos(t)/2+0.5, (sin(t)*cos(t))/2+0.5, 1.0)
	Processing.pushMatrix()
	Processing.rotateY(pi/10*t)
	Processing.rotateX(pi/10*t)
	Processing.translate(0.2, -0.5, 0)
	Processing.box(0.15)
	Processing.popMatrix()
	Processing.noStroke()
	Processing.fillcol(1, 1, 1, 1)
	Processing.rect(-0.1, 0.6, 0.2, 0.2, tex)
	if Processing.keyPress(window, GLFW.KEY_SPACE)
		Processing.save("screenshot.tiff")
		println("key pressed and screenshot saved.")
	end
	if Processing.mousePress(window, GLFW.MOUSE_BUTTON_LEFT)
		println("bye!")
		break
	end
	ftbf[Processing.frameCount()+1] = time_ns() - st
	Processing.animate(window)
	ft[Processing.frameCount()] = time_ns() - st
	t += 1/60
	x += 1
end

Processing.endDrawing(window)


module SecureSessions

using HttpServer
using MbedTLS

export
# Secure cookies
    create_secure_session_cookie,
    get_session_cookie_data,
    invalidate_cookie!,
# Password hashing
    StoredPassword,
    password_is_valid,
    password_is_valid,
# utils
    username_is_permissible,
    password_is_permissible

include("utils.jl")
include("secure_cookies.jl")
include("password_hash.jl")
include("pbkdf2.jl")

# Globals
session_timeout = 5 * 60 * 1000        # Duration of a session's validity in milliseconds
key_length      = 32                   # Key length for AES 256-bit cipher in CBC mode
block_size      = 16                   # IV  length for AES 256-bit cipher in CBC mode
const_key       = csrng(key_length)    # Symmetric key for encrypting secret_keys (with 256-bit encryption)
const_iv        = csrng(block_size)    # IV for encrypting secret_keys
http_only       = true
encrypted_sessions_only = true
timeout_str     = string(convert(Int64, 0.001 * session_timeout))    # Session timeout in seconds, represented as a string

end # module


using ArcadeLearningEnvironment
using ArgParse
using Cambrian
using CartesianGeneticProgramming
using Dates
using IICGP
using Random

# function extension
import Cambrian.mutate

```
Playing Atari games using DualCGP on screen input values

This uses a single Game seed, meaning an unfair deterministic Atari
environment. To evolve using a different game seed per generation, add in
reset_expert=true and seed=evo.gen below.
```

s = ArgParseSettings()
@add_arg_table! s begin
    "--cfg"
    help = "configuration script"
    default = "cfg/test_dual.yaml"
    "--game"
    help = "game rom name"
    default = "assault"
    "--seed"
    help = "random seed for evolution"
    arg_type = Int
    default = 0
    "--out"
    help = "output directory"
    arg_type = String
    default = dirname(@__DIR__)
    "--ind"
    help = "individual for evaluation"
    arg_type = String
    default = ""
end
args = parse_args(ARGS, s)
const game = args["game"]
const seed = args["seed"]
const resdir = args["out"]
Random.seed!(seed)

main_cfg, enco_cfg, cont_cfg, reducer, bootstrap = IICGP.dualcgp_config(
    args["cfg"], game
)
const max_frames = main_cfg["max_frames"]
const stickiness = main_cfg["stickiness"]
const grayscale = main_cfg["grayscale"]
const downscale = main_cfg["downscale"]
const logid = enco_cfg.id
const state_ref = get_state_ref(game, seed)

function play_atari(
    encoder::CGPInd,
    reducer::Reducer,
    controller::CGPInd,
    lck::ReentrantLock;
    rom=game,
    seed=seed,
    rom_state_ref=state_ref,
    max_frames=max_frames,
    grayscale=grayscale,
    downscale=downscale,
    stickiness=stickiness
)
    # Random.seed!(seed)
    mt = MersenneTwister(seed)
    game = Game(rom, seed, lck=lck, state_ref=rom_state_ref)
    IICGP.reset!(reducer) # zero the buffers
    reward = 0.0
    frames = 0
    prev_action = Int32(0)
    while ~game_over(game.ale)
        if rand(mt) > stickiness || frames == 0
            s = get_state(game, grayscale, downscale)
            output = IICGP.process(encoder, reducer, controller, s)
            action = game.actions[argmax(output)]
        else
            action = prev_action
        end
        reward += act(game.ale, action)
        frames += 1
        prev_action = action
        if frames > max_frames
            break
        end
    end
    close!(game)
    [reward]
end

if length(args["ind"]) > 0
    ind = CGPInd(cfg, read(args["ind"], String))
    ftn = fitness(ind, inps, outs)
    println("Fitness: ", ftn)
else
    # Define mutate and fit functions
    function mutate(ind::CGPInd, ind_type::String)
        if ind_type == "encoder"
            return goldman_mutate(enco_cfg, ind, init_function=IPCGPInd)
        elseif ind_type == "controller"
            return goldman_mutate(cont_cfg, ind)
        end
    end
    lck = ReentrantLock()
    fit(encoder::CGPInd, controller::CGPInd) = play_atari(encoder, reducer,
                                                          controller, lck)
    # Create an evolution framework
    e = IICGP.DualCGPEvolution(enco_cfg, cont_cfg, fit, resdir,
                               encoder_init_function=IPCGPInd, logid=logid,
                               bootstrap=bootstrap, game=game)
    # Run evolution
    init_backup(logid, resdir, args["cfg"])
    run!(e)
end


module OrdinaryDiffEq

  using Reexport
  @reexport using DiffEqBase

  using Logging

  using MuladdMacro, SparseArrays, FastClosures

  using LinearAlgebra

  # Interfaces
  import DiffEqBase: solve!, step!, initialize!, isadaptive

  # Internal utils
  import DiffEqBase: ODE_DEFAULT_NORM, ODE_DEFAULT_ISOUTOFDOMAIN, ODE_DEFAULT_PROG_MESSAGE, ODE_DEFAULT_UNSTABLE_CHECK

  using DiffEqBase: DiffEqArrayOperator, DEFAULT_UPDATE_FUNC

  using DiffEqBase: TimeGradientWrapper, UJacobianWrapper, TimeDerivativeWrapper, UDerivativeWrapper

  import RecursiveArrayTools: chain, recursivecopy!

  using UnPack, GenericSVD, ForwardDiff, RecursiveArrayTools,
        DataStructures, FiniteDiff, ArrayInterface

  import ForwardDiff.Dual

  using ExponentialUtilities

  using NLsolve
  # Required by temporary fix in not in-place methods with 12+ broadcasts
  # `MVector` is used by Nordsieck forms
  import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray

  # Integrator Interface
  import DiffEqBase: resize!,deleteat!,addat!,full_cache,user_cache,u_cache,du_cache,
                     resize_non_user_cache!,deleteat_non_user_cache!,addat_non_user_cache!,
                     terminate!,get_du, get_dt,get_proposed_dt,set_proposed_dt!,
                     u_modified!,savevalues!,add_tstop!,add_saveat!,set_reltol!,
                     set_abstol!, postamble!, last_step_failed,
                     isautodifferentiable

  using DiffEqBase: check_error!, @def, @.. , _vec, _reshape

  using DiffEqBase: AbstractNLSolverAlgorithm, AbstractNLSolverCache, NLStatus
  using DiffEqBase: nlsolve_f, qrdelete!, qradd!, build_jac_config, resize_jac_config!
  using DiffEqBase: Convergence, Divergence
  const TryAgain = DiffEqBase.SlowConvergence

  import DiffEqBase: calculate_residuals, calculate_residuals!, unwrap_cache, @tight_loop_macros, islinear, timedepentdtmin

  import SparseDiffTools
  import SparseDiffTools: matrix_colors, forwarddiff_color_jacobian!,
                          forwarddiff_color_jacobian, ForwardColorJacCache,
                          default_chunk_size, getsize

  using MacroTools, Adapt

  const CompiledFloats = Union{Float32,Float64,
    ForwardDiff.Dual{ForwardDiff.Tag{T,W},K,3} where {T,W<:Union{Float64,Float32},
                                                        K<:Union{Float64,Float32}}}

  include("misc_utils.jl")
  include("algorithms.jl")
  include("alg_utils.jl")

  include("nlsolve/type.jl")
  include("nlsolve/utils.jl")
  include("nlsolve/nlsolve.jl")
  include("nlsolve/functional.jl")
  include("nlsolve/newton.jl")

  include("generic_rosenbrock.jl")

  include("caches/basic_caches.jl")
  include("caches/low_order_rk_caches.jl")
  include("caches/high_order_rk_caches.jl")
  include("caches/low_storage_rk_caches.jl")
  include("caches/ssprk_caches.jl")
  include("caches/feagin_caches.jl")
  include("caches/verner_caches.jl")
  include("caches/sdirk_caches.jl")
  include("caches/firk_caches.jl")
  include("caches/kencarp_kvaerno_caches.jl")
  include("caches/linear_caches.jl")
  include("caches/linear_nonlinear_caches.jl")
  include("caches/symplectic_caches.jl")
  include("caches/rosenbrock_caches.jl")
  include("caches/rkn_caches.jl")
  include("caches/adams_bashforth_moulton_caches.jl")
  include("caches/nordsieck_caches.jl")
  include("caches/bdf_caches.jl")
  include("caches/rkc_caches.jl")
  include("caches/extrapolation_caches.jl")
  include("caches/prk_caches.jl")
  include("caches/pdirk_caches.jl")
  include("caches/dae_caches.jl")

  include("tableaus/low_order_rk_tableaus.jl")
  include("tableaus/high_order_rk_tableaus.jl")
  include("tableaus/symplectic_tableaus.jl")
  include("tableaus/verner_tableaus.jl")
  include("tableaus/feagin_tableaus.jl")
  include("tableaus/rosenbrock_tableaus.jl")
  include("tableaus/sdirk_tableaus.jl")
  include("tableaus/firk_tableaus.jl")
  include("tableaus/rkn_tableaus.jl")
  include("tableaus/rkc_tableaus.jl")

  include("integrators/type.jl")
  include("integrators/controllers.jl")
  include("integrators/integrator_utils.jl")
  include("cache_utils.jl")
  include("integrators/integrator_interface.jl")
  include("initialize_dae.jl")
  include("wrappers.jl")

  include("perform_step/fixed_timestep_perform_step.jl")
  include("perform_step/symplectic_perform_step.jl")
  include("perform_step/rkn_perform_step.jl")
  include("perform_step/split_perform_step.jl")
  include("perform_step/linear_perform_step.jl")
  include("perform_step/exponential_rk_perform_step.jl")
  include("perform_step/explicit_rk_perform_step.jl")
  include("perform_step/low_order_rk_perform_step.jl")
  include("perform_step/high_order_rk_perform_step.jl")
  include("perform_step/verner_rk_perform_step.jl")
  include("perform_step/feagin_rk_perform_step.jl")
  include("perform_step/low_storage_rk_perform_step.jl")
  include("perform_step/ssprk_perform_step.jl")
  include("perform_step/sdirk_perform_step.jl")
  include("perform_step/kencarp_kvaerno_perform_step.jl")
  include("perform_step/firk_perform_step.jl")
  include("perform_step/rosenbrock_perform_step.jl")
  include("perform_step/composite_perform_step.jl")
  include("perform_step/adams_bashforth_moulton_perform_step.jl")
  include("perform_step/nordsieck_perform_step.jl")
  include("perform_step/bdf_perform_step.jl")
  include("perform_step/rkc_perform_step.jl")
  include("perform_step/extrapolation_perform_step.jl")
  include("perform_step/prk_perform_step.jl")
  include("perform_step/pdirk_perform_step.jl")
  include("perform_step/dae_perform_step.jl")

  include("dense/generic_dense.jl")
  include("dense/interpolants.jl")
  include("dense/rosenbrock_interpolants.jl")
  include("dense/stiff_addsteps.jl")
  include("dense/low_order_rk_addsteps.jl")
  include("dense/verner_addsteps.jl")
  include("dense/high_order_rk_addsteps.jl")

  include("derivative_utils.jl")
  include("nordsieck_utils.jl")
  include("adams_utils.jl")
  include("bdf_utils.jl")
  include("rkc_utils.jl")
  include("derivative_wrappers.jl")
  include("iterator_interface.jl")
  include("constants.jl")
  include("composite_solution.jl")
  include("solve.jl")
  include("initdt.jl")
  include("interp_func.jl")
  include("composite_algs.jl")

  #General Functions
  export solve, solve!, init, step!

  export OrdinaryDiffEqAlgorithm

  #Callback Necessary
  export addsteps!, ode_interpolant, terminate!, savevalues!, copyat_or_push!, isfsal

  export constructDormandPrince

  # Reexport the Alg Types

  export FunctionMap, Euler, Heun, Ralston, Midpoint, RK4, ExplicitRK, OwrenZen3, OwrenZen4, OwrenZen5,
         BS3, BS5, RK46NL, DP5, Tsit5, DP8, Vern6, Vern7, Vern8, TanYam7, TsitPap8,
         Vern9, Feagin10, Feagin12, Feagin14, CompositeAlgorithm, Anas5, RKO65, FRK65, PFRK87, RKM

  export SSPRK22, SSPRK33, KYKSSPRK42, SSPRK53, SSPRK53_2N1, SSPRK53_2N2, SSPRK53_H, SSPRK63, SSPRK73, SSPRK83, SSPRK43, SSPRK432,
         SSPRKMSVS32, SSPRKMSVS43, SSPRK932, SSPRK54, SSPRK104

  export ORK256, CarpenterKennedy2N54, SHLDDRK64, HSLDDRK64, DGLDDRK73_C, DGLDDRK84_C, DGLDDRK84_F, NDBLSRK124, NDBLSRK134, NDBLSRK144,
         CFRLDDRK64, TSLDDRK74,SHLDDRK52,SHLDDRK_2N,
         CKLLSRK43_2,CKLLSRK54_3C,CKLLSRK95_4S,CKLLSRK95_4C,CKLLSRK95_4M,
         CKLLSRK54_3C_3R, CKLLSRK54_3M_3R, CKLLSRK54_3N_3R, CKLLSRK85_4C_3R, CKLLSRK85_4M_3R, CKLLSRK85_4P_3R,
         CKLLSRK54_3N_4R, CKLLSRK54_3M_4R, CKLLSRK65_4M_4R, CKLLSRK85_4FM_4R, CKLLSRK75_4M_5R,
         ParsaniKetchesonDeconinck3S32, ParsaniKetchesonDeconinck3S82,
         ParsaniKetchesonDeconinck3S53, ParsaniKetchesonDeconinck3S173,
         ParsaniKetchesonDeconinck3S94, ParsaniKetchesonDeconinck3S184,
         ParsaniKetchesonDeconinck3S105, ParsaniKetchesonDeconinck3S205,
         KYK2014DGSSPRK_3S2

  export RadauIIA3, RadauIIA5

  export ImplicitEuler, ImplicitMidpoint, Trapezoid, TRBDF2, SDIRK2, SDIRK22,
         Kvaerno3, KenCarp3, Cash4, Hairer4, Hairer42, SSPSDIRK2, Kvaerno4,
         Kvaerno5, KenCarp4, KenCarp47, KenCarp5, KenCarp58, ESDIRK54I8L2SA, SFSDIRK4, SFSDIRK5, CFNLIRK3, SFSDIRK6, SFSDIRK7, SFSDIRK8
         Kvaerno5, KenCarp4, KenCarp5, ESDIRK54I8L2SA, SFSDIRK4, SFSDIRK5, CFNLIRK3, SFSDIRK6, SFSDIRK7, SFSDIRK8

  export MagnusMidpoint, LinearExponential, MagnusLeapfrog, LieEuler, CayleyEuler, MagnusGauss4, MagnusNC6, MagnusGL6, MagnusGL8, MagnusNC8, MagnusGL4,
         MagnusAdapt4, RKMK2, RKMK4, LieRK4, CG2, CG3

  export Rosenbrock23, Rosenbrock32, RosShamp4, Veldd4, Velds4, GRK4T, GRK4A,
         Ros4LStab, ROS3P, Rodas3, Rodas4, Rodas42, Rodas4P, Rodas5,
         RosenbrockW6S4OS, ROS34PW1a, ROS34PW1b, ROS34PW2, ROS34PW3

  export LawsonEuler, NorsettEuler, ETD1, ETDRK2, ETDRK3, ETDRK4, HochOst4, Exp4, EPIRK4s3A, EPIRK4s3B,
         EPIRK5s3, EXPRB53s3, EPIRK5P1, EPIRK5P2, ETD2, Exprb32, Exprb43

  export SymplecticEuler, VelocityVerlet, VerletLeapfrog, PseudoVerletLeapfrog,
         McAte2, Ruth3, McAte3, CandyRoz4, McAte4, McAte42, McAte5,
         CalvoSanz4, Yoshida6, KahanLi6, McAte8, KahanLi8, SofSpa10

  export SplitEuler

  export Nystrom4, Nystrom4VelocityIndependent, Nystrom5VelocityIndependent,
         IRKN3, IRKN4, DPRKN6, DPRKN8, DPRKN12, ERKN4, ERKN5

  export ROCK2, ROCK4, RKC, IRKC, ESERK4, ESERK5, SERK2

  export AB3, AB4, AB5, ABM32, ABM43, ABM54

  export VCAB3, VCAB4, VCAB5, VCABM3, VCABM4, VCABM5

  export VCABM

  export IMEXEuler, CNAB2, CNLF2

  export AN5, JVODE, JVODE_Adams, JVODE_BDF

  export ABDF2, QNDF1, QBDF1, QNDF2, QBDF2, QNDF, QBDF

  export SBDF2, SBDF3, SBDF4

  export MEBDF2

  export AutoSwitch, AutoTsit5, AutoDP5,
         AutoVern6, AutoVern7, AutoVern8, AutoVern9

  export AitkenNeville, ExtrapolationMidpointDeuflhard, ExtrapolationMidpointHairerWanner, ImplicitEulerExtrapolation,
         ImplicitDeuflhardExtrapolation, ImplicitHairerWannerExtrapolation, ImplicitEulerBarycentricExtrapolation

  export KuttaPRK2p5, PDIRK44, DImplicitEuler, DABDF2

  export ShampineCollocationInit, BrownFullBasicInit, NoInit
end # module


using MuZero

#=
@boardgame Connect4 (6,7) (21,21)
@actions [1,2,3,4]
@pieces [
    (name=:black, text="B", icon=(symbol="●", color="black")),
    (name=:red, text="R", icon=(symbol="●", color="red"))
]
=#

struct Connect4State <: AbstractBoardState
    board_state::BoardState{42}
end

Connect4State() = Connect4State(BoardState{42}())

#@todo change this to player "A" and "B"
const PLAYER_1_PIECES = [Piece(:black, "B", icon_symbol="●", icon_color="black", icon_size=60)]
const PLAYER_2_PIECES = [Piece(:red, "R", icon_symbol="●", icon_color="red", icon_size=60)]
const ALL_PIECES = [PLAYER_1_PIECES..., PLAYER_2_PIECES...]

MuZero.get_pieces(::Type{Connect4State}) = ALL_PIECES
MuZero.num_rows(::Type{Connect4State}) = 6
MuZero.num_cols(::Type{Connect4State}) = 7
MuZero.image_scale(::Type{Connect4State}) = 750
MuZero.image_title_offset(::Type{Connect4State}) = (75,45)
MuZero.actions(::Type{Connect4State}) = [1,2,3,4,5,6,7]

struct Connect4StateSVG
    s::Connect4State
end

struct Connect4 <: AbstractBoardGame
    history::Vector{<:Integer}
    states::Vector{Connect4State}
    rewards::Vector{Float32}
    function Connect4(history, states, rewards)
        @assert length(history) == length(rewards)
        @assert length(states) == (length(history) + 1)
        new(history, states, rewards)
    end
end

function Connect4()
    Connect4(
        Vector{Int8}(),
        [Connect4State()],
        Vector{Float32}(),
    )
end

Connect4(actions::Vector{<:Integer}) = init_game(Connect4, actions)

MuZero.getstatetype(::Type{Connect4}) = Connect4State

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

function MuZero.get_marker(g::Connect4, player::Integer, action::Integer)
    col_idx = action
    col = g[:, col_idx]
    row_idx = findlast(x->x==0, col)
    color = player == PLAYER_1 ? :black : :red
    piece = get_piece(g, color)
    Marker(row_idx, col_idx, piece)
end

function MuZero.get_winner(s::Connect4State)
    for player in [PLAYER_1, PLAYER_2]
        target = player * ones(UInt8, 4)
        for row_idx in 1:num_rows(s)
            for col_idx in 1:4
                if s[row_idx, col_idx:(col_idx+3)] == target
                    return player
                end
            end
        end
        for col_idx in 1:num_cols(s)
            for row_idx in 1:3
                if s[row_idx:(row_idx+3), col_idx] == target
                    return player
                end
            end
        end
    end
    nothing
end

function MuZero.legal_actions(g::Connect4)
    if game_is_won(g)
        return []
    end
    filter(col_idx->iszero(g[1,col_idx]), 1:num_cols(g))
end

function MuZero.getconfig(::Type{Connect4})
    config = make_board_game_config(
        Connect4,
        network_type=Incremental,
        max_moves=42,
        root_dirichlet_alpha=0.3,
        action_space_size=7, # Number of all possible actions in the game
        support_size = 10, # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))
        observation_shape = (6, 7, 3),
        stacked_observations = 0, # Number of previous observations and previous actions to add to the current observation
        hidden_state_size = 32, # The 1-dimensional size of the encoded game state
        checkpoint_interval = 100, #  Number of training steps before using the model for self-playing
        num_actors = 3,
        weight_decay = 1e-3,
        lr_init=0.002f0,
        save_interval = 10000,
        num_unroll_steps=1,
    )
    config.batch_size = 10
    config.num_simulations = 10
    config.training_steps = 12000
    config
end

# @todo Remove this requirement
function MuZero.gettestfn(::Type{Connect4})
    function fn(n::Network)
        0.5
    end
end

function MuZero.policy_testset(t::Type{Connect4})
    d = Dict()
    d[Connect4([1,7,1,7,1,7])] = [1]
    d
end


using SafeTestsets

@safetestset "Testing Finsts" begin
    using  VisualSearchACTR, Test
    ex = Experiment()
    target,present = initialize_trial!(ex)
    visual_objects = ex.populate_visicon(ex, target..., present)
    T = typeof(visual_objects)
    visual_location = VisualLocation(buffer=T)
    visicon = visual_objects
    visual_location.iconic_memory = visual_objects
    target_chunk = Chunk(;target...)
    goal = Goal(buffer=target_chunk)
    visual = Visual(buffer=T)
    actr = ACTRV(;T=Parm, goal, visual_location, visual, visicon)
    iconic_memory = get_iconic_memory(actr)
    map(x->x.attended=true, iconic_memory)
    map(x->x.attend_time = 0.0, iconic_memory)
    iconic_memory[1].attend_time = 1.1
    iconic_memory[2].attend_time = 1.1
    iconic_memory[3].attend_time = 1.1
    iconic_memory[4].attend_time = 1.1
    iconic_memory[6].attend_time = 2.0
    VisualSearchACTR.update_finst!(actr)
    @test iconic_memory[6].attend_time == 2.0
    @test iconic_memory[6].attended
    @test sum(x->x.attended, iconic_memory) == actr.parms.n_finst
    @test maximum(x->x.attend_time, iconic_memory) <= actr.parms.finst_span
end

@safetestset "Testing Iconic Decay" begin
    using  VisualSearchACTR, Test
    ex = Experiment()
    target,present = initialize_trial!(ex)
    visual_objects = ex.populate_visicon(ex, target..., present)
    T = typeof(visual_objects)
    visual_location = VisualLocation(buffer=T)
    visicon = visual_objects
    visual_location.iconic_memory = visual_objects
    target_chunk = Chunk(;target...)
    goal = Goal(buffer=target_chunk)
    visual = Visual(buffer=T)
    actr = ACTRV(;T=Parm, goal, visual_location, visual, persistence=1.0,
        visicon, time=4.0)
    iconic_memory = get_iconic_memory(actr)
    visible_objects = iconic_memory[1:5]
    map(x->x.features.color.visible=true, visible_objects[1:2])
    map(x->x.features.color.fixation_time=5.0, visible_objects[1:2])
    map(x->x.features.shape.visible=true, visible_objects[3:end])
    map(x->x.features.shape.fixation_time=6.0, visible_objects[3:end])
    map(x->VisualSearchACTR.object_visibility!(x), iconic_memory)
    @test sum(x->x.visible, iconic_memory) == 5
    VisualSearchACTR.update_decay!(actr)
    map(x->VisualSearchACTR.object_visibility!(x), iconic_memory)
    @test sum(x->x.visible, iconic_memory) == 5
    actr.time = 6.5
    VisualSearchACTR.update_decay!(actr)
    map(x->VisualSearchACTR.object_visibility!(x), iconic_memory)
    @test sum(x->x.visible, iconic_memory) == 3
end

@safetestset "Testing Activation" begin
    using  VisualSearchACTR, Test

    vals = [(color = :gray, shape = :p, x = 338, y =515),
            (color = :black, shape = :q, x = 351 , y = 193),
            (color = :gray, shape = :q, x = 511 , y = 462),
            (color = :black, shape = :p, x = 511 , y = 400)]

    experiment = Experiment(n_trials=10^1, trace=true, visible=false)

    target = (color=:gray,shape=:p)
    visual_objects = [VisualObject(features=(color=Feature(;value=x.color), shape=Feature(;value=x.shape)),
        width=32.0, location=[x.x-0.,x.y-0.]) for x in vals]

    T = typeof(visual_objects)
    visual_location = VisualLocation(buffer=T)
    visicon = visual_objects
    visual_location.iconic_memory = visual_objects
    target_chunk = Chunk(;target...)
    goal = Goal(buffer=target_chunk)
    visual = Visual(buffer=T)
    actr = ACTRV(;T=Parm, goal, visual_location, visual, visicon, σ=0.0)
    iconic_memory = get_iconic_memory(actr)

    # model = Model(;target=target, iconic_memory=visicon, noise=0.0)
    VisualSearchACTR.compute_angular_size!(actr)
    actr.visual.focus = [324.0,324.0]
    VisualSearchACTR.update_decay!(actr)
    VisualSearchACTR.update_finst!(actr)
    VisualSearchACTR.update_visibility!(actr)
    VisualSearchACTR.compute_activations!(actr)

    correct_top_down = [2,0,1,1]
    @test all(x->x[2].topdown_activation == x[1], zip(correct_top_down, iconic_memory))

    vo1 = iconic_memory[1]
    vo2 = iconic_memory[2]
    distance = VisualSearchACTR.compute_distance(vo1, vo2)
    bottomup_activation1 = VisualSearchACTR.bottomup_activation(vo1.features, vo2.features, distance)
    @test bottomup_activation1 == 1/sqrt(distance)*2

    vo1 = iconic_memory[1]
    vo3 = iconic_memory[3]
    distance = VisualSearchACTR.compute_distance(vo1, vo3)
    bottomup_activation2 = VisualSearchACTR.bottomup_activation(vo1.features, vo3.features, distance)
    @test bottomup_activation2 == 1/sqrt(distance)

    vo1 = iconic_memory[1]
    vo4 = iconic_memory[4]
    distance = VisualSearchACTR.compute_distance(vo1, vo4)
    bottomup_activation3 = VisualSearchACTR.bottomup_activation(vo1.features, vo4.features, distance)
    @test bottomup_activation3 == 1/sqrt(distance)

    @test vo1.bottomup_activation == (bottomup_activation1 + bottomup_activation2 + bottomup_activation3)
end

@safetestset "Testing Feature Search" begin
    using VisualSearchACTR, Test, DataFrames, GLM, Statistics
    include("simulation.jl")
    Random.seed!(95025181)
    set_sizes = [1,2,5,10,20,30]
    Δτ = .8
    topdown_weight = .60
    noise = true
    rnd_time = true
    results = run_simulation(set_sizes; fun=feature_set, Δτ, topdown_weight, noise, rnd_time)

    df_present = filter(x->x[:target_present] ==:present && x[:response] ==:present, results)
    ols_present = lm(@eval(@formula(rt_mean ~ distractors)), df_present)
    β0,β1 = coef(ols_present)
    @test β0 ≈ 0.36 rtol = .05
    @test β1 ≈ 0.0 atol = .01
    df_absent = filter(x->x[:target_present] ==:absent && x[:response] ==:absent, results)
    ols_absent = lm(@eval(@formula(rt_mean ~ distractors)), df_absent)
    β0,β1 = coef(ols_absent)
    @test β0 ≈ 0.45 rtol = .05
    @test β1 ≈ 0.0 atol = .01
end

@safetestset "Testing Conjunctive Search" begin
    using VisualSearchACTR, Test, DataFrames, GLM, Statistics
    include("simulation.jl")
    Random.seed!(52484)
    set_sizes = [1,2,5,10,20,30]
    Δτ = .4
    topdown_weight = .60
    noise = true
    rnd_time = true
    results = run_simulation(set_sizes; fun=conjunctive_set, Δτ, topdown_weight, noise, rnd_time)

    df_present = filter(x->x[:target_present] ==:present && x[:response] ==:present, results)
    ols_present = lm(@eval(@formula(rt_mean ~ distractors)), df_present)
    β0,β1 = coef(ols_present)
    @test β0 ≈ 0.34 rtol = .05
    @test β1 ≈ 0.020 atol = .01
    df_absent = filter(x->x[:target_present] ==:absent && x[:response] ==:absent, results)
    ols_absent = lm(@eval(@formula(rt_mean ~ distractors)), df_absent)
    β0,β1 = coef(ols_absent)
    @test β0 ≈ 0.38 rtol = .05
    @test β1 ≈ 0.067 atol = .01
end

@safetestset "Gamma Parms" begin
    using VisualSearchACTR, Test, Distributions
    import VisualSearchACTR: gamma_parms

    θ = gamma_parms(3, 2)
    @test mean(Gamma(θ...)) ≈ 3 rtol = .0005
    @test std(Gamma(θ...)) ≈ 2 rtol = .0005

    θ = gamma_parms(.2, 1)
    @test mean(Gamma(θ...)) ≈ .2 rtol = .0005
    @test std(Gamma(θ...)) ≈ 1 rtol = .0005
end

@safetestset "Run model" begin
    using VisualSearchACTR, Test
    experiment = Experiment(set_size=10,  n_trials=2,
        trace=true, speed =.5)
    run_condition!(experiment)
    @test true
end


#Defines the control system type, including its discrete time behaviour
#include("DiscreteTimeSystem.jl")
#= ControlFunction should input the "measured state" of the vehicle obtained from the
    Navigation system, the attitude/rate commands from the guidance system, and time.
=#

#Define ControlSystem as a subtype of the DiscreteTimeSystem
abstract type AbstractControlSystem <: DiscreteTimeSystem end

########### Template function ###########################
mutable struct ControlSystem{T <: Number} <: AbstractControlSystem
    Active::Bool #Activation Boolean
    Ts::T #Sample Time
    ControlRelation::Function #Encodes
    Tprev::T # Storage for previous activation values
end
ControlSystem() = ControlSystem(false, 0.0, x -> nothing, 0.0)
ControlSystem(Ts::Number) = ControlSystem(false, Float64(Ts), x -> nothing, 0.0)
ControlSystem(Ts::Number, csfunc::Function) = ControlSystem(true, Float64(Ts), csfunc, 0.0)
ControlSystem(Active::Bool, Ts::Number, csfunc::Function) = ControlSystem(Active, Float64(Ts), csfunc, 0.0)

#Set the callback affect condition of the ControlSystem
#placeholder!
function affect!(int,dts::Type{<:AbstractControlSystem})
    for cu in user_cache(int)
        cu.u = [dts.ControlRelation(int.u)...]
    end
  dts.Tprev = int.t
end

#=cb = ContinuousCallback((t, u, int) -> condition(det,t,u,int),
                          (int) -> affect!(int,det))
=#


# some definitions used throughout TIcon

"""
  Abstract type for _AssembleElementData.  See that type for details.
"""
abstract type AssembleElementData end

#------------------------------------------------------------------------------
# sparsity pattern enums

"""
  disc_type constant for SparseMatrixCSC constructor, indicating the face nodes
  of an element are connected to the face nodes of the element that shares
  the face.
"""
global const INVISCID=1

"""
  disc_type constant for SparseMatrixCSC constructor, indicating all nodes
  of an element are connected to all face nodes of its neighbours
"""
global const VISCOUS=2

"""
  disc_type constant for SparseMatrixCSC constructor, indicating all nodes
  of an element are connected to all nodes of its neighbors
"""
global const COLORING=3

"""
  disc_type constant for sparsity pattern of viscous terms when the T4 = 0.
  In this case, all nodes in the stencil of the interpolation operator are
  connected to all nodes of the adjacent element.
"""
global const VISCOUSTIGHT=4

#------------------------------------------------------------------------------
# Shock Capturing stuff - see the docs in the Euler module

abstract type AbstractShockSensor end


@testset "Cross validating LMHMM on ZX..." begin
        @time ctb = ZhuXian()

        nb = 10
        batches = KongYiji.foldbatch(ctb, nb)
        tbs = KongYiji.HmmScoreTable[]
        @showprogress 1 "Cross Validating ooHMM..." for (tr, te) in batches
                hmm = KongYiji.HMM(tr, "", 0.)
                lm = KongYiji.LM(tr; epochs=20, rnnType=:lstm, numLayers=1, bidirectional=false)
                m = hmm + lm
                x = KongYiji.rawsents(te)
                z = KongYiji.wordsents(te)
                @time y = m(x; beams=500, recover=false, withpos=false)
                push!(tbs, KongYiji.HmmScoreTable(z, y))
                println(tbs[end])
        end
        hmmscoretable = sum(tbs)
        println(hmmscoretable)
end


########################################################################
# Training
########################################################################

# Implementation of the FedAvg and FedProx algorithm
function fedAvgAndProx(
    server::Union{FedProxServer, FedProxServerNN},
    clients::Union{Vector{FedProxClient}, Vector{FedProxClientNN} },
    numRounds::Int64
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        select!(server)
        sendModel!(server)
        # objValue = getObjValue(server)
        # @printf("Round : %4d, obj: %6.4e\n", t, objValue)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            update!(client)
        end
        @printf("aggregate\n")
        aggregate!(server)
        # Print log
        sendModelToAllClients!(server)
        @printf("getObj\n")
        objValue = getObjValue(server)
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end

# Implementation of the Scaffold algorithm.
function Scaffold(
    server::Union{ScaffoldServer, ScaffoldServerNN},
    clients::Union{Vector{ScaffoldClient}, Vector{ScaffoldClientNN} },
    numRounds::Int64
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        select!(server)
        sendModel!(server)
        # objValue = getObjValue(server)
        # @printf("Round : %4d, obj: %6.4e\n", t, objValue)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            update!(client)
        end
        aggregate!(server)
        # Print log
        sendModelToAllClients!(server)
        objValue = getObjValue(server)
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end

# Implementation of the FedDCD algorithm (both exact and inexact)
function fedDCD(
    server::FedDCDServer,
    clients::Vector{FedDCDClient},
    numRounds::Int64
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        # @printf("Round %d\n", t)
        select!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            update!(client)
        end
        aggregate!(server)
        sendModel!(server)
        # Print log
        objValue = getObjValue(server)
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end

function fedDCD(
    server::FedDCDServerNN,
    clients::Vector{FedDCDClientNN},
    numRounds::Int64
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        # @printf("Round %d\n", t)
        select!(server)
        synchronize!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            update!(client)
        end
        aggregate!(server)
        sendModel!(server)
        decay_lr!(server)
        # Print log
        objValue = getObjValue(server)
        # acc = 0.0
        # num_data = 0
        # for c in server.clients
        #     n = size(c.Xtrain, 1)
        #     num_data += n
        #     acc += n*accuracy(c.XtrainT, c.Ytrain, server.W)
        # end
        # acc /= num_data
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end

# Implementation of the accelerated FedDCD algorithm
function accfedDCD(
    server::AccFedDCDServer,
    clients::Vector{AccFedDCDClient},
    numRounds::Int64;
    objMin::Float64 = 0.0
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        # @printf("Round %d\n", t)
        # update v
        Threads.@threads for client in server.clients
            updatev!(client)
        end
        # first inner round
        select!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            updateW!(client)
        end
        aggregate!(server)
        sendModel!(server, 1)
        # update u
        Threads.@threads for client in server.clients
            updateu!(client)
        end
        # second inner round
        select!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            updateW!(client)
        end
        aggregate!(server)
        sendModel!(server, 2)
        # Print log
        objValue = getObjValue(server)
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if objValue < objMin
            break
        end
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end

function accfedDCD(
    server::AccFedDCDServerNN,
    clients::Vector{AccFedDCDClientNN},
    numRounds::Int64;
    objMin::Float64 = 0.0
)
    # Connect clients with server
    server.clients = clients
    # Training process
    objList = zeros(Float64, 0)
    testAccList = zeros(Float64, 0)
    startTime = time()
    @printf("Start training!\n")
    for t = 1:numRounds
        # update v
        Threads.@threads for client in server.clients
            updatev!(client)
        end
        # first inner round
        select!(server)
        synchronize!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            updateW!(client)
        end
        aggregate!(server)
        sendModel!(server, 1)
        # update u
        Threads.@threads for client in server.clients
            updateu!(client)
        end
        # second inner round
        select!(server)
        synchronize!(server)
        Threads.@threads for idx in server.selectedIndices
            client = server.clients[idx]
            updateW!(client)
        end
        aggregate!(server)
        sendModel!(server, 2)
        # Print log
        objValue = getObjValue(server)
        acc = accuracy(server.Xtest, server.Ytest, server.W)
        @printf("Round : %4d, obj: %6.4e, acc: % 3.2f %%, time: %4.3f s\n", t, objValue, acc*100, time()-startTime)
        push!(objList, objValue)
        push!(testAccList, acc)
        if objValue < objMin
            break
        end
        if isnan(objValue)
            break
        end
    end
    endTime = time()
    @printf("Finished training, time elapsed: %.4e\n", endTime - startTime)
    return server.W, objList, testAccList
end


# ----------------------------------------------------------------------------------- #
# Copyright (c) 2016 Varnerlab
# School of Chemical Engineering Purdue University
# W. Lafayette IN 46907 USA

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------------- #
function DataFile(TSTART,TSTOP,Ts)
# ----------------------------------------------------------------------------------- #
# DataFile.jl was generated using the Kwatee code generation system.
# DataFile: Stores model parameters as key - value pairs in a Julia Dict()
# Username: jeffreyvarner
# Type: GRN-JULIA
# Version: 1.0
# Generation timestamp: 02-26-2016 14:31:22
#
# Input arguments:
# TSTART  - Time start
# TSTOP  - Time stop
# Ts - Time step
#
# Return arguments:
# data_dictionary  - Data dictionary instance (holds model parameters)
# ----------------------------------------------------------------------------------- #

# Initial conditions (default, no initiator)
initial_condition_array = zeros(19)
initial_condition_array[1] = 0.0      # 1 zymosan
initial_condition_array[2] = 1.9      # 2 C4
initial_condition_array[3] = 0.322    # 3 C2
initial_condition_array[8] = 7.57     # 4 C3
initial_condition_array[14] = 0.195   # 5 C5
initial_condition_array[17] = 2.23    # 6 Factor H
initial_condition_array[18] = 0.417   # 7 C4BP

# Load the parameters from disk -
# parameter_array = readdlm("./best_solution.txt")

# JV: initialize a placeholder parameter array.
# The actual parameter ensemble is contained in the data subdir
# pc_array*.dat = number_of_parameters x number_of_trials (each col is a parameter soln)
# rank_array*.dat = number_of_trials x 1  (holds the Pareto rank for each col in pc_array*.dat)
parameter_array = zeros(28,1)

# ---------------------------- DO NOT EDIT BELOW THIS LINE -------------------------- #
data_dictionary = Dict();
data_dictionary["PARAMETER_ARRAY"] = parameter_array;
data_dictionary["INITIAL_CONDITION_ARRAY"] = initial_condition_array;
# ----------------------------------------------------------------------------------- #
return data_dictionary;
end


#Input.jl
"""
Handling input for DEB-IBM simulations.
Contains functions for parameter conversions and corrections, as well as dictionaries with AMP values and default simulation settings.
"""

using DocStringExtensions

#### ----- default global dictionaries ---- ####

# parameters for population experiment with D. magna conducted in early 2019
population_dict = Dict(
    "remove_juveniles" => [], # weekdays on which to remove juveniles
    "renew_medium" => [0, 3], # weekdays on which to remove medium
    "frct_medium_renewal" => 0.25, # fraction of medium to replace per renewal
    "N0_juvenile" => 8, # initial juveniles
    "N0_adult"    => 2, # initial adults
    "timestep"    => 24, # number of timesteps per day
    "tmax"        => 42, # max. simulated time (days)
    "food_level"  => 8e7, # food level (# cells/L/day)
    "age_range_init" => 24, # maximum initial age of juveniles (timesteps)
    "stressor_conc" => [0., 0., 0.]
)

# settings of a life-table experiment as conducted in late 2019/early 2020
life_table_dict = Dict(
    "remove_juveniles" => [0, 2, 5], # remove juveniles 3x per week
    "renew_medium"     => [0, 2, 5], # renew media on the same days
    "frct_medium_renewal" => 1.00, # 100% of media are renewed
    "N0_juvenile" => 1, # simulation starts with 1 juvenila
    "N0_adult"    => 0,
    "timestep"    => 24, # number of timesteps per day
    "tmax"        => 21, # test lasts for 21 days
    "food_level"  => 5., # food level in mgC//L3/day
    "age_range_init" => 24, # maximum initial age of juveniles (timesteps)
    "stressor_conc" => [0., 0., 0.]
)

# reference scenario for ecorelevance project
ecorelevance_reference_dict = Dict(
    "remove_juveniles" => [],
    "renew_medium" => [0,3], # renew media on mondays and thursdays
    "frct_medium_renewal" => 0.25, # renew 25% media
    "N0_juvenile" => 3, # start with 3 juveniles
    "N0_adult" => 1, # and one adult
    "age_range_init" => 48,
    "timestep" => 24,
    "tmax" => 56,
    "food_level" => 2.5/1000, # food level in mgC/cm^3/day
    "stressor_1_conc" => 0, # concentration of stressor 1
    "stressor_2_conc" => 0, # concentration of stressor 2
    "stressor_3_conc" => 0  # concentration of stressor 3
)

"""
Conversion of AMP parameters for use in IBM.
See user manual of the generic DEB-IBM implementation for explanations.
$(SIGNATURES)
"""
function convert_parameters(fld)
    # maintenance rate follows from volume-specific somatic structure
    # and volume-specific costs for structure
    fld["k_M_rate"]    = fld["p_M"]  / fld["E_G"]
     # compound parameters following from scaling by assimilation rate
    fld["g"]           = fld["E_G"]  * fld["v_rate"] / (fld["kap"] * fld["p_Am"])
    fld["U_Hb"]        = fld["E_Hb"] / fld["p_Am"]
    fld["U_Hp"]        = fld["E_Hp"] / fld["p_Am"]
    # functional response parameters derived from primary parameters
    fld["J_XAm_rate"]  = fld["p_Am"] / fld["y_EX"]
    fld["K"]           = fld["J_XAm_rate"] / fld["F_m"]
    return fld
end

"""
Estimate embryonal parameters using bisection method.
"""
function estimate_egg_weight(
        timestep::Int64,
        v_rate::Float64,
        U_Hb::Float64,
        kap::Float64,
        g::Float64,
        k_J_rate::Float64,
        k_M_rate::Float64
)
    embryo_timestep = timestep * 1000
    lower_bound = 0
    upper_bound = 1
    maxiter = 100
    egg_weight = 0.5
    U_H_embryo = 0
    U_E_embryo = 0
    e_scaled_embryo = 0
    e_ref = 1
    L_embryo = 0.00001

    for i in 1:maxiter
        egg_weight = .5 * (lower_bound + upper_bound)
        L_embryo = 0.00001
        U_E_embryo = egg_weight
        U_H_embryo = 0
        e_scaled_embryo = v_rate * (U_E_embryo/L_embryo)^3

        while (U_H_embryo < U_Hb) & (e_scaled_embryo >1)
             # calculate reserve density based on energy conductance, scaled reserves and structural volume
            e_scaled_embryo = v_rate * (U_E_embryo/L_embryo^3)
            # calculate corresponding mobilization flux
            S_C_embryo      = L_embryo^2 * (g * e_scaled_embryo / ( g + e_scaled_embryo)) * (1 + (L_embryo / (g * (v_rate / (g * k_M_rate)))))
            # calculate change in state varaibles
            dU_E_embryo     = -1 * S_C_embryo
            dU_H_embryo     = (1-kap)*S_C_embryo - k_J_rate*U_H_embryo
            dL_embryo       = ((1/3) * (((v_rate / (g*(L_embryo^2))) * S_C_embryo) - k_M_rate*L_embryo))
            # update state variables
            U_E_embryo = U_E_embryo + dU_E_embryo / embryo_timestep
            U_H_embryo = U_H_embryo + dU_H_embryo / embryo_timestep
            L_embryo   = L_embryo   + dL_embryo   / embryo_timestep
            e_scaled_embryo = v_rate * (U_E_embryo / L_embryo^3)
        end

        # in order to accept an estimate:
        # 1. estimated reserve density of embryo has to be close to one
        # (to fulfil the criterion that embryos are mostly reserve)
        # 2. estimated maturity has to be larger or equal to maturity at birth (by definition)
        if (e_scaled_embryo < 0.01+e_ref) & (e_scaled_embryo > -0.01+e_ref) & (U_H_embryo >= U_Hb)
            @goto escape_label # jump to @label escape_label
        end
        # if estimate has not been accpeted:
        # adjust upper/lower bound based on result of previous estimation
        if U_H_embryo > U_Hb
            upper_bound = egg_weight
        else
            lower_bound = egg_weight
        end
        #if i == maxiter
        #    error("Embryo submodel is not converging. Smaller timesteps could help.")
        #end
    end
    @label escape_label

    return egg_weight, L_embryo, U_E_embryo, U_H_embryo
end

"""
Estimation of intitial state variables of an Embryo, determining the amount of energy
that has to be invested per offspring.
$(SIGNATURES)
"""
function estimate_egg_weight(fleas_dict::Dict)
    # estimate embryonal parameters
    fleas_dict["egg_weight"],
    fleas_dict["L_embryo"],
    fleas_dict["U_E_embryo"],
    fleas_dict["U_H_embryo"] = estimate_egg_weight(
        24,
        fleas_dict["v_rate"],
        fleas_dict["U_Hb"],
        fleas_dict["kap"],
        fleas_dict["g"],
        fleas_dict["k_J_rate"],
        fleas_dict["k_M_rate"])
    return fleas_dict
end

#### ---- default flea dictionaries ---- ####
"""
Parameter values for Daphnia magna according to AMP database.
Last checked: Jan 2020
Toxicity parameters from Daphnia magna models (Pereira et al. 2019; Karel Vlaeminck, unpublished).
"""
DM_AMP = Dict(
    "species" => "DM_AMP",
    "cv"      => 0.1,

    # AMP parameters
    "p_Am" => 313.169, #
    "F_m"  => 30., # filtration rate (cm^3/day*cm^2)
    "E_Hb"       => 0.05464,
    "E_Hp"       => 1.09,
    "p_M"        => 1200,
    "E_G"        => 4400,
    "k_J_rate"   => 0.2537,
    "v_rate"     => 0.1858,
    "kap"        => 0.5809,
    "kap_R"      => 0.95,
    "shape_factor"       => 0.264,

    "sG"  => -0.3,
    "h_a" => 0.0002794,
    "time_between_molts" => 2.324,

    # embryonal parameters
    "L0"  => 0.0001,
    "egg_weight" => 3.771781921386719E-4,
    "L_embryo"   => 0.023340451224020977,
    "U_E_embryo" => 6.039673323323045E-5,

    "starvation_hazard_rate" => 0.,
    "crit_mass" => 0.,

    # Toxicity parameters
    # Toxicity parameters
    "pmoas"  => ["Growth", "Maintenance", "Assimilation", "Reproduction"],
    "k_d"    => [ 0., 0., 0.],
    "h_max"  => [ 0., 0., 0.],
    "h_beta" => [ 0., 0., 0.],
    "h_ED50" => [ 0., 0., 0.],
    "S_max"  => [ 0., 0., 0.],
    "S_beta" => [ 0., 0., 0.],
    "S_ED50" => [ 0., 0., 0.]
)

"""
Parameter values for Daphnia magna, fed with mixture of Pseudokirchneriella subcapitata and Chlamydomonas reinhardtii, simulated as a single food source.
Food densities in # cells/L.
"""
DM1_2019 = Dict(
    "species" => "DM1_2019",
    "cv"      => 0.1,

    # corrected parameters

    "p_Am" => 260.4, # (CI = (260.3, 309.5))
    "F_m"  => 10.0, # (CI = (8, 11))
    "y_EX" => 3.7e-6, # (CI = (3e-6, 4e-6))
    "mortality_constant" => 0.41, # (CI = (0.3, 0.63))

    # AMP parameters

    "E_Hb"       => 0.05464,
    "E_Hp"       => 1.09,
    "p_M"        => 1200,
    "E_G"        => 4400,
    "k_J_rate"   => 0.2537,
    "v_rate"     => 0.1858,
    "kap"        => 0.5809,
    "kap_R"      => 0.95,
    "shape_factor"       => 0.264,

    "sG"  => -0.3,
    "h_a" => 0.0002794,
    "time_between_molts" => 2.324,

    # embryonal parameters
    "L0"  => 0.0001,
    "egg_weight" => 3.771781921386719E-4,
    "L_embryo"   => 0.023340451224020977,
    "U_E_embryo" => 6.039673323323045E-5,

    "starvation_hazard_rate" => 0.,
    "crit_mass" => 0.,

    "toxicity_particle" => Dict(
        "stressors" => ["Cu", "Ni", "Zn"],
        "pMoAs" => ["Growth", "Maintenance", "Assimilation", "Reproduction"],
        # one k_d for Cu, Ni and Zn
        "k_d" => [0., 0., 0.],
        # one parameter for every stressor (rows) and every pMoA (columns)
        "S_max" => hcat(
            [0., 0., 0., 0.],
            [0., 0., 0., 0.],
            [0., 0., 0., 0.]
        ),
        "S_ED50" => fill(Inf, 4, 3),
        "S_beta" => fill(2., 4, 3),
        "h_max" => [0., 0., 0.],
        "h_ED50" => [Inf, Inf, Inf],
        "h_beta" => [2., 2., 2.]
    )
)

DL_AMP = Dict(
    "species" => "DL_AMP",
    "cv"      => 0.1,

    # AMP parameters
    "E_Hb"       => 0.01746,
    "E_Hp"       => 1.401,
    "p_M"        => 5686.83,
    "E_G"        => 4400.59,
    "k_J_rate"   => 0.002,
    "v_rate"     => 0.020455,
    "kap"        => 0.5568,
    "kap_R"      => 0.95,
    "shape_factor"       => 0.18946,
    "mortality_constant" => 0,
    "crit_mass"          => 0.4,
    "starvation_hazard_rate" => 0.9,
    "sG"  => 0.001,
    "h_a" => 8.603e-5,
    "L0"  => 0.0001,

    # food-dependent parameters
    "p_Am" => 558.735, # J/d*cm^2
    "y_EX" => 1e-3, # yield of reserve (J) on food (# cells/cm^3)
    "F_m"  => 270.833, # filtration rate (cm^3/day*cm^2)

    # embryonal parameters
    "egg_weight" => 3.771781921386719E-4,
    "L_embryo"   => 0.023340451224020977,
    "U_E_embryo" => 6.039673323323045E-5,

    # reproductive timing
    "time_between_molts" => 2.324,

    # Toxicity parameters
    # compound 1
    "compound_1" => "Copper",
    "pmoa_1"     => "Growth",
    "h_b_1"      => 0/24, # background hazard rate
    "k_d_1"      => 0, # dominant rate constant
    "k_k_1"      => 0, # ?
    "m_n_1"      => 0, # ?
    "s_s_1"      => 0,
    "s_d_1"      => 0, # ?
    "A_eq_1" => 0,
    "B_eq_1" => 0,
    "C_eq_1" => 0,

    # compound 2
    "compound_2" => "Nickel",
    "pmoa_2"     => "Maintenance",
    "h_b_2"      => 0,
    "k_d_2"      => 0,
    "k_k_2"      => 0,
    "m_n_2"      => 0,
    "s_s_2"      => 0,
    "s_d_2"      => 0,
    "A_eq_2"     => 0,
    "B_eq_2"     => 0,
    "C_eq_2"     => 0,

    # compound 3
    "compound_3" => "Zinc",
    "pmoa_3"     => "Assimilation",
    "h_b_3"      => 0,
    "k_d_3"      => 0,
    "k_k_3"      => 0,
    "m_n_3"      => 0,
    "s_s_3"      => 0,
    "s_d_3"      => 0,
    "A_eq_3"     => 0,
    "B_eq_3"     => 0,
    "C_eq_3"     => 0
)

DL1 = Dict(
    "species" => "DL_AMP",
    "cv"      => 0.1,

    # AMP parameters
    "E_Hb"       => 0.01746,
    "E_Hp"       => 1.401,
    "p_M"        => 5686.83,
    "E_G"        => 4400.59,
    "k_J_rate"   => 0.002,
    "v_rate"     => 0.020455,
    "kap"        => 0.331941,
    "kap_R"      => 0.95,
    "shape_factor"       => 0.18946,
    "mortality_constant" => 0.93567,
    "sG"  => 0.001,
    "h_a" => 8.603e-5,
    "L0"  => 0.0001,

    # food-dependent parameters
    "p_Am" => 651.572, # 2020-03-24 J/d*cm^2
    "y_EX" => 71312.7, # yield of reserve (J) on food (# cells/cm^3)
    "F_m"  => 606.58, # filtration rate (cm^3/day*cm^2)

    # embryonal parameters
    "egg_weight" => 3.771781921386719E-4,
    "L_embryo"   => 0.023340451224020977,
    "U_E_embryo" => 6.039673323323045E-5,

    # other
    "time_between_molts" => 2.324,
    "crit_mass"          => 0.442002,
    "starvation_hazard_rate" => 0.784624,

    # Toxicity parameters
    # compound 1
    "compound_1" => "Copper",
    "pmoa_1"     => "Growth",
    "h_b_1"      => 0/24, # background hazard rate
    "k_d_1"      => 0, # dominant rate constant
    "k_k_1"      => 0, # ?
    "m_n_1"      => 0, # ?
    "s_s_1"      => 0,
    "s_d_1"      => 0, # ?
    "A_eq_1" => 0,
    "B_eq_1" => 0,
    "C_eq_1" => 0,

    # compound 2
    "compound_2" => "Nickel",
    "pmoa_2"     => "Maintenance",
    "h_b_2"      => 0,
    "k_d_2"      => 0,
    "k_k_2"      => 0,
    "m_n_2"      => 0,
    "s_s_2"      => 0,
    "s_d_2"      => 0,
    "A_eq_2"     => 0,
    "B_eq_2"     => 0,
    "C_eq_2"     => 0,

    # compound 3
    "compound_3" => "Zinc",
    "pmoa_3"     => "Assimilation",
    "h_b_3"      => 0,
    "k_d_3"      => 0,
    "k_k_3"      => 0,
    "m_n_3"      => 0,
    "s_s_3"      => 0,
    "s_d_3"      => 0,
    "A_eq_3"     => 0,
    "B_eq_3"     => 0,
    "C_eq_3"     => 0
)

DL1_ECOREL = Dict(
    "species" => "DL_AMP",
    "cv"      => 0.1,

    # AMP parameters
    "E_Hb"       => 0.01746,
    "E_Hp"       => 1.401,
    "p_M"        => 5686.83,
    "E_G"        => 4400.59,
    "k_J_rate"   => 0.002,
    "v_rate"     => 0.020455,
    "kap"        => 0.331941,
    "kap_R"      => 0.95,
    "shape_factor"       => 0.18946,
    "sG"  => 0.001,
    "h_a" => 8.603e-5,
    "L0"  => 0.0001,

    # food-dependent parameters
    "p_Am" => 651.572, # 2020-03-24 J/d*cm^2
    "y_EX" => 71312.7, # yield of reserve (J) on food (# cells/cm^3)
    "F_m"  => 606.58, # filtration rate (cm^3/day*cm^2)

    # embryonal parameters
    "egg_weight" => 3.771781921386719E-4,
    "L_embryo"   => 0.023340451224020977,
    "U_E_embryo" => 6.039673323323045E-5,

    # reproductive timing
    "time_between_molts" => 2.324,

    # starvation
    "mortality_constant" => 0.93567,
    "starvation_hazard_rate" => 0.784624,
    "crit_mass" => 0.442002,

    # Toxicity parameters
    # compound 1
    "compound_1" => "Copper",
    "pmoa_1"     => "Growth",
    "h_b_1"      => 0/24, # background hazard rate
    "k_d_1"      => 0, # dominant rate constant
    "k_k_1"      => 0, # ?
    "m_n_1"      => 0, # ?
    "s_s_1"      => 0,
    "s_d_1"      => 0, # ?
    "A_eq_1" => 0,
    "B_eq_1" => 0,
    "C_eq_1" => 0,

    # compound 2
    "compound_2" => "Nickel",
    "pmoa_2"     => "Maintenance",
    "h_b_2"      => 0,
    "k_d_2"      => 0,
    "k_k_2"      => 0,
    "m_n_2"      => 0,
    "s_s_2"      => 0,
    "s_d_2"      => 0,
    "A_eq_2"     => 0,
    "B_eq_2"     => 0,
    "C_eq_2"     => 0,

    # compound 3
    # Zinc toxicity parameters based on mean µg dissolved Zn/L
    "compound_3" => "Zinc",
    "pmoa_3"     => "Maintenance",
    "k_d_3"      => 0.727032,
    "m_n_3"      => 344.759,
    "k_k_3"      => 2.67,
    "s_s_3"      => 0.01,
    "A_eq_3"     => 298.85,
    "B_eq_3"     => 2.11,
    "C_eq_3"     => 9.40,

    # TODO: remove these parameters entirely from all code
    "h_b_3" => 0,
    "s_d_3" => 0
)


using SnpArrays
using Knockoffs
using DelimitedFiles
using Random
using LinearAlgebra
using Distributions
using ProgressMeter
using MendelIHT
using VCFTools
using StatsBase
using CSV
using DataFrames
using Printf
using GLMNet
BLAS.set_num_threads(1)

global qctools_exe = "/scratch/users/bbchu/qctool/build/release/qctool_v2.0.7"
global snpknock2_exe = "/scratch/users/bbchu/knockoffgwas/snpknock2/bin/snpknock2"
global rapid_exe = "/scratch/users/bbchu/RaPID/RaPID_v.1.7"
global partition_exe = "/scratch/users/bbchu/knockoffgwas/knockoffgwas/utils/partition.R";

"""
    simulate_pop_structure(n, p)

Simulate genotypes with K = 2 populations. 1% of SNPs will have different allele
frequencies between the populations.

# Inputs
- `plinkfile`: Output plink file name.
- `n`: Number of samples
- `p`: Number of SNPs

# Output
- `x1`: n×p matrix of the 1st haplotype for each sample. Each row is a haplotype
- `x2`: n×p matrix of the 2nd haplotype for each sample. `x = x1 + x2`
- `populations`: Vector of length `n` indicating population membership for eachsample.
- `diff_markers`: Indices of the differentially expressed alleles.

# Reference
https://www.nature.com/articles/nrg2813
"""
function simulate_pop_structure(n::Int, p::Int)
    # first simulate genotypes treating all samples equally
    x1 = BitMatrix(undef, n, p)
    x2 = BitMatrix(undef, n, p)
    pmeter = Progress(p, 0.1, "Simulating genotypes...")
    @inbounds for j in 1:p
        d = Bernoulli(rand())
        for i in 1:n
            x1[i, j] = rand(d)
            x2[i, j] = rand(d)
        end
        next!(pmeter)
    end
    # assign populations and simulate 0.01p unually differentiated markers
    populations = rand(1:2, n)
    diff_markers = sample(1:p, Int(0.01p), replace=false)
    @inbounds for j in diff_markers
        pop1_allele_freq = 0.4rand()
        pop2_allele_freq = pop1_allele_freq + 0.6
        pop1_dist = Bernoulli(pop1_allele_freq)
        pop2_dist = Bernoulli(pop2_allele_freq)
        for i in 1:n
            d = isone(populations[i]) ? pop1_dist : pop2_dist
            x1[i, j] = rand(d)
            x2[i, j] = rand(d)
        end
    end
    return x1, x2, populations, diff_markers
end

"""
    simulate_IBD(h1::BitMatrix, h2::BitMatrix, populations::Vector{Int}, k::Int)

Simulate recombination events. Parent haplotypes `h1` and `h2` will be used to generate
`k` children, then both parent and children haplotypes will be returned.

In offspring simulation, we first randomly sample 2 parents from the same population.
Then generate offspring individuals by copying segments of the parents haplotype
directly to the offspring to represent IBD segments. The number of segments (i.e. places of
recombination) is 1 to 5 per sample chosen uniformly across all SNPs.

# Inputs
- `h1`: `n × p` matrix of the 1st haplotype for each parent. Each row is a haplotype
- `h2`: `n × p` matrix of the 2nd haplotype for each parent. `H = h1 + h2`
- `populations`: `populations[i]` is the population (represented as integer) of sample `i`.
- `k`: Total number of offsprings

# Output
- `H1`: `n+k × p` matrix of the 1st haplotype. The first `n` haplotypes are from parents
    and the next `k` haplotypes are the offsprings. Each row is a haplotype
- `H2`: `n+k × p` matrix of the 2nd haplotype. `x = x1 + x2`

# References
https://journals.plos.org/plosgenetics/article?id=10.1371/journal.pgen.1003520
"""
function simulate_IBD(h1::BitMatrix, h2::BitMatrix, populations::Vector{Int}, k::Int)
    n, p = size(h1)
    unique_populations = unique(populations)
    # randomly designate gender for parents
    sex = bitrand(n)
    male_idx = findall(x -> x == true, sex)
    female_idx = findall(x -> x == false, sex)
    # simulate new samples
    x1 = falses(k, p)
    x2 = falses(k, p)
    fathers = Int[]
    mothers = Int[]
    pmeter = Progress(k, 0.1, "Simulating IBD segments...")
    for i in 1:k
        # assign parents (mom has to be from same population as dad)
        dad = rand(male_idx)
        mom = 0
        while true
            mom = rand(female_idx)
            populations[mom] == populations[dad] && break
        end
        push!(fathers, dad)
        push!(mothers, mom)
        # recombination
        recombine!(@view(x1[i, :]), @view(x2[i, :]), @view(h1[dad, :]),
                   @view(h2[dad, :]), @view(h1[mom, :]), @view(h2[mom, :]))
        # update progress
        next!(pmeter)
    end
    # combine offsprings and parents
    H1 = [h1; x1]
    H2 = [h2; x2]
    return H1, H2, fathers, mothers
end

function recombination_segments(breakpoints::Vector{Int}, snps::Int)
    start = 1
    result = UnitRange{Int}[]
    for bkpt in breakpoints
        push!(result, start:bkpt)
        start = bkpt + 1
    end
    push!(result, breakpoints[end]+1:snps)
    return result
end

function recombine!(child_h1, child_h2, dad_h1, dad_h2, mom_h1, mom_h2)
    p = length(child_h1)
    recombinations = rand(1:5)
    breakpoints = sort!(sample(1:p, recombinations, replace=false))
    segments = recombination_segments(breakpoints, p)
    for segment in segments
        dad_hap = rand() < 0.5 ? dad_h1 : dad_h2
        mom_hap = rand() < 0.5 ? mom_h1 : mom_h2
        copyto!(@view(child_h1[segment]), @view(dad_hap[segment]))
        copyto!(@view(child_h2[segment]), @view(mom_hap[segment]))
    end
end

function write_plink(outfile::AbstractString, x1::AbstractMatrix, x2::AbstractMatrix)
    n, p = size(x1)
    x = SnpArray(outfile * ".bed", n, p)
    for j in 1:p, i in 1:n
        c = x1[i, j] + x2[i, j]
        if c == 0
            x[i, j] = 0x00
        elseif c == 1
            x[i, j] = 0x02
        elseif c == 2
            x[i, j] = 0x03
        else
            error("matrix entries should be 0, 1, or 2 but was $c!")
        end
    end
    # create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(outfile * ".bim", "w") do f
        for i in 1:p
            println(f, "1\tsnp$i\t0\t$(100i)\t1\t2")
        end
    end
    # create .fam file structure: https://www.cog-genomics.org/plink2/formats#fam
    open(outfile * ".fam", "w") do f
        for i in 1:n
            println(f, "$i\t1\t0\t0\t1\t-9")
        end
    end
    return nothing
end

function make_partition_mapfile(filename, p::Int)
    map_cM = LinRange(0.0, Int(p / 10000), p)
    open(filename, "w") do io
        println(io, "Chromosome\tPosition(bp)\tRate(cM/Mb)\tMap(cM)")
        for i in 1:p
            println(io, "chr1\t", 100i, '\t', 0.01rand(), '\t', map_cM[i])
        end
    end
end

function make_rapid_mapfile(filename, p::Int)
    map_cM = LinRange(0.0, Int(p / 10000), p)
    open(filename, "w") do io
        for i in 1:p
            println(io, i, '\t', map_cM[i])
        end
    end
end

function process_rapid_output(inputfile, outputfile)
    writer = open(outputfile, "w")
    df = readdlm(inputfile)
    println(writer, "CHR ID1 HID1 ID2 HID2 BP.start BP.end site.start site.end cM FAM1 FAM2")
    for r in eachrow(df)
        chr, id1, id2, hap1, hap2, start_pos, end_pos, genetic_len, start_site, end_site =
            Int(r[1]), Int(r[2]), Int(r[3]), Int(r[4]), Int(r[5]), Int(r[6]), Int(r[7]),
            r[8], Int(r[9]), Int(r[10])
        println(writer, chr, ' ', id1, ' ', hap1, ' ', id2, ' ', hap2, ' ',
            start_pos, ' ', end_pos, ' ', start_site, ' ', end_site, ' ',
            genetic_len, ' ', 1, ' ', 1)
    end
    close(writer)
end

function make_bgen_samplefile(filename, n)
    open(filename, "w") do io
        println(io, "ID_1 ID_2 missing sex")
        println(io, "0 0 0 D")
        for i in 1:n
            println(io, "$i 1 0 1")
        end
    end
end

function simulate_genotypes(n, p, offsprings, seed)
    Random.seed!(seed)
    h1, h2, populations, diff_markers = simulate_pop_structure(n, p)

    # simulate random mating to get IBD segments
    x1, x2 = simulate_IBD(h1, h2, populations, offsprings)

    # write phased genotypes to VCF format
    write_vcf("sim.phased.vcf.gz", x1, x2)

    # write unphased genotypes to PLINK binary format
    outfile = "sim"
    write_plink(outfile, x1, x2)

    # save pop1/pop2 index and unually differentiated marker indices
    writedlm("populations.txt", populations)
    writedlm("diff_markers.txt", diff_markers)

    # generate fake map file
    make_partition_mapfile("sim.partition.map", p)

    # also generate QC file that contains all SNPs and all samples
    snpdata = SnpData("sim")
    snpIDs = snpdata.snp_info[!, :snpid]
    sampleIDs = Matrix(snpdata.person_info[!, 1:2])
    writedlm("variants_qc.txt", snpIDs)
    writedlm("samples_qc.txt", sampleIDs)
end

function make_knockoffs(n, p, offsprings, seed)
    # fake map file for rapid
    make_rapid_mapfile("sim.rapid.map", p)

    # run rapid
    Random.seed!(seed)
    vcffile = "sim.phased.vcf.gz"
    mapfile = "sim.rapid.map"
    outfolder = "rapid"
    d = 3    # minimum IBD length in cM
    w = 3    # number of SNPs per window
    r = 10   # number of runs
    s = 2    # Minimum number of successes to consider a hit
    @time rapid(rapid_exe, vcffile, mapfile, d, outfolder, w, r, s)

    # unzip and postprocess rapid output to suit snpknock2
    run(pipeline(`gunzip -c ./rapid/results.max.gz`, stdout="./rapid/results.max"))
    process_rapid_output("./rapid/results.max", "sim.snpknock.ibdmap")

    # convert VCF to BGEN format
    outfile = "sim.bgen"
    run(`$qctools_exe -g $vcffile -og $outfile`)
    make_bgen_samplefile("sim.sample", n + offsprings)

    # snpknock2 arguments
    bgenfile = "sim"
    sample_qc = "samples_qc.txt"
    variant_qc = "variants_qc.txt"
    mapfile = "sim.partition.map"
    partfile = "sim.partition.txt"
    ibdfile = "sim.snpknock.ibdmap"
    K = 10
    cluster_size_min = 1000
    cluster_size_max = 10000
    hmm_rho = 1
    hmm_lambda = 1e-3
    windows = 0
    n_threads = Threads.nthreads()
    compute_references = true
    generate_knockoffs = true
    outfile = "sim.knockoffs"
    @time snpknock2(snpknock2_exe, bgenfile, sample_qc, variant_qc, mapfile, partfile, ibdfile,
        K, cluster_size_min, cluster_size_max, hmm_rho, hmm_lambda, windows, n_threads,
        seed, compute_references, generate_knockoffs, outfile)
end

# there is 100 causal snps, 50 on normally differentiated SNPs and 50 on abnormally diff snps
function simulate_beta_and_y(x::AbstractMatrix, diff_markers, seed)
    n, p = size(x)
    Random.seed!(seed)
    h2 = 0.5 # heritability
    k = 100 # number of causal snps
    d = Normal(0, sqrt(h2 / (2k))) # from paper: Efficient Implementation of Penalized Regression for Genetic Risk Prediction
    # simulate β
    β = zeros(p)
    norm_markers = setdiff(1:p, diff_markers)
    norm_causal_snps = sample(norm_markers, k >> 1, replace=false)
    β[norm_causal_snps] .= rand(d, k >> 1)
    shuffle!(diff_markers)
    diff_causal_markers = diff_markers[1:k >> 1]
    β[diff_causal_markers] .= rand(d, k >> 1)
    # simulate y
    ϵ = Normal(0, 1 - h2)
    y = x * β + rand(ϵ, n)
    # save results
    writedlm("y_true.txt", y)
    writedlm("beta_true.txt", β)
    writedlm("normal_markers.txt", norm_markers)
    writedlm("normal_causal_markers.txt", norm_causal_snps)
    writedlm("diff_markers.txt", diff_markers)
    writedlm("diff_causal_markers.txt", diff_causal_markers)
    return y, β
end

function R2(X::AbstractMatrix, y::AbstractVector, β̂::AbstractVector)
    μ = y - X * β̂
    tss = y .- mean(y)
    return 1 - dot(μ, μ) / dot(tss, tss)
end

function TP(correct_groups, signif_groups)
    return length(signif_groups ∩ correct_groups) / length(correct_groups)
end

function FDR(correct_groups, signif_groups)
    FP = length(signif_groups) - length(signif_groups ∩ correct_groups) # number of false positives
    # FPR = FP / (FP + TN) # https://en.wikipedia.org/wiki/False_positive_rate#Definition
    FDR = FP / length(signif_groups)
    return FDR
end

function tune_k(y::AbstractVector, xko_la::AbstractMatrix, original::Vector{Int},
    knockoff::Vector{Int}, groups::Vector{Int}, fdr::Float64, best_k::Int
    )
    # do a grid search for best sparsity level
    best_β = Float64[]
    best_err = Inf
    for cur_k in best_k:5:round(Int, 1.5best_k)
        result = fit_iht(y, xko_la, k=cur_k, init_beta=true, max_iter=500)
        W = coefficient_diff(result.beta, groups, original, knockoff)
        τ = threshold(W, fdr, :knockoff)
        detected = count(x -> x ≥ τ, W)
        if abs(detected - best_k) < best_err
            best_β = copy(result.beta)
            best_err = abs(detected - best_k)
        end
        println("wrapped CV says best_k = $best_k; using k = $cur_k detected $detected")
        GC.gc()
    end
    return best_β
end

function get_signif_groups(β, groups)
    correct_groups = Int[]
    for i in findall(!iszero, β)
        g = groups[i]
        g ∈ correct_groups || push!(correct_groups, g)
    end
    return correct_groups
end

function summarize_result(groups, β_true, β_iht, β_iht_knockoff, β_iht_knockoff_cv, β_lasso, β_lasso_knockoff)
    length(groups) == length(β_true) || error("summarize_result: groups and β_true have different length!")
    df = DataFrame(Metric = String[], IHT = Float64[], IHT_ko = Float64[],
        IHT_ko_cv = Float64[], LASSO = Float64[], LASSO_ko = Float64[])

    # count non-zero entries of β returned from cross validation
    push!(df, hcat("beta_non_zero_count", count(!iszero, β_iht),
        count(!iszero, vec(readdlm("iht.knockoff.beta"))),
        count(!iszero, vec(readdlm("iht.knockoff.cv.beta"))),
        count(!iszero, β_lasso),
        count(!iszero, vec(readdlm("lasso.knockoff.beta")))
        ))
    # count non-zero entries after knockoff filter
    push!(df, hcat("beta_selected", count(!iszero, β_iht),
        count(!iszero, β_iht_knockoff),
        count(!iszero, β_iht_knockoff_cv), count(!iszero, β_lasso),
        count(!iszero, β_lasso_knockoff)
        ))
    # count TP proportion
    correct_snps = findall(!iszero, vec(readdlm("beta_true.txt")))
    correct_groups = get_signif_groups(β_true, groups)
    push!(df, hcat("TPP", TP(correct_snps, findall(!iszero, β_iht)),
        TP(correct_groups, get_signif_groups(β_iht_knockoff, groups)),
        TP(correct_groups, get_signif_groups(β_iht_knockoff_cv, groups)),
        TP(correct_groups, get_signif_groups(β_lasso, groups)),
        TP(correct_groups, get_signif_groups(β_lasso_knockoff, groups))
        ))
    # count FDR
    push!(df, hcat("FDR", FDR(correct_snps, findall(!iszero, β_iht)),
        FDR(correct_groups, get_signif_groups(β_iht_knockoff, groups)),
        FDR(correct_groups, get_signif_groups(β_iht_knockoff_cv, groups)),
        FDR(correct_groups, get_signif_groups(β_lasso, groups)),
        FDR(correct_groups, get_signif_groups(β_lasso_knockoff, groups))
        ))
    @show df
    CSV.write("summary.txt", df)
end

function one_simulation(n, p, offsprings, seed)
    cur_dir = pwd() * "/"

    # save all simulated data under data/sim$seed
    isdir("data") || mkdir("data")
    isdir("data/sim$seed") || mkdir("data/sim$seed")
    cd("data/sim$seed")
    data_dir = cur_dir * "data/sim$seed/"

    # simulate phased genotypes with pop structure and cryptic relatedness
    simulate_genotypes(n, p, offsprings, seed)

    # partition genotypes to groups (for making group knockoffs)
    partition(partition_exe, "sim", "sim.partition.map", "variants_qc.txt", "sim.partition.txt")

    # make knockoffs
    make_knockoffs(n, p, offsprings, seed)

    # loop over group resolution
    for group in 0:3
        x = SnpArray(data_dir * "knockoffs/sim.knockoffs_res$(group).bed")
        snpid = SnpData(data_dir * "knockoffs/sim.knockoffs_res$(group)").snp_info.snpid
        knockoff_idx = endswith.(snpid, ".k")
        original = findall(knockoff_idx .== false)
        knockoff = findall(knockoff_idx)
        groups = repeat(Vector{Int}(readdlm(data_dir * "knockoffs/sim.knockoffs_res$(group)_grp.txt",
            header=true)[1][:, 2]), inner=2) .+ 1
        xla = convert(Matrix{Float64}, @view(x[:, original]), center=true, scale=true, impute=true)
        xko_la = convert(Matrix{Float64}, x, center=true, scale=true, impute=true)
        diff_markers = Vector{Int}(vec(readdlm(data_dir * "diff_markers.txt")))
        group_dir = cur_dir * "res$group/"
        isdir(group_dir) || mkdir(group_dir)
        cd(group_dir)

        for fdr in [0.05, 0.1, 0.25, 0.5]
            top_dir = group_dir * "fdr$fdr/"
            new_dir = group_dir * "fdr$fdr/sim$seed"
            isdir(top_dir) || mkdir(top_dir)
            isdir(new_dir) || mkdir(new_dir)
            cd(new_dir)

            #
            # simulate beta and y
            #
            Random.seed!(seed)
            y, β = simulate_beta_and_y(xla, diff_markers, seed)

            #
            # Run standard IHT
            #
            Random.seed!(seed)
            path = 10:10:200
            mses = cv_iht(y, xla, path=path, init_beta=true)
            GC.gc()
            Random.seed!(seed)
            k_rough_guess = path[argmin(mses)]
            dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
            mses_new = cv_iht(y, xla, path=dense_path, init_beta=true)
            GC.gc()
            Random.seed!(seed)
            result = fit_iht(y, xla, k=dense_path[argmin(mses_new)], init_beta=true, max_iter=500)
            @show result
            β_iht = result.beta
            writedlm("iht.beta", β_iht)
            GC.gc()

            #
            # Run standard lasso
            #
            Random.seed!(seed)
            cv = glmnetcv(xla, y, nfolds=5, parallel=true)
            β_lasso = GLMNet.coef(cv)
            writedlm("lasso.beta", β_lasso)

            #
            # run knockoff IHT
            #
            Random.seed!(seed)
            path = 10:10:200
            mses = cv_iht(y, xko_la, path=path, init_beta=true)
            GC.gc()
            Random.seed!(seed)
            k_rough_guess = path[argmin(mses)]
            dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
            mses_new = cv_iht(y, xko_la, path=dense_path, init_beta=true)
            GC.gc()
            Random.seed!(seed)
            result = fit_iht(y, xko_la, k=dense_path[argmin(mses_new)],
                init_beta=true, max_iter=500)
            @show result
            β_iht_knockoff = extract_beta(result.beta, fdr, groups, original, knockoff)
            writedlm("iht.knockoff.beta", result.beta)
            writedlm("iht.knockoff.postfilter.beta", β_iht_knockoff)

            #
            # run knockoff IHT with wrapped cross validation
            #
            Random.seed!(seed)
            path = 10:10:200
            z = ones(Float64, size(xla, 1))
            mses = cv_iht_knockoff(y, xko_la, z, original, knockoff, fdr, path=path,
                init_beta=true, group_ko=groups)
            GC.gc()
            Random.seed!(seed)
            k_rough_guess = path[argmin(mses)]
            dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
            mses_new = cv_iht_knockoff(y, xko_la, z, original, knockoff, fdr,
                path=dense_path, init_beta=true, group_ko=groups)
            GC.gc()
            # adjust sparsity level so it best matches sparsity chosen by ko filter
            Random.seed!(seed)
            best_k = dense_path[argmin(mses_new)]
            best_β = tune_k(y, xko_la, original, knockoff, groups, fdr, best_k)
            β_iht_knockoff_cv = extract_beta(best_β, fdr, groups, original, knockoff)
            writedlm("iht.knockoff.cv.beta", best_β)
            writedlm("iht.knockoff.cv.postfilter.beta", β_iht_knockoff_cv)

            #
            # Run knockoff lasso
            #
            Random.seed!(seed)
            cv = glmnetcv(xko_la, y, nfolds=5, parallel=true)
            β_lasso_knockoff = extract_beta(GLMNet.coef(cv), fdr, groups, original, knockoff)
            writedlm("lasso.knockoff.beta", GLMNet.coef(cv))
            writedlm("lasso.knockoff.postfilter.beta", β_lasso_knockoff)

            #
            # process results
            #
            summarize_result(groups[original], β, β_iht, β_iht_knockoff, β_iht_knockoff_cv, β_lasso, β_lasso_knockoff)
            GC.gc()
        end
        GC.gc()
    end
end

n = 5000
p = 50000
offsprings = 500
seed = parse(Int, ARGS[1])
one_simulation(n, p, offsprings, seed)


export PandaDyn
using PandaRobot

# Useless for now...

mutable struct PandaDyn <: DynamicsModel
  # state: q
  x_dim # joints
  u_dim

  # Parameters that can be updated
  B
end
function PandaDyn()
  x_dim = 9 # 9 joints
  u_dim = 9

  PandaDyn(x_dim,u_dim,[])
end

######
# Cost
######

##################
# Initializations
##################