# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# 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.

"""
    Optimizer(; <keyword arguments>)

Return a progressive-hedging optimizer.

...
# Arguments
- `subproblem_optimizer::AbstractOptimizer`: MathOptInterface solver capable of solving quadratic programs.
- `penalty::AbstractPenalizer = Fixed()`: Specify penalty update procedure (Fixed, Adaptive)
- `execution::AbstractExecuter = Serial`: Specify how algorithm should be executed (Serial, Synchronous, Asynchronous). Distributed variants requires worker cores.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`Linearized`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
- <keyword arguments>: Algorithm specific parameters, consult individual docstrings (see above list) for list of possible arguments and default values.
...
"""
mutable struct Optimizer <: AbstractStructuredOptimizer
    subproblem_optimizer
    sub_params::Dict{MOI.AbstractOptimizerAttribute, Any}
    execution::AbstractExecution
    penalizer::AbstractPenalizer
    penaltyterm::AbstractPenaltyTerm
    parameters::ProgressiveHedgingParameters{Float64}

    status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    dual_status::MOI.ResultStatusCode
    raw_status::String
    solve_time::Float64
    progressivehedging::Union{AbstractProgressiveHedging, Nothing}

    function Optimizer(; subproblem_optimizer = nothing,
                       execution::AbstractExecution = nworkers() == 1 ? Serial() : Synchronous(),
                       penalty::AbstractPenalizer = Fixed(),
                       penaltyterm::AbstractPenaltyTerm = Quadratic(),
                       kw...)
        return new(subproblem_optimizer,
                   Dict{MOI.AbstractOptimizerAttribute, Any}(),
                   execution,
                   penalty,
                   penaltyterm,
                   ProgressiveHedgingParameters{Float64}(; kw...),
                   MOI.OPTIMIZE_NOT_CALLED,
                   MOI.NO_SOLUTION,
                   MOI.NO_SOLUTION,
                   "Progressive-hedging optimizer has not been run.",
                   NaN,
                   nothing)
    end
end

# Interface #
# ========================== #
function supports_structure(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    return true
end

function default_structure(::UnspecifiedInstantiation, optimizer::Optimizer)
    if optimizer.execution isa Serial && nworkers() == 1
        return ScenarioDecomposition()
    else
        return DistributedScenarioDecomposition()
    end
end

function check_loadable(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    if optimizer.subproblem_optimizer === nothing
        msg = "Subproblem optimizer not set. Consider setting `SubProblemOptimizer` attribute."
        throw(UnloadableStructure{Optimizer, ScenarioDecompositionStructure}(msg))
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{ScenarioProblems}})
    if !(optimizer.execution isa Serial)
        @warn "Distributed execution policies are not compatible with a single-core scenario-decomposition structure. Switching to `Serial` execution by default."
        MOI.set(optimizer, Execution(), Serial())
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}})
    if optimizer.execution isa Serial
        @warn "Serial execution not compatible with distributed scenario-decomposition structure. Switching to `Synchronous` execution by default."
        MOI.set(optimizer, Execution(), Synchronous())
    end
    return nothing
end

function load_structure!(optimizer::Optimizer, structure::ScenarioDecompositionStructure, x₀::AbstractVector)
    # Sanity check
    check_loadable(optimizer, structure)
    # Restore structure if optimization has been run before
    restore_structure!(optimizer)
    # Ensure that execution policy is compatible
    ensure_compatible_execution!(optimizer, structure)
    # Set subproblem optimizers
    set_subproblem_optimizer!(structure, optimizer.subproblem_optimizer)
    # Set subproblem optimizer attributes
    for (attr, value) in optimizer.sub_params
        MOI.set(scenarioproblems(structure), attr, value)
    end
    # Create new progressive-hedging algorithm
    optimizer.progressivehedging = ProgressiveHedgingAlgorithm(structure,
                                                               x₀,
                                                               optimizer.execution,
                                                               optimizer.penalizer,
                                                               optimizer.penaltyterm;
                                                               type2dict(optimizer.parameters)...)
    return nothing
end

function restore_structure!(optimizer::Optimizer)
    if optimizer.progressivehedging !== nothing
        restore_subproblems!(optimizer.progressivehedging)
    end
    return nothing
end

function MOI.optimize!(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    # Run progressive-hedging procedure
    optimizer.status = optimizer.progressivehedging()
    # Check if optimal
    if optimizer.status == MOI.OPTIMAL
        optimizer.primal_status = MOI.FEASIBLE_POINT
        optimizer.dual_status = MOI.FEASIBLE_POINT
        optimizer.raw_status = "Progressive-hedging procedure converged to optimal solution."
    end
    # Extract solve time
    optimizer.solve_time = optimizer.progressivehedging.progress.tlast - optimizer.progressivehedging.progress.tinit
    return nothing
end

function num_iterations(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return num_iterations(optimizer.progressivehedging)
end

function optimizer_name(optimizer::Optimizer)
    return "$(str(optimizer.execution))Progressive-hedging with $(str(optimizer.penalizer))"
end

# MOI #
# ========================== #
function MOI.get(optimizer::Optimizer, ::MOI.Silent)
    return !MOI.get(optimizer, MOI.RawParameter("log"))
end

function MOI.set(optimizer::Optimizer, attr::MOI.Silent, flag::Bool)
    MOI.set(optimizer, MOI.RawParameter("log"), !flag)
    optimizer.sub_params[attr] = flag
    return nothing
end

function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    return getfield(optimizer.parameters, name)
end

function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    setfield!(optimizer.parameters, name, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TimeLimitSec)
    limit = MOI.get(optimizer, MOI.RawParameter("time_limit"))
    return isinf(limit) ? nothing : limit
end

function MOI.set(optimizer::Optimizer, ::MOI.TimeLimitSec, limit::Union{Real, Nothing})
    limit = limit === nothing ? Inf : limit
    MOI.set(optimizer, MOI.RawParameter("time_limit"), limit)
    return
end

function MOI.get(optimizer::Optimizer, ::PrimalTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₁"))
end

function MOI.set(optimizer::Optimizer, ::PrimalTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₁"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::DualTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₂"))
end

function MOI.set(optimizer::Optimizer, ::DualTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₂"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MasterOptimizer)
    return MOI.get(optimizer, SubProblemOptimizer())
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizer)
    if optimizer.subproblem_optimizer === nothing
        return MOI.get(optimizer, MasterOptimizer())
    end
    return MOI.OptimizerWithAttributes(optimizer.subproblem_optimizer, collect(optimizer.sub_params))
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizer, optimizer_constructor)
    optimizer.subproblem_optimizer = optimizer_constructor
    # Clear any old parameters
    empty!(optimizer.sub_params)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute)
    if !haskey(optimizer.sub_params, attr)
        error("Subproblem optimizer attribute $(attr) has not been set.")
    end
    return optimizer.sub_params[attr]
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute, value)
    optimizer.sub_params[attr] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, param::RawSubProblemOptimizerParameter)
    moi_param = MOI.RawParameter(param.name)
    if !haskey(optimizer.sub_params, moi_param)
        error("Subproblem optimizer attribute $(param.name) has not been set.")
    end
    return optimizer.sub_params[moi_param]
end

function MOI.set(optimizer::Optimizer, param::RawSubProblemOptimizerParameter, value)
    moi_param = MOI.RawParameter(param.name)
    optimizer.sub_params[moi_param] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Execution)
    return optimizer.execution
end

function MOI.set(optimizer::Optimizer, ::Execution, execution::AbstractExecution)
    optimizer.execution = execution
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Penalizer)
    return optimizer.penalizer
end

function MOI.set(optimizer::Optimizer, ::Penalizer, penalizer::AbstractPenalizer)
    optimizer.penalizer = penalizer
    return nothing
end

function MOI.get(optimizer::Optimizer, ::PenaltyTerm)
    return optimizer.penaltyterm
end

function MOI.set(optimizer::Optimizer, ::PenaltyTerm, penaltyterm::AbstractPenaltyTerm)
    optimizer.penaltyterm = penaltyterm
    return nothing
end

function MOI.get(optimizer::Optimizer, param::ExecutionParameter)
    return MOI.get(optimizer.execution, param)
end

function MOI.set(optimizer::Optimizer, param::ExecutionParameter, value)
    MOI.set(optimizer.execution, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, param::PenalizationParameter)
    return MOI.get(optimizer.penalizer, param)
end

function MOI.set(optimizer::Optimizer, param::PenalizationParameter, value)
    MOI.set(optimizer.penalizer, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    return optimizer.status
end

function MOI.get(optimizer::Optimizer, ::MOI.PrimalStatus)
    return optimizer.primal_status
end

function MOI.get(optimizer::Optimizer, ::MOI.DualStatus)
    return optimizer.dual_status
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.raw_status
end

function MOI.get(optimizer::Optimizer, ::MOI.VariablePrimal, index::MOI.VariableIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return decision(optimizer.progressivehedging, index)
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintPrimal, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintPrimal(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintDual, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintDual(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return objective_value(optimizer.progressivehedging)
end

function MOI.get(optimizer::Optimizer, ::MOI.DualObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    if optimizer.status == MOI.OPTIMAL
        return objective_value(optimizer.progressivehedging)
    elseif optimizer.status == MOI.INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? Inf : -Inf
    elseif optimizer.status == MOI.DUAL_INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? -Inf : Inf
    else
        return NaN
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.solve_time
end

function MOI.get(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute})
    # Fallback to through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute}, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentModelAttribute)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, index)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentModelAttribute, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.is_empty(optimizer::Optimizer)
    return optimizer.progressivehedging === nothing
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::Optimizer, ::MOI.RawParameter) = true
MOI.supports(::Optimizer, ::AbstractStructuredOptimizerAttribute) = true
MOI.supports(::Optimizer, ::RelativeTolerance) = false
MOI.supports(::Optimizer, ::RawInstanceOptimizerParameter) = true
MOI.supports(::Optimizer, ::AbstractProgressiveHedgingAttribute) = true

# High-level attribute setting #
# ========================== #
"""
    get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)

Return the value associated with the penalization-specific attribute named `name` in `stochasticprogram`.

See also: [`set_penalization_attribute`](@ref), [`set_penalization_attributes`](@ref).
"""
function get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)
    return return MOI.get(optimizer(stochasticprogram), RawPenalizationParameter(name))
end
"""
    set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

Sets the penalization-specific attribute identified by `name` to `value`.

"""
function set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawPenalizationParameter(String(name)), value)
end
"""
    set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

Given a list of `attribute => value` pairs or a collection of keyword arguments, calls
`set_penalization_attribute(stochasticprogram, attribute, value)` for each pair.

"""
function set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_penalization_attribute(stochasticprogram, name, value)
    end
end
function set_penalization_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_penalization_attribute(stochasticprogram, name, value)
    end
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


########################################################################
# 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


#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)


"""
    readsiggen(file_path::String; T::Type)

Reads a [SigGen](https://github.com/radforddc/icpc_siggen) configuration file (ending in `.config`) in 'file_path'
and returns a dictionary of all parameters. Non-existing parameteres are set to `0`.

## Arguments
* `file_path::String`: File path leading to the [SigGen](https://github.com/radforddc/icpc_siggen) configuration file.

## Keywords
* `T::Type`: Type of the parameters in the output dictionary. Default is `Float64`.
"""
function readsiggen(file_path::String; T::Type=Float64)

    detector_name = split(basename(file_path), ".config")[1]

    config = Dict("name"       => detector_name,
        "verbosity_level"      => 0,
        "xtal_length"          => 0,
        "xtal_radius"          => 0,
        "top_bullet_radius"    => 0, # not implemented in later conversions
        "bottom_bullet_radius" => 0, # not implemented in later conversions
        "pc_length"            => 0,
        "pc_radius"            => 0,
        "bulletize_PC"         => 0,
        "wrap_around_radius"   => 0,
        "ditch_depth"          => 0,
        "ditch_thickness"      => 0,
        "bottom_taper_length"  => 0,
        "hole_length"          => 0,
        "hole_radius"          => 0,
        "outer_taper_length"   => 0,
        "inner_taper_length"   => 0,
        "outer_taper_width"    => 0,
        "inner_taper_width"    => 0,
        "taper_angle"          => 0,
        "taper_length"         => 0,
        "Li_thickness"         => 0,
        "energy"               => 0,
        "xtal_grid"            => 0,
        "impurity_z0"          => 0,
        "impurity_gradient"    => 0,
        "impurity_surface"     => 0,
        "xtal_HV"              => 0,
        "max_iterations"       => 0,
        "write_field"          => 0,
        "write_WP"             => 0,
        "drift_name"           => "",
        "field_name"           => "",
        "wp_name"              => "",
        "xtal_temp"            => 0,
        "preamp_tau"           => 0,
        "time_steps_calc"      => 0,
        "step_time_calc"       => 0,
        "step_time_out"        => 0,
        "charge_cloud_size"    => 0,
        "use_diffusion"        => 0,
        "surface_drift_vel_factor" => 0);

    file = open(file_path) do f
        temp          = readlines(f)
        config_file   = Dict()
        for line in temp
            if line  != "" && string(line[1]) != "#"
                line  = strip(split(line, "#")[1])
                name  = ""
                value = ""
                if length(split(line, "\t")) > 1
                    name  = split(line, "\t")[1]
                    value = strip(split(line, "\t")[2])
                else
                    name  = split(line, " ")[1]
                    value = strip(split(line, " ")[end])
                end

                config_file[name] = value

            end
        end
        return config_file
    end

    config_keys = keys(config)
    file_keys   = keys(file)

    for key in file_keys
        if key in config_keys
            if key == "wp_name" || key == "drift_name" || key == "field_name"
                config[key] = string(file[key])
            elseif key == "taper_angle" && parse(Float64, file["taper_angle"]) != 0
                config["outer_taper_width"] = tan(deg2rad(parse(T, file[key]))) * parse(T, file["outer_taper_length"])
                config["inner_taper_width"] = tan(deg2rad(parse(T, file[key]))) * parse(T, file["inner_taper_length"])
                config[key]                 = parse(T, file[key])
            else
                config[key] = parse(T, file[key])
            end
        end
    end
    return config
end

"""
    siggentodict(config::Dict; units::Dict)

Converts the dictionary containing the parameters from a [SigGen](https://github.com/radforddc/icpc_siggen) configuration file
to a dictionary that can be understood by SolidStateDetectors.jl.

## Arguments
* `config::Dict`: Dictionary containing [SigGen](https://github.com/radforddc/icpc_siggen) parameters (output of `readsiggen()``).

## Keywords
* `units::Dict`: Units used in [SigGen](https://github.com/radforddc/icpc_siggen) configuration file (set to `"mm"`, `"deg"`, `"V"` and `"K"`).
    The dictionary needs the fields `"length"`, `"angle"`, `"potential"` and `"temperature"`.
"""
function siggentodict(config::Dict;
        units::Dict = Dict("length" => "mm","angle" => "deg","potential" => "V","temperature" => "K"))

    #>----------------Settings START-------------------------------------------<

    #>----------------Setting the z and r limits-------------------------------<
    z_limit      = ceil(config["xtal_length"]; sigdigits=1);
    if z_limit - config["xtal_length"] <= 0
        z_limit += 10
    end
    r_limit      = ceil(config["xtal_radius"]; sigdigits=1);
    if r_limit - config["xtal_radius"] <= 0
        r_limit += 10
    end

    #>----------------Setting the Axes & Grid dictionaries---------------------<
    axes         = Dict("r"           => Dict("to"         => r_limit,
                                              "boundaries" => "inf"),
                        "phi"         => Dict("from"       => 0,
                                              "to"         => 0,
                                              "boundaries" => "periodic"),
                        "z"           => Dict("from"       => -10,
                                              "to"         => z_limit,
                                              "boundaries" => Dict("left" => "inf", "right" => "inf"))
                        )

    grid         = Dict("coordinates" => "cylindrical",
                        "axes"        => axes);

    #>----------------Settings END---------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Volume START---------------------------------------------<

    #>----------------Main Volume -> Initial Cylinder--------------------------<
    # r and L are given by the input file. Cylindrical symmetry is set to default
    main_vol     = Dict("translate"   => Dict("z"    => config["xtal_length"]/2,
                        "tube"        => Dict(
                        "name"        => "Initial Cylinder",
                        "r"           => Dict("from" => 0.0, "to" => config["xtal_radius"]),
                        "phi"         => Dict("from" => 0.0, "to" => 360.0),
                        "h"           => config["xtal_length"])));

    #>----------------Upper Cone / Taper on top--------------------------------<
    # If not existing, everything will be set to 0
    upper_cone   = Dict("translate"   => Dict("z"      => config["xtal_length"] -
                                            config["outer_taper_length"]/2,
                        "cone"        => Dict(
                        "name"        => "Upper Cone",
                        "r"           => Dict("bottom" => Dict("from" => config["xtal_radius"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0)),
                                              "top"    => Dict("from" => config["xtal_radius"] -
                                                                            config["outer_taper_width"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0))),
                        "phi"         => Dict("from"   => 0.0, "to" => 360.0),
                        "h"           => config["outer_taper_length"])));

    #>----------------Lower Cone / Taper on bottom-----------------------------<
    # 45-degree taper at bottom of ORTEC-type crystal
    lower_cone   = Dict("translate"   => Dict("z"      => 0.4*config["taper_length"],
                        "cone"        => Dict(
                        "name"        => "Lower Cone",
                        "r"           => Dict("top"    => Dict("from" => config["xtal_radius"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0)),
                                              "bottom" => Dict("from" => config["xtal_radius"] -
                                                                            1.2 * config["taper_length"],
                                                               "to"   => ceil(config["xtal_radius"] + 1.0))),
                        "phi"         => Dict("from"   => 0.0, "to" => 360.0),
                        "h"           => 1.2*config["taper_length"])));

    #>----------------Ditch around the bottom contact--------------------------<
    # If not existing, everything will be set to 0
    ditch        = Dict("translate"   => Dict("z"    => config["ditch_depth"]/2,
                        "tube"        => Dict(
                        "name"        => "Ditch",
                        "r"           => Dict("from" => config["wrap_around_radius"] -
                                                            config["ditch_thickness"],
                                              "to"   => config["wrap_around_radius"]),
                        "phi"         => Dict("from" => 0.0, "to" => 360.0),
                        "h"           => config["ditch_depth"])));

    #>----------------Borehole on bottom as contact----------------------------<
    # If not existing, everything will be set to 0
    borehole = if config["bulletize_PC"] == 1
        if config["pc_length"] < config["pc_radius"]
            Dict( "union"   => [
            Dict( "tube"    => Dict(
                  "r"       => config["pc_radius"] - config["pc_length"],
                  "h"       => 1.2*config["pc_length"],
                  "origin"  => Dict("z" => 0.4*config["pc_length"]))),
            Dict( "torus"   => Dict(
                  "r_torus" => config["pc_radius"] - config["pc_length"],
                  "r_tube"  => config["pc_length"]))])
        elseif config["pc_length"] > config["pc_radius"]
            Dict( "union"  => [
            Dict( "tube"   => Dict(
                  "r"      => config["pc_radius"],
                  "h"      => 1.2*config["pc_length"] - config["pc_radius"],
                  "origin" => Dict("z" => (0.8*config["pc_length"] - config["pc_radius"])/2))),
            Dict( "sphere" => Dict(
                  "r"      => config["pc_radius"],
                  "origin" => Dict("z" => config["pc_length"] - config["pc_radius"])))])
        else
            Dict( "sphere" => Dict(
                       "r" => config["pc_radius"]))
        end
    else # not bulletized
            Dict( "tube"   => Dict(
                  "name"   => "Borehole",
                  "r"      => Dict("from" => 0.0, "to" => config["pc_radius"]),
                  "h"      => config["pc_length"],
                  "origin" => Dict("z" => config["pc_length"]/2)));
    end

    #>----------------Hole on top of the detector------------------------------<
    # If not existing, everything will be set to 0
    hole = if config["inner_taper_width"] != 0 && config["inner_taper_length"] != 0
        Dict("translate"   => Dict("z"      => config["xtal_length"] -
                               config["inner_taper_length"]/2,
             "cone"        => Dict(
             "name"        => "Inner Cone",
             "r"           => Dict("top"    => Dict("from" => 0,
                                                   "to"   => config["hole_radius"] +
                                                                config["inner_taper_width"]),
                                  "bottom" => Dict("from" => 0,
                                                   "to"   => config["hole_radius"])),
             "phi"         => Dict("from"   => 0.0, "to" => 360.0),
             "h"           => config["inner_taper_length"])));

    else
        Dict("translate"   => Dict("z"    => config["xtal_length"] -
                                               0.8*config["hole_length"]/2,
             "tube"        => Dict(
             "name"        => "Top Hole",
             "r"           => Dict("from" => 0.0, "to" => config["hole_radius"]),
             "phi"         => Dict("from" => 0.0, "to" => 360.0),
             "h"           => 1.2*config["hole_length"])));
    end

    #>----------------Subtract volumes-----------------------------------------<
    geometry = main_vol;
    if config["outer_taper_length"] != 0 || config["outer_taper_width"] != 0
        geometry = Dict("difference" => [geometry, upper_cone])
    end
    if config["inner_taper_length"] != 0 || config["inner_taper_width"] != 0 || config["hole_radius"] != 0 || config["hole_length"] != 0
        geometry = Dict("difference" => [geometry, hole])
    end
    if config["taper_length"] != 0
        geometry = Dict("difference" => [geometry, lower_cone])
    end
    if config["ditch_thickness"] != 0
        geometry = Dict("difference" => [geometry, ditch])
    end
    if config["bulletize_PC"] == 1 || (config["pc_radius"] != 0 && config["pc_length"] >= 0.1) # below this, it is the thickness of the layer
        geometry = Dict("difference" => [geometry, borehole])
    end

    #>----------------Volume END-----------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Contacts START-------------------------------------------<

    #>----------------Area/volume of point contact-----------------------------<
    #>----------------Contact 1------------------------------------------------<
    # Depending on whether the contact is a borehole or not

    geometry_1 = if config["bulletize_PC"] == 1.0
        if config["pc_length"] < config["pc_radius"]
            Dict( "union"      => [
            Dict( "tube"       => Dict(
                  "r"          => config["pc_radius"] - config["pc_length"],
                  "h"          => 0,
                  "origin"     => Dict("z" => config["pc_length"]))),
            Dict( "torus"      => Dict(
                  "r_torus"    => config["pc_radius"] - config["pc_length"],
                  "r_tube"     => Dict("from" => config["pc_length"],
                                       "to"   => config["pc_length"]),
                  "theta"      => Dict("from" => 0.0, "to" => 90.0)))])
        elseif config["pc_length"] > config["pc_radius"]
            Dict( "union"      => [
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"], "to" => config["pc_radius"]),
                  "h"          => config["pc_length"] - config["pc_radius"],
                  "origin"     => Dict("z" => (config["pc_length"] - config["pc_radius"])/2))),
            Dict( "difference" => [
            Dict( "sphere"     => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]),
                  "origin"     => Dict("z" => config["pc_length"] - config["pc_radius"]))),
            Dict( "tube"       => Dict(
                  "r"          => 1.2*config["pc_radius"],
                  "h"          => 1.2*config["pc_radius"],
                  "origin"     => Dict("z" => config["pc_length"] - 1.6*config["pc_radius"])))])])
        else
            Dict( "difference" => [
            Dict( "sphere"     => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]))),
            Dict( "tube"       => Dict(
                  "r"          => 1.2*config["pc_radius"],
                  "h"          => 1.2*config["pc_radius"],
                  "origin"     => Dict("z" => -0.6 * config["pc_radius"])))])
        end
    elseif config["pc_radius"] != 0 && config["pc_length"] >= 5 # below this, it is the thickness of the layer
            Dict( "union"      => [
            # borehole_wall
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["pc_radius"]),
                  "h"          => config["pc_length"],
                  "origin"     => Dict("z" => config["pc_length"]/2))),
            # borehole_top
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => 0.0,
                                       "to"   => config["pc_radius"]),
                  "h"          => 0.0,
                  "origin"     => Dict("z" => config["pc_length"]))),
            # borehole_around
            Dict( "tube"       => Dict(
                  "r"          => Dict("from" => config["pc_radius"],
                                       "to"   => config["wrap_around_radius"] - config["ditch_thickness"]),
                  "h"          => 0.0))])
    elseif config["pc_radius"] != 0 && config["pc_length"] < 5
            Dict( "tube"       => Dict(
                  "r"          => config["pc_radius"],
                  "h"          => config["pc_length"],
                  "origin"     => Dict("z" => config["pc_length"]/2)))
    end

    #>----------------Define contact 1-----------------------------------------<

    contact_1 = Dict("material"    => "HPGe",
                     "id"          => 1,
                     "potential"   => 0.0,
                     "geometry"    => geometry_1);

    #>----------------Contact 2------------------------------------------------<
    geometry_2 = Dict[]

    if config["wrap_around_radius"] != 0.0
        bottom = Dict("tube" => Dict(
                      "r"    => Dict("from" => config["wrap_around_radius"],
                                       "to" => config["xtal_radius"] - config["taper_length"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => 0.0))
        push!(geometry_2, bottom)
    end

    wall   = Dict("translate" => Dict("z" => config["xtal_length"]/2 + config["taper_length"]/2 - config["outer_taper_length"]/2,
                  "tube" => Dict(
                  "r"    => Dict("from" => config["xtal_radius"],
                                 "to"   => config["xtal_radius"]),
                  "phi"  => Dict("from" => 0.0,
                                   "to" => 360.0),
                  "h"    => config["xtal_length"] - config["taper_length"] - config["outer_taper_length"])))
    push!(geometry_2, wall)

    top    = Dict("translate" => Dict("z" => config["xtal_length"],
                  "tube" => Dict(
                  "r"    => Dict("from" => config["hole_radius"] + config["inner_taper_width"],
                                   "to" => config["xtal_radius"] - config["outer_taper_width"]),
                  "phi"  => Dict("from" => 0.0,
                                   "to" => 360.0),
                  "h"    => 0.0)))
    push!(geometry_2, top)
    if config["hole_radius"] != 0 && config["hole_length"] != 0 && config["inner_taper_width"] == 0 && config["inner_taper_length"] == 0
        hole   = Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"]/2,
                      "tube" => Dict(
                      "r"    => Dict("from" => config["hole_radius"],
                                       "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => config["hole_length"])))
        inner_hole_bottom =  Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"],
                      "tube" => Dict(
                      "r"    => Dict("from" => 0.0,
                                       "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0,
                                       "to" => 360.0),
                      "h"    => 0.0)))
        push!(geometry_2, hole)
        push!(geometry_2, inner_hole_bottom)
    end
    if config["hole_radius"] != 0 && config["hole_length"] != 0 && config["inner_taper_width"] != 0 && config["inner_taper_length"] != 0
        inner_taper = Dict("translate" => Dict("z" => config["xtal_length"] - config["inner_taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["hole_radius"] + config["inner_taper_width"],
                                                      "to"   => config["hole_radius"] + config["inner_taper_width"]),
                                     "bottom" => Dict("from" => config["hole_radius"],
                                                      "to"   => config["hole_radius"])),
                      "phi"  => Dict("from" => 0.0, "to" => 360.0),
                      "h"    => config["inner_taper_length"])))
        inner_hole_bottom = Dict("translate" => Dict("z" => config["xtal_length"] - config["hole_length"],
                      "tube" => Dict(
                      "r"    => Dict("from" => 0.0, "to" => config["hole_radius"]),
                      "phi"  => Dict("from" => 0.0, "to" => 360.0),
                      "h"    => 0.0)))
        push!(geometry_2, inner_taper)
        push!(geometry_2, inner_hole_bottom)
    end
    if config["outer_taper_width"] != 0 && config["outer_taper_length"] != 0
        outer_taper = Dict("translate" => Dict("z" => config["xtal_length"] - config["outer_taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["xtal_radius"] - config["outer_taper_width"],
                                                      "to"   => config["xtal_radius"] - config["outer_taper_width"]),
                                     "bottom" => Dict("from" => config["xtal_radius"],
                                                      "to"   => config["xtal_radius"])),
                      "phi"  => Dict("from"   => 0.0, "to" => 360.0),
                      "h"    => config["outer_taper_length"])))
        push!(geometry_2, outer_taper)
    end
    if config["taper_length"] != 0
        taper = Dict("translate" => Dict("z" => config["taper_length"]/2,
                      "cone" => Dict(
                      "r"    => Dict("top"    => Dict("from" => config["xtal_radius"],
                                                      "to"   => config["xtal_radius"]),
                                     "bottom" => Dict("from" => config["xtal_radius"] - config["taper_length"],
                                                      "to"   => config["xtal_radius"] - config["taper_length"])),
                      "phi"   => Dict("from"   => 0.0, "to" => 360.0),
                      "h"     => config["taper_length"])))
        push!(geometry_2, taper)
    end

    #>----------------Define contact 2-----------------------------------------<

    contact_2 = Dict("material"    => "HPGe",
                     "id"          => 2,
                     "potential"   => config["xtal_HV"],
                     "geometry"    => Dict("union" => geometry_2));

    #>----------------Contacts END---------------------------------------------<
    #<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><
    #>----------------Final configuration--------------------------------------<

    impurity_z0            = config["impurity_z0"]*1e10
    impurity_gradient      = config["impurity_gradient"]*1e10
    if units["length"]     == "mm"
        impurity_z0        /= 1e3
        impurity_gradient  /= 1e4
    elseif units["length"] == "m"
        impurity_z0        *= 1e6
        impurity_gradient  *= 1e8
    end

    objects = Dict("semiconductor" => Dict(
                    "material"    => "HPGe",
                    "temperature" => config["xtal_temp"],
                    "impurity_density" => Dict(
                                "name"     => "cylindrical",
                                "r"        => Dict("init"     => 0.0,
                                                   "gradient" => 0.0),
                                "z"        => Dict("init"     => impurity_z0,
                                                   "gradient" => impurity_gradient)),
                    "geometry"    => geometry),
                    "contacts" => [ contact_1, contact_2 ]);

    conf = Dict("name"    => config["name"],
                "units"   => units,
                "grid"    => grid,
                "medium"  => "vacuum",
                "detectors" => [objects]);

    return conf
end