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

function assure_ordering!(ring, exps, coeffs, metainfo)
    if ring.ord != metainfo.computeord
        sort_input_to_change_ordering(exps, coeffs, metainfo.computeord)
    end
    ring.ord = metainfo.computeord
end

# TODO: change project structure

function input_statistics(exps)
    sz = length(exps)
    deg = maximum(sum(e[1]) for e in exps)
    sz, deg
end

function select_tablesize(ring, exps)
    nvars = ring.nvars
    sz = length(exps)

    tablesize = 2^10
    if nvars > 4
        tablesize = 2^14
    end
    if nvars > 7
        tablesize = 2^16
    end

    if sz < 3
        tablesize = div(tablesize, 2)
    end
    if sz < 2
        tablesize = div(tablesize, 2)
    end

    tablesize
end

function cleanup_gens!(ring, gens_ff, prime)
    gens_ff.ch = prime
    ring.ch = prime
    normalize_basis!(gens_ff)
end

# Mutate everything!

#------------------------------------------------------------------------------
# Finite field groebner

function groebner_ff(
            ring::PolyRing,
            exps::Vector{Vector{Vector{UInt16}}},
            coeffs::Vector{Vector{UInt64}},
            reduced::Bool,
            meta::GroebnerMetainfo) where {Rng<:Random.AbstractRNG}

    # select hashtable size
    tablesize = select_tablesize(ring, exps)

    basis, ht = initialize_structures(
                        ring, exps, coeffs, meta.rng, tablesize)

    f4!(ring, basis, ht, reduced, meta.linalg, meta.rng)

    gbexps = hash_to_exponents(basis, ht)
    gbexps, basis.coeffs
end

#------------------------------------------------------------------------------
# Rational field groebner

function groebner_qq(
            ring::PolyRing,
            exps::Vector{Vector{ExponentVector}},
            coeffs::Vector{Vector{CoeffQQ}},
            reduced::Bool,
            meta::GroebnerMetainfo) where {Rng<:Random.AbstractRNG}

    # we can mutate coeffs and exps here

    # select hashtable size
    tablesize = select_tablesize(ring, exps)
    @info "Selected tablesize $tablesize"

    # initialize hashtable and finite field generators structs
    gens_temp_ff, ht = initialize_structures_ff(ring, exps,
                                        coeffs, meta.rng, tablesize)
    gens_ff = copy_basis_thorough(gens_temp_ff)

    # now hashtable is filled correctly,
    # and gens_temp_ff exponents are correct and in correct order.
    # gens_temp_ff coefficients are filled with random stuff and
    # gens_temp_ff.ch is 0

    # to store integer and rational coefficients of groebner basis
    coeffaccum  = CoeffAccum()
    # to store BigInt buffers and reduce overall memory usage
    coeffbuffer = CoeffBuffer()

    # scale coefficients of input to integers
    coeffs_zz = scale_denominators(coeffbuffer, coeffs)

    # keeps track of used prime numbers
    primetracker = PrimeTracker(coeffs_zz)

    #=
        exps, coeffs and coeffs_zz must be *not changed* during whole computation
    =#

    i = 1

    # copy basis so that we initial exponents dont get lost
    gens_ff = copy_basis_thorough(gens_temp_ff)

    prime = nextluckyprime!(primetracker)
    @info "$i: selected lucky prime $prime"

    # perform reduction and store result in gens_ff
    reduce_modulo!(coeffbuffer, coeffs_zz, gens_ff.coeffs, prime)

    # do some things to ensure generators are correct
    cleanup_gens!(ring, gens_ff, prime)

    # compute groebner basis in finite field
    #=
    Need to make sure input invariants in f4! are satisfied, f4.jl for details
    =#

    tracer = Tracer()

    pairset = initialize_pairset()

    # ADDED
    global F4TIME
    # ADDED
    F4TIME += @elapsed f4trace!(ring, gens_ff, tracer, pairset, ht, reduced, meta.linalg, meta.rng)
    # f4!(ring, gens_ff, ht, reduced, meta.linalg)

    # reconstruct into integers
    @info "CRT modulo ($(primetracker.modulo), $(prime))"

    # ADDED
    global RECTIME
    # ADDED
    RECTIME += @elapsed reconstruct_crt!(coeffbuffer, coeffaccum, primetracker, gens_ff.coeffs, prime)
    # reconstruct_crt!(coeffbuffer, coeffaccum, primetracker, gens_ff.coeffs, prime)

    # reconstruct into rationals
    @info "Reconstructing modulo $(primetracker.modulo)"
    # ADDED
    RECTIME += @elapsed reconstruct_modulo!(coeffbuffer, coeffaccum, primetracker)
    # reconstruct_modulo!(coeffbuffer, coeffaccum, primetracker)

    correct = false
    # ADDED
    global CORRTIME
    t = time()
    if correctness_check!(coeffaccum, coeffbuffer, primetracker, meta,
                            ring, exps, coeffs, coeffs_zz, gens_temp_ff, gens_ff, ht)
        @info "Success!"
        CORRTIME += time() - t
        correct = true
    end
    CORRTIME += time() - t

    gap = 1
    primegaps = (1, 1, 1, 1, 1)
    primemult = 2
    # if first prime was not successfull
    while !correct

        if i <= length(primegaps)
            gap = primegaps[i]
        else
            gap = gap*primemult
        end

        for j in 1:gap
            i += 1

            # copy basis so that we initial exponents dont get lost
            gens_ff = copy_basis_thorough(gens_temp_ff)
            prime = nextluckyprime!(primetracker)
            @info "$i: selected lucky prime $prime"
            # perform reduction and store result in gens_ff
            reduce_modulo!(coeffbuffer, coeffs_zz, gens_ff.coeffs, prime)
            # do some things to ensure generators are correct
            cleanup_gens!(ring, gens_ff, prime)
            # compute groebner basis in finite field
            #=
            Need to make sure input invariants in f4! are satisfied, f4.jl for details
            =#
            # ADDED
            global F4TIME
            # ADDED

            # @error "tracer" tracer

            F4TIME += @elapsed f4trace!(ring, gens_ff, tracer, pairset, ht, reduced, meta.linalg, meta.rng)
            # f4!(ring, gens_ff, ht, reduced, meta.linalg)
            # reconstruct into integers
            @info "CRT modulo ($(primetracker.modulo), $(prime))"
            # ADDED
            global RECTIME
            # ADDED

            if meta.linalg === :prob
                # check if probabilistic linear algebra failed
                shape_control(coeffaccum.gb_coeffs_qq, gens_ff.coeffs)
            end

            # ADDED
            # prevcoeffs = deepcopy(coeffaccum.gb_coeffs_zz)

            RECTIME += @elapsed reconstruct_crt!(coeffbuffer, coeffaccum,
                                        primetracker, gens_ff.coeffs, prime)
            # reconstruct_crt!(coeffbuffer, coeffaccum, primetracker, gens_ff.coeffs, prime)

            #=
            -----------------------------------
            Given r1, r2, and m1, m2, so that
                r1 ≡ 1/b1 (mod m1)
                r2 ≡ 1/b2 (mod m2)
            Determine if b1 = b2
            -----------------------------------

            Easy:
            b1 is inverse of r1 mod m1;
            b2 is inverse of r2 mod m2;

            then check equality
            =#

            #=
            -----------------------------------
            We go further.
            Given r1, r2, and m1, m2, so that
                r1 ≡ a1/b1 (mod m1)
                r2 ≡ a2/b2 (mod m2)
            Determine if a1 = a2 and b1 = b2
            -----------------------------------
            r1*b1 ≡ a1 (mod m1),
            r2*b2 ≡ a2 (mod m2),

            r1*b1 * r2*b2 = a1*a2 (mod m1*m2)
            =#

            #=
            if j != gap
                if crt_issame_check!(coeffbuffer, coeffaccum,
                                gens_ff.coeffs, primetracker.modulo, prime)
                    @info "CRT same check passed!"
                    break
                end
            end
            =#
        end

        # reconstruct into rationals
        @info "Reconstructing modulo $(primetracker.modulo)"
        # ADDED
        stats = @timed reconstruct_modulo!(coeffbuffer, coeffaccum, primetracker)
        RECTIME += stats.time

        # ADDED TODO
        if !stats.value
            continue
        end

        # reconstruct_modulo!(coeffbuffer, coeffaccum, primetracker)
        # ADDED
        global CORRTIME
        t = time()
        if correctness_check!(coeffaccum, coeffbuffer, primetracker, meta,
                                ring, exps, coeffs, coeffs_zz, gens_temp_ff, gens_ff, ht)
            @info "Success!"
            CORRTIME += time() - t
            correct = true
        end
        CORRTIME += time() - t
        #=
        if correctness_check!(coeffaccum, coeffbuffer, primetracker, meta,
                                ring, exps, coeffs, coeffs_zz, gens_temp_ff, gens_ff, ht)
            @info "Success!"
            break
        end
        =#

        if meta.linalg === :prob
            if i > 1024
                coeff_control(primetracker, coeffaccum)
            end
        end
    end

    #=
    if meta.usefglm
        fglm_f4!(ring, gens_ff, ht)
    end
    =#

    # ADDED
    global PRIMES
    PRIMES = i

    # normalize_coeffs!(gbcoeffs_qq)
    gb_exps = hash_to_exponents(gens_ff, ht)
    gb_exps, coeffaccum.gb_coeffs_qq
end

#######################################
# Finite field isgroebner

# UWU!
function isgroebner_ff(
            ring::PolyRing,
            exps::Vector{Vector{ExponentVector}},
            coeffs::Vector{Vector{CoeffFF}},
            meta)

    isgroebner_f4(ring, exps, coeffs, meta.rng)
end

#######################################
# Rational field groebner

# TODO
function isgroebner_qq(
            ring::PolyRing,
            exps::Vector{Vector{ExponentVector}},
            coeffs::Vector{Vector{CoeffQQ}},
            meta)

    # if randomized result is ok
    if !meta.guaranteedcheck
        coeffs_zz = scale_denominators(coeffs)
        goodprime = nextgoodprime(coeffs_zz, Int[], 2^30 + 3)
        ring_ff, coeffs_ff = reduce_modulo(coeffs_zz, ring, goodprime)
        isgroebner_f4(ring_ff, exps, coeffs_ff, meta.rng)
    else # if proved result is needed, compute in rationals
        isgroebner_f4(ring, exps, coeffs, meta.rng)
    end
end


# 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 an L-shaped optimizer.
...
# Arguments
- `master_optimizer::AbstractOptimizer`: MathOptInterface solver capable of solving linear (and possibly quadratic) programs.
- `subproblem_optimizer::AbstractOptimizer`: Optionally specify a different solver for the subproblems.
- `feasibility_cuts::Bool = false`: Specify if feasibility cuts should be used
- `regularize::AbstractRegularizer = DontRegularize()`: Specify regularization procedure (DontRegularize, RegularizedDecomposition/RD/WithRegularizedDecomposition, TrustRegion/TR/WithTrustRegion, LevelSet/LV/WithLevelSets).
- `aggregate::AbstractAggregator = DontAggregate()`: Specify aggregation procedure (DontAggregate, Aggregate, PartialAggregate, DynamicAggregate, ClusterAggregate, GranulatedAggregate, HybridAggregate)
- `consolidate::AbstractConsolidator = DontConsolidate()`: Specify consolidation procedure (DontConsolidate, Consolidate)
- `execution::Execution = Serial`: Specify how algorithm should be executed (Serial, Synchronous, Asynchronous). Distributed variants requires worker cores.
- <keyword arguments>: Algorithm specific parameters, See `?LShaped` for list of possible arguments and default values.
...
"""
mutable struct Optimizer <: AbstractStructuredOptimizer
    master_optimizer
    subproblem_optimizer
    master_params::Dict{MOI.AbstractOptimizerAttribute, Any}
    sub_params::Dict{MOI.AbstractOptimizerAttribute, Any}
    feasibility_strategy::AbstractFeasibilityStrategy
    integer_strategy::AbstractIntegerStrategy
    execution::AbstractExecution
    regularizer::AbstractRegularizer
    aggregator::AbstractAggregator
    consolidator::AbstractConsolidator
    parameters::LShapedParameters{Float64}

    status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    dual_status::MOI.ResultStatusCode
    raw_status::String
    solve_time::Float64

    lshaped::Union{AbstractLShaped, Nothing}

    function Optimizer(; master_optimizer = nothing,
                       execution::AbstractExecution = nworkers() == 1 ? Serial() : Synchronous(),
                       feasibility_strategy::AbstractFeasibilityStrategy = IgnoreFeasibility(),
                       integer_strategy::AbstractIntegerStrategy = IgnoreIntegers(),
                       regularize::AbstractRegularizer = DontRegularize(),
                       aggregate::AbstractAggregator = DontAggregate(),
                       consolidate::AbstractConsolidator = DontConsolidate(),
                       subproblem_optimizer = nothing, kw...)
        return new(master_optimizer,
                   subproblem_optimizer,
                   Dict{MOI.AbstractOptimizerAttribute, Any}(),
                   Dict{MOI.AbstractOptimizerAttribute, Any}(),
                   feasibility_strategy,
                   integer_strategy,
                   execution,
                   regularize,
                   aggregate,
                   consolidate,
                   LShapedParameters{Float64}(; kw...),
                   MOI.OPTIMIZE_NOT_CALLED,
                   MOI.NO_SOLUTION,
                   MOI.NO_SOLUTION,
                   "L-shaped optimizer has not been run.",
                   NaN,
                   nothing)
    end
end

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

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

function check_loadable(optimizer::Optimizer, ::StageDecompositionStructure)
    if optimizer.master_optimizer === nothing
        msg = "Master optimizer not set. Consider setting `MasterOptimizer` attribute."
        throw(UnloadableStructure{Optimizer, StageDecompositionStructure}(msg))
    end
    return nothing
end

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

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

function load_structure!(optimizer::Optimizer, structure::StageDecompositionStructure, 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 master optimizer
    set_master_optimizer!(structure, optimizer.master_optimizer)
    # Default subproblem optimizer to master optimizer if
    # none have been set
    if optimizer.subproblem_optimizer === nothing
        optimizer.subproblem_optimizer = optimizer.master_optimizer
    end
    # Set subproblem optimizers
    set_subproblem_optimizer!(structure, optimizer.subproblem_optimizer)
    # Supply the subproblem optimizer if the integer strategy requires it and none have been set
    if optimizer.integer_strategy isa Convexification && optimizer.integer_strategy.parameters.optimizer === nothing
        optimizer.integer_strategy.parameters.optimizer = MOI.OptimizerWithAttributes(optimizer.subproblem_optimizer, collect(optimizer.sub_params))
    end
    # Set any given master/sub optimizer attributes
    for (attr, value) in optimizer.master_params
        MOI.set(backend(structure.first_stage), attr, value)
    end
    for (attr, value) in optimizer.sub_params
        MOI.set(scenarioproblems(structure), attr, value)
    end
    # Create new L-shaped algorithm
    optimizer.lshaped = LShapedAlgorithm(structure,
                                         x₀,
                                         optimizer.feasibility_strategy,
                                         optimizer.integer_strategy,
                                         optimizer.execution,
                                         optimizer.regularizer,
                                         optimizer.aggregator,
                                         optimizer.consolidator;
                                         type2dict(optimizer.parameters)...)
    return nothing
end

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

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

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

function optimizer_name(optimizer::Optimizer)
    optimizer_str = "$(str(optimizer.execution))$(str(optimizer.regularizer))"
    aggregate_str = str(optimizer.aggregator)
    if aggregate_str != ""
        return string(optimizer_str, " with ", aggregate_str)
    else
        return optimizer_str
    end
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.master_params[attr] = 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(LShapedParameters))
        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(LShapedParameters))
        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, ::RelativeTolerance)
    return MOI.get(optimizer, MOI.RawParameter("τ"))
end

function MOI.set(optimizer::Optimizer, ::RelativeTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("τ"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MasterOptimizer)
    if optimizer.master_optimizer === nothing
        return nothing
    end
    return MOI.OptimizerWithAttributes(optimizer.master_optimizer, collect(optimizer.master_params))
end

function MOI.set(optimizer::Optimizer, ::MasterOptimizer, optimizer_constructor)
    optimizer.master_optimizer = optimizer_constructor
    # Clear any old parameters
    empty!(optimizer.master_params)
    return nothing
end

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

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

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

function MOI.set(optimizer::Optimizer, param::RawMasterOptimizerParameter, value)
    moi_param = MOI.RawParameter(param.name)
    optimizer.master_params[moi_param] = value
    return nothing
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, ::FeasibilityStrategy)
    return optimizer.feasibility_strategy
end

function MOI.set(optimizer::Optimizer, ::FeasibilityStrategy, strategy::AbstractFeasibilityStrategy)
    optimizer.feasibility_strategy = strategy
    return nothing
end

function MOI.get(optimizer::Optimizer, ::IntegerStrategy)
    return optimizer.integer_strategy
end

function MOI.set(optimizer::Optimizer, ::IntegerStrategy, strategy::AbstractIntegerStrategy)
    optimizer.integer_strategy = strategy
    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, ::Regularizer)
    return optimizer.regularizer
end

function MOI.set(optimizer::Optimizer, ::Regularizer, regularizer::AbstractRegularizer)
    optimizer.regularizer = regularizer
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Aggregator)
    return optimizer.aggregator
end

function MOI.set(optimizer::Optimizer, ::Aggregator, aggregator::AbstractAggregator)
    optimizer.aggregator = aggregator
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Consolidator)
    return optimizer.consolidator
end

function MOI.set(optimizer::Optimizer, ::Consolidator, consolidator::AbstractConsolidator)
    optimizer.consolidator = consolidator
    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::RegularizationParameter)
    return MOI.get(optimizer.regularizer, param)
end

function MOI.set(optimizer::Optimizer, param::RegularizationParameter, value)
    MOI.set(optimizer.regularizer, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, param::AggregationParameter)
    return MOI.get(optimizer.aggregator, param)
end

function MOI.set(optimizer::Optimizer, param::AggregationParameter, value)
    MOI.set(optimizer.aggregator, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, param::ConsolidationParameter)
    return MOI.get(optimizer.consolidator, param)
end

function MOI.set(optimizer::Optimizer, param::ConsolidationParameter, value)
    MOI.set(optimizer.consolidator, 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, attr::MOI.ListOfVariableIndices)
    if optimizer.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    list = MOI.get(optimizer.lshaped.structure, attr)
    # Remove the master variables
    filter!(vi -> !(vi in optimizer.lshaped.master_variables), list)
    # Remove any auxiliary variables from regularization
    filter_variables!(optimizer.lshaped.regularization, list)
    return list
end

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

function MOI.get(optimizer::Optimizer, attr::MOI.ListOfConstraintIndices)
    if optimizer.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    list = MOI.get(optimizer.lshaped.structure, attr)
    # Remove any cut constraints
    filter_cuts!(optimizer.lshaped, list)
    # Remove any penaltyterm constraints
    filter_constraints!(optimizer.lshaped.regularization, list)
    return list
end

function MOI.get(optimizer::Optimizer, ::MOI.ObjectiveValue)
    if optimizer.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return objective_value(optimizer.lshaped)
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 first-stage optimizer through structure
    if optimizer.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.lshaped.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to first-stage optimizer through structure
    if optimizer.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.lshaped.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.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.lshaped.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.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.lshaped.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.lshaped === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.lshaped.structure, attr, ci, value)
end

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

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

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

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

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

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

function MOI.is_empty(optimizer::Optimizer)
    return optimizer.lshaped === 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, ::RawInstanceOptimizerParameter) = true
MOI.supports(::Optimizer, ::AbstractLShapedAttribute) = true

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

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

See also: [`set_regularization_attribute`](@ref), [`set_regularization_attributes`](@ref).
"""
function get_regularization_attribute(stochasticprogram::StochasticProgram, name::String)
    return return MOI.get(optimizer(stochasticprogram), RawRegularizationParameter(name))
end
"""
    set_regularization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

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

"""
function set_regularization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawRegularizationParameter(String(name)), value)
end
"""
    set_regularization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

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

"""
function set_regularization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_regularization_attribute(stochasticprogram, name, value)
    end
end
function set_regularization_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_regularization_attribute(stochasticprogram, name, value)
    end
end
"""
    get_aggregation_attribute(stochasticprogram::StochasticProgram, name::String)

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

See also: [`set_aggregation_attribute`](@ref), [`set_aggregation_attributes`](@ref).
"""
function get_aggregation_attribute(stochasticprogram::StochasticProgram, name::String)
    return MOI.get(optimizer(stochasticprogram), RawAggregationParameter(name))
end
"""
    set_aggregation_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

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

"""
function set_aggregation_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawAggregationParameter(String(name)), value)
end
"""
    set_aggregation_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

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

"""
function set_aggregation_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_aggregation_attribute(stochasticprogram, name, value)
    end
end
function set_aggregation_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_aggregation_attribute(stochasticprogram, name, value)
    end
end
"""
    get_consolidation_attribute(stochasticprogram::StochasticProgram, name::String)

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

See also: [`set_consolidation_attribute`](@ref), [`set_consolidation_attributes`](@ref).
"""
function get_consolidation_attribute(stochasticprogram::StochasticProgram, name::String)
    return return MOI.get(optimizer(stochasticprogram), RawConsolidationParameter(name))
end
"""
    set_consolidation_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

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

"""
function set_consolidation_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawConsolidationParameter(String(name)), value)
end
"""
    set_consolidation_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

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

"""
function set_consolidation_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_consolidation_attribute(stochasticprogram, name, value)
    end
end
function set_consolidation_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_consolidation_attribute(stochasticprogram, name, value)
    end
end


using DifferentiableStateSpaceModels,
    SparseArrays,
    LinearAlgebra,
    Parameters,
    Test,
    TimerOutputs,
    BenchmarkTools

@testset "Dense SW First Order" begin
    m = @include_example_module(Examples.SW07)

    p = [
        10,
        0.51,
        10,
        0,
        0.7,
        0.742,
        0,
        0,
        0.24,
        0.2696,
        6.0144,
        0.025,
        1.5,
        0.6361,
        1.5,
        0.3243,
        0.8087,
        0.47,
        0.6,
        1.9423,
        1.5,
        1.488,
        0.2347,
        0.0593,
        0.8762,
        0.9977,
        0.5799,
        0.9957,
        0.7165,
        0,
        0,
        0,
        0.3982,
        0.18,
    ]
    generate_perturbation(m, p)
    reset_timer!()
    sol = generate_perturbation(m, p)
    print_timer()
    @test sol.retcode == :Success
end


using RecipesBase

################################################################################
# Destriping results

@doc raw"""
    mutable struct DestripingData{T <: Real, O <: Healpix.Order}

Structure containing the result of a destriping operation. It contains
the following fields:

- `skymap`: a `PolarizedMap` containing the maximum-likelihood map
- `hitmap`: a map containing the weights of each pixel (this is not
  the hit count, as it is normalized over the value of ``σ^2`` for
  each sample in the TOD)
- `nobs_matrix`: an array of [`NobsMatrixelement`](@ref) objects,
  which can be used to determine which pixels were the most
  troublesome in the reconstruction of tehe I/Q/U components
- `baselines`: the computed baselines. These must be kept as a list of
  `RunLengthArray` objects
- `stopping_factors`: sequence of stopping factors, one for
  each iteration of the CG algorithm
- `threshold`: upper limit on the tolerable error for the Conjugated
  Gradient method (optional, default is `1e-9`). See
  [`calc_stopping_factor`](@ref).
- `max_iterations`: maximum number of allowed iterations for the
  Conjugated Gradient algorithm (optional, default is `1000`).
- `use_preconditioner`: a Boolean flag telling whether a
  preconditioner should be used or not (optional, default is
  `true`). Usually you want this to be `true`.

The structure provides only one constructor, which accepts the
following parameters:

- `nside`: resolution of `skymap` and `hitmap`

- `obs_list`: vector of [`Observation`](@ref); it is used to determine
  which pixels in the map are going to be observed

- `runs`: list of vectors (one for each element in `obs_list`),
  containing the number of samples to be included in each baseline for
  each observation

The constructor accepts the following keywords as well:

- `threshold` (default: `1e-9`)

- `max_iterations` (default: 1000)

- `use_preconditioner` (default: `true`)

The function [`destripe!`](@ref) use `DestripingData` to return the
result of its computation.

"""
mutable struct DestripingData{T <: Real, O <: Healpix.Order}
    skymap::Healpix.PolarizedMap{T, O}
    hitmap::Healpix.Map{T, O}
    nobs_matrix::Vector{NobsMatrixElement{T}}
    baselines::Vector{RunLengthArray{Int, T}}
    stopping_factors::Vector{T}
    threshold::Float64
    max_iterations::Int
    use_preconditioner::Bool

    function DestripingData{T,O}(
        nside,
        obs_list::Vector{Observation{T}},
        runs;
        threshold = 1e-9,
        max_iterations = 1000,
        use_preconditioner = true,
        comm = nothing,
    ) where {T <: Real, O <: Healpix.Order}

        @assert length(obs_list) == length(runs)
        @assert threshold > 0.0
        @assert max_iterations > 0

        npix = Healpix.nside2npix(nside)
        nobs_matrix = [NobsMatrixElement{T}() for i in 1:npix]
        compute_nobs_matrix!(nobs_matrix, obs_list, comm = comm)

        new(
            Healpix.PolarizedMap{T, O}(
                zeros(npix),
                zeros(npix),
                zeros(npix),
            ),
            Healpix.Map{T,O}(nside),
            nobs_matrix,
            [RunLengthArray{Int,T}(x, zeros(length(x))) for x in runs],
            T[],  # List of stopping factors
            threshold,
            max_iterations,
            use_preconditioner,
        )
    end

end

function Base.show(io::IO, d::DestripingData)
    print(io, "DestripingData(NSIDE=$(d.hitmap.resolution.nside))")
end

function Base.show(io::IO, ::MIME"text/plain", d::DestripingData)
    # Count how many pixels have been observed in the map
    obspix = 0
    for elem in d.nobs_matrix
        elem.neglected || (obspix += 1)
    end

    # Put all the baselines of each observation into one long list
    # (without allocating memory!)
    flattened_baselines = Iterators.flatten([values(x) for x in d.baselines])

    println(
        io,
        @sprintf(
            """Destriping data:
NSIDE of the map........................ %d
Number of observed pixels............... %d (%.1f%%)
Number of iterations.................... %d/%d
Best convergence factor................. %s (threshold: %e)
Number of observations.................. %d
Average value of the baselines.......... %e
RMS of the baselines.................... %e
Preconditioner.......................... %s
""",
            d.hitmap.resolution.nside,
            obspix,
            100.0 * obspix / length(d.nobs_matrix),
            length(d.stopping_factors),
            d.max_iterations,
            isempty(d.stopping_factors) ? "none" : @sprintf("%e", minimum(d.stopping_factors)),
            d.threshold,
            length(d.baselines),
            mean(flattened_baselines),
            std(flattened_baselines),
            "$(d.use_preconditioner)",
        )
    )
end

@doc raw"""
    reset_maps!(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}

Set the skymap and the hitmap in `d` to zero. Nothing is done on the
other fields.

"""
function reset_maps!(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}
    d.skymap.i .= 0.0
    d.skymap.q .= 0.0
    d.skymap.u .= 0.0

    d.hitmap .= 0.0
end

@recipe function plot_destriping_data(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}
    seriestype := :path
    xlabel := "Iteration number"
    ylabel := "Stopping factor"
    linecolor := :black

    isempty(d.stopping_factors) && return ([0.0], [0.0])
    num_of_steps = length(d.stopping_factors)

    @series begin
        # Horizontal line showing the threshold
        seriestype := :path
        linecolor := :red
        ([1, num_of_steps], [d.threshold, d.threshold])
    end

    ((1:num_of_steps), d.stopping_factors)
end

@doc raw"""
    plot(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}

Plot the convergence plot for a destriping solution.
"""
plot_destriping_data


using Pope, HDF5
using Pope.LJH
using Test
using ReferenceMicrocalFiles

@testset "BufferedHDF5Dataset2D" begin
    for name in ["ds","1/ds"]
        h5 = h5open(tempname(),"w")
        chunksize = 100
        # create an extendible dataset to hold reduced pulses with 6 element basis
        buffer = Pope.BufferedHDF5Dataset2D{Float32}(h5,name,6,1000)
        write(buffer, Vector{Float32}(1:6))
        write(buffer, Vector{Float32}(7:12))
        write(buffer, [Vector{Float32}(13:18), Vector{Float32}(19:24)])
        Pope.write_to_hdf5(buffer)
        @test read(buffer.ds) == reshape(Vector{Float32}(1:24),(6,4))
        close(h5)
        rm(h5.filename)
    end
end

@testset "BasisAnalyzer" begin
    FrameTime, PretrigNSamples, NRow, row = 9.6e-6, 100, 30, 1
    # rowcount of LJH2 files = framecount*NRow+row, framecount=frame1index
    r = LJH.LJHRecord{FrameTime, PretrigNSamples, NRow}(1:1000,10*NRow+row,20)
    r3 = LJH.LJH3Record{FrameTime}(1:1000,PretrigNSamples,10,20)
    projectors = rand(6,1000)
    analyzer = Pope.BasisAnalyzer(projectors, projectors')
    dp = analyzer(r)
    dp3 = analyzer(r3)
    @test Pope.nbases(analyzer) == 6
    @test Pope.nsamples(analyzer) == 1000
    @test dp.reduced == dp3.reduced
    @test dp.residual_std == dp3.residual_std
    @test dp.timestamp_usec == dp3.timestamp_usec == LJH.timestamp_usec(r)
    @test dp.first_rising_sample == dp.first_rising_sample == PretrigNSamples
    @test dp.nsamples == dp3.nsamples == length(r)
end

@testset "BasisBufferedWriter" begin
    h5 = h5open(tempname(),"w")
    g = HDF5.g_create(h5,"1")
    b = Pope.BasisBufferedWriter(g, 6, 1000, 0.001, start=true)
    dataproduct = Pope.BasisDataProduct(collect(1:6),1,2,3,4,5)
    write(b,dataproduct)
    Pope.stop(b)
    fetch(b)
    @test read(h5["1/reduced"]) == reshape(1:6,(6,1))
    @test read(h5["1/residual_std"]) == [dataproduct.residual_std]
    @test read(h5["1/frame1index"]) == [dataproduct.frame1index]
    @test read(h5["1/timestamp_usec"]) == [dataproduct.timestamp_usec]
    @test read(h5["1/first_rising_sample"]) == [dataproduct.first_rising_sample]
    @test read(h5["1/nsamples"]) == [dataproduct.nsamples]
    close(h5)
    rm(h5.filename)
end


## Loading modules
using Distributed
using DataFrames
using CSV
addprocs(6)

## Settin up directories

@everywhere dir = @__DIR__ ;
datadir=dir*"/data/oildata.csv" ;
resultsdir=dir*"/results/"

## Load modules in parallel
@everywhere include(dir*"/modules/MyModules.jl")

@everywhere begin
    using .MyModules
    using Random
    using SharedArrays
    ## Setting seed
    Random.seed!(66636)
end

@everywhere Param=Par(nfast=1000)

## Big 6

pt, qt = load_data_select(datadir)
set="Big6"

# Run tests

# Run base case
include("Main.jl")

# A) Test alternative transition kernel

include("RP_Cour_ME_sem_algo.jl")

# B) Test alternative acceptance rule with current
# transition kernel

# include("RP_Cour_ME_sem_alt.jl")

# C) Annual test

include("Main_annual.jl")

## OPEC

OPEC=[:saudiarabia, :iran, :venezuela, :uae, :nigeria, :kuwait]
pt, qt = load_data_select(datadir, select=OPEC)
set="OPEC"

# Run tests

# Run base case
include("Main.jl")
GC.gc()
# A) Test alternative transition kernel
include("RP_Cour_ME_sem_algo.jl")
GC.gc()
# C) Annual test
include("Main_annual.jl")
GC.gc()
## Non-OPEC
NonOPEC=[:unitedstates, :china, :mexico, :uk, :canada, :norway]
pt, qt = load_data_select(datadir, select=NonOPEC)
set="NonOPEC"

# Run tests

# Run base case
include("Main.jl")
GC.gc()
# A) Test alternative transition kernel
include("RP_Cour_ME_sem_algo.jl")
GC.gc()
# C) Annual test
include("Main_annual.jl")
GC.gc()
## OPEC no Saudia Arabia
OPEC_NSA=[:iran, :venezuela, :uae, :nigeria, :kuwait, :iraq]
pt, qt = load_data_select(datadir, select=OPEC_NSA)
set="OPEC_NSA"

# Run tests

# Run base case
include("Main.jl")
GC.gc()
# A) Test alternative transition kernel
include("RP_Cour_ME_sem_algo.jl")
GC.gc()
# C) Annual test
include("Main_annual.jl")
GC.gc()


using ShaderAbstractions, LinearAlgebra
using ShaderAbstractions: VertexArray
using Test

struct Bla <: ShaderAbstractions.AbstractContext end

import GeometryTypes, AbstractPlotting, GeometryBasics

m = GeometryTypes.GLNormalMesh(GeometryTypes.Sphere(GeometryTypes.Point3f0(0), 1f0))

mvao = VertexArray(m)
instances = VertexArray(positions = rand(GeometryBasics.Point{3, Float32}, 100))

x = ShaderAbstractions.InstancedProgram(
    Bla(), "hi",
    mvao,
    instances,

    model = GeometryTypes.Matf4f0(I),
    view = GeometryTypes.Matf4f0(I),
    projection = GeometryTypes.Matf4f0(I),

)

x.program.source |> println


function move_downloads()

    dl_dir = "/Users/mzgubic/Downloads/"
    for fname in readdir(dl_dir) # TODO: change to ~

        # order matters
        if occursin("Income", fname)
            println(fname)
            mv(dl_dir*fname, "data/income/"*fname, force=true)

        elseif occursin("Net", fname)
            println(fname)
            mv(dl_dir*fname, "data/networth/"*fname, force=true)

        elseif occursin("Expenses", fname)
            println(fname)
            mv(dl_dir*fname, "data/expense/"*fname, force=true)

        elseif occursin("exchange", fname)
            println(fname)
            mv(dl_dir*fname, "data/fx/"*fname, force=true)
        end

    end

end

function is_individual(fpath)::Bool
    return !occursin("Repeated", fpath)
end

function step(row)
    if row.frequency == "DAILY"
        return Dates.Day(1)
    elseif row.frequency == "MONTHLY"
        return Dates.Month(1)
    else
        throw("Error")
    end
end

function load_repeated(fpath::String)::DataFrame

    # repeated df
    rep_df = DataFrame(CSV.File(fpath))

    # individual df
    df = DataFrame(date=Date[],
                   amount=Float64[],
                   currency=String[],
                   category=String[],
                   description=String[])

    # loop over repeated and extract individuals
    for row in eachrow(rep_df)
        datestep = step(row)
        daterange = range(row.start, row.end, step=datestep)
        amount = row.amount / length(daterange)
        for date in daterange
            push!(df, [date, amount, row.currency, row.category, row.description])
        end
    end

    return df

end

function load_expenses(fpath::String)::DataFrame

    # treat individual and repeated expenses separately
    println(fpath)
    if is_individual(fpath)
        return DataFrame(CSV.File(fpath))
    else
        return load_repeated(fpath)
    end

end

function load_expenses()::DataFrame
    expense_dir = "data/expense/"
    dfs = [load_expenses(expense_dir*fname) for fname in readdir(expense_dir)]
    reverse!(dfs)
    return vcat(dfs...)
end

function load_incomes()::DataFrame
    income_dir = "data/income/"
    return DataFrame(CSV.File(income_dir*"Expenses - Income.csv"))
end

function load_net()::DataFrame
    net_dir = "data/networth/"
    return DataFrame(CSV.File(net_dir*"Expenses - NetWorth.csv", header=2))
end

function parse_from_to(fpath::String)::Tuple{String, String}

    fname = split(fpath, "/")[end]
    currencies = split(fname, "-exchange-rate")[1]

    first = split(currencies, "-")[1]
    second = split(currencies, "-")[2]

    d = Dict("euro"=>"EUR",
             "british"=>"GBP",
             "dollar"=>"USD")

    return d[first], d[second]
end

function load_fxs(fpath::String)::DataFrame

    # load the raw file
    header = 16 # ignore the first 15 lines in the file
    df = DataFrame(CSV.File(fpath, header=header))
    rename!(df, Dict(Symbol(" value")=>:value))

    # ignore old exchange rates
    old = Date(2015, 01, 01) # make this optional input (supply from expenses/income data)
    df = df[df[:, "date"] .> old, :]

    # add from to info
    from, to = parse_from_to(fpath)
    insertcols!(df, 3, :to_curr=>repeat([to], size(df)[1]))
    insertcols!(df, 3, :from_curr=>repeat([from], size(df)[1]))

    # add a df with reverse exchange rate as well
    df_reverse = copy(df)
    df_reverse[:, "value"] = 1 ./ df[:, "value"]
    rename!(df_reverse, Dict(:to_curr=>:from_curr, :from_curr=>:to_curr))

    return vcat(df, df_reverse)

end

function load_fxs()::DataFrame
    fx_dir = "data/fx/"
    dfs = [load_fxs(fx_dir*fname) for fname in readdir(fx_dir)]
    return vcat(dfs...)
end

function xchange_rate(from_curr::String, to_curr::String, date::Date, fx::DataFrame)::Float64

    # get the subset of rates from the particular currency combination
    xrates = @from i in fx begin
        @where (i.from_curr == from_curr) && (i.to_curr == to_curr)
        @select i
        @collect DataFrame
    end

    # if currency pairs do not exist
    if size(xrates)[1] == 0
        println("WARNING: No exchange rate between $(from_curr) and $(to_curr), using 1.")
        return 1.0
    end

    # try getting an exact rate
    exact_rate = xrates[xrates.date .== date, "value"]
    if length(exact_rate) == 1
        return exact_rate[1]

    # otherwise get an estimate from the closest available day
    else
        days = [abs(Dates.value(d)) for d in xrates.date-date]
        return xrates.value[argmin(days)]
    end

end

function exchange!(exp_df::DataFrame, to_curr::String, fx::DataFrame)::DataFrame

    # loop over all expenses
    for exp in eachrow(exp_df)

        # and convert the spending if needed
        if exp.currency != to_curr

            # determine the xchange rate
            xrate = xchange_rate(exp.currency, to_curr, exp.date, fx)

            # change the amount and the currency
            exp.amount = exp.amount * xrate
            exp.currency = to_curr

        end

    end

    return exp_df
end


include("../imports.jl")

@testset ExtendedTestSet "`Scheduler`" begin
    nbatches = 16
    schedule = Schedule([0, 2], [1e-4, 1e-6], Animations.linear())
    learner = testlearner(coeff = 3, Recorder(), Scheduler(LearningRate => schedule), nbatches = nbatches)

    epoch!(learner, TrainingPhase())
    @test isapprox(last(learner.cbstate.hyperparams[:LearningRate])[2], ((1e-4 + 1e-6) / 2), atol = 0.1)

    epoch!(learner, TrainingPhase())
    @test isapprox(last(learner.cbstate.hyperparams[:LearningRate])[2], 1e-6, atol =0.1)
end


function _randomping()

    ping = true

    function loop()
        count = 1
        t0 = time()
        while ping
            sleep(1)
            msg = "Ping $(count): $(round(time()-t0,digits=2)) seconds"
            e = Pages.Element(id = string(count),innerHTML = msg)
            e.style["color"] = rand(["red","green","blue"])
            rand() < .5 && Pages.append(e)

            count += 1
            tstart = time()
        end
    end

    function start()
        ping = true

        @async loop()
        return
    end

    function stop()
        ping = false
        return
    end

    return (start=start,stop=stop)
end

randomping = _randomping()


function build_arch(name::String,A)
    num_fifos = 2

    f = open(name * ".json")
    s = JSON.parse(f)

    arch = TopLevel{A,2}(s["name"])
    row = s["dimensions"]["row"]
    col = s["dimensions"]["column"]
    processor_add = s["processor_tile"]["add"]
    processor_remove = s["processor_tile"]["remove"]
    one_port = s["memory"]["1_port"]
    two_port = s["memory"]["2_port"]
    num_links = s["configurations"]["num_links"]
    input_handler = s["input_handler"]
    output_handler = s["output_handler"]
    special_mem_connection = s["configurations"]["special_mem_connection"]

    # memories connected in a non-asap way are dealt with below
    offset = Array{Tuple{Int64,Int64},1}()
    if(special_mem_connection)
      tmp = s["configurations"]["offsets"]
      for off in tmp
        push!(offset, stringtotuple(off))
      end
    end

    close(f)

    # preallocate arrays to store Tuple addresses
    in_addr = Array{Tuple{Int64,Int64},1}()
    out_addr = Array{Tuple{Int64,Int64},1}()
    proc = Array{Tuple{Int64,Int64},1}()

    # push the tuples into their respective arrays
    for handler in input_handler
      push!(in_addr,stringtotuple(handler))
    end
    for handler in output_handler
      push!(out_addr,stringtotuple(handler))
    end

    # get all the processors
    # add/remove processors to and from the processor array
    for r in 0:row-1, c in 0:col-1
      match = false
      for rem in processor_remove
        if(rem == "($r,$c)")
          match = true
        end
      end
      match && continue
      push!(proc,(r,c))
    end
    for p in processor_add
      push!(proc,stringtotuple(p))
    end

    one_port_mem = Array{Tuple{Int64,Int64},1}()
    two_port_mem = Array{Tuple{Int64,Int64},1}()
    memory_1port = Array{AsapMapper.MemoryLocation{2},1}()
    memory_2port = Array{AsapMapper.MemoryLocation{2},1}()
    extra_two_port_mem = Array{Tuple{Int64,Int64},1}()

    # get all mems
    count = 1
    for mem in one_port
      addr = stringtotuple(mem)
      push!(one_port_mem,addr)
      if(special_mem_connection)
        push!(memory_1port, MemoryLocation(addr,[offset[count]]))
        count += 1
      else
        push!(memory_1port, MemoryLocation(addr,[(-1,0)]))
      end
    end
    for mem in two_port
      addr = stringtotuple(mem)
      push!(two_port_mem,addr)
      if(special_mem_connection)
        second_address = other_address(addr,offset,count)
        push!(memory_2port, MemoryLocation(addr,second_address,[offset[count],offset[count+1]]))
        push!(extra_two_port_mem,second_address)
        count += 2
      else
        push!(memory_2port, MemoryLocation(addr,[(-1,0),(-1,1)]))
      end
    end
    mem_neighbor_addrs = vcat(mem_neighbor(memory_2port),mem_neighbor(memory_1port))

    ####################
    # Normal Processor #
    ####################
    # Get a processor tile and instantiate it.
    processor = build_processor_tile(num_links)
    for p in proc
      overlap = false
      one_match = false
      two_match = false
      ex_two_match = false
      for mem_neighbor_addr in mem_neighbor_addrs
        if(mem_neighbor_addr == CartesianIndex(p[1],p[2]))
          overlap = true
        end
      end
      for mem in one_port_mem
        if(mem == (p[1],p[2]))
          one_match = true
        end
      end
      for mem in two_port_mem
        if(mem == (p[1],p[2]))
          two_match = true
        end
      end
      for mem in extra_two_port_mem
        if(mem == (p[1],p[2]))
          ex_two_match = true
        end
      end
      (overlap || one_match || two_match || ex_two_match) && continue
      add_child(arch, processor, CartesianIndex(p[1],p[2]))
    end

    ####################
    # Memory Processor #
    ####################
    # Instantiate a memory processor at every address neighboring a memory.
    memory_processor = build_processor_tile(num_links, include_memory = true)
    for mem_neighbor_addr in mem_neighbor_addrs
      add_child(arch, memory_processor, mem_neighbor_addr)
    end

    #################
    # 2 Port Memory #
    #################
    memory_twoport = build_memory(2)
    for mem in two_port_mem
          add_child(arch, memory_twoport, CartesianIndex(mem[1],mem[2]))
    end

    #################
    # 1 Port Memory #
    #################
    memory_oneport = build_memory(2)
    for mem in one_port_mem
          add_child(arch, memory_oneport, CartesianIndex(mem[1],mem[2]))
    end

    #################
    # Input Handler #
    #################
    input_handler = build_input_handler(1)
    for handler in in_addr
      add_child(arch, input_handler, CartesianIndex(handler[1],handler[2]))
    end

    ##################
    # Output Handler #
    ##################
    output_handler = build_output_handler(1)
    for handler in out_addr
      add_child(arch, output_handler, CartesianIndex(handler[1],handler[2]))
    end

    connect_processors(arch, num_links)
    connect_io(arch, num_links)
    if (special_mem_connection)
      if(one_port_mem != nothing)
        connect_memories_cluster(arch, memory_1port)
      end
      if(two_port_mem != nothing)
        connect_memories_cluster(arch, memory_2port)
      end
    else
      connect_memories(arch)
    end

    return arch
end

function stringtotuple(string)

  row,column = split(string,",")
  row = parse(Int64,split(row,"(")[2])
  column = parse(Int64,split(column,")")[1])

  return (row,column)
end

function other_address(addr, offset, count)
  if(offset[count][1] != offset[count+1][1])
    second_addr = (addr[1]+offset[count+1][1],addr[2])
  elseif(offset[count][2] != offset[count+1][2])
    second_addr =  (addr[1],addr[2]+offset[count+1][2])
  end
  return second_addr
end


__precompile__()
module FINUFFT

## Export
export nufft1d1, nufft1d2, nufft1d3
export nufft2d1, nufft2d2, nufft2d3
export nufft3d1, nufft3d2, nufft3d3

export nufft1d1!, nufft1d2!, nufft1d3!
export nufft2d1!, nufft2d2!, nufft2d3!
export nufft3d1!, nufft3d2!, nufft3d3!

export nufftf1d1!, nufftf1d2!, nufftf1d3!
export nufftf2d1!, nufftf2d2!, nufftf2d3!
export nufftf3d1!, nufftf3d2!, nufftf3d3!

export finufft_default_opts
export nufft_opts
export nufft_c_opts # backward-compability

## External dependencies
using finufft_jll

const libfinufft = finufft_jll.libfinufft

const BIGINT = Int64 # defined in src/finufft.h

## FINUFFT opts struct from src/finufft.h
"""
    mutable struct nufft_opts
        modeord             :: Cint
        chkbnds             :: Cint
        debug               :: Cint
        spread_debug        :: Cint
        showwarn            :: Cint
        nthreads            :: Cint
        fftw                :: Cint
        spread_sort         :: Cint
        spread_kerevalmeth  :: Cint
        spread_kerpad       :: Cint
        upsampfac           :: Cdouble
        spread_thread       :: Cint
        maxbatchsize        :: Cint
        spread_nthr_atomic  :: Cint
        spread_max_sp_size  :: Cint
    end

Options struct passed to the FINUFFT library.

# Fields

## Data handling opts

    modeord :: Cint

(type 1,2 only):    0: CMCL-style increasing mode order,
                    1: FFT-style mode order

    chkbnds :: Cint

0: don't check NU pts in [-3pi,3pi),
1: do (<few % slower)

## Diagnostic opts

    debug :: Cint

0: silent,
1: some timing/debug,
2: more

    spread_debug :: Cint

0: silent,
1: some timing/debug,
2: tonnes

    showwarn :: Cint

0: don't print warnings to stderr,
1: do

## Algorithm performance opts

    nthreads :: Cint

number of threads to use, or 0 uses all available

    fftw :: Cint

plan flags to FFTW (`FFTW_ESTIMATE`=64, `FFTW_MEASURE`=0,...)

    spread_sort :: Cint

0: don't sort,
1: do,
2: heuristic choice

    spread_kerevalmeth :: Cint

0: exp(sqrt()) spreading kernel,
1: Horner piecewise poly (faster)

    spread_kerpad :: Cint

option only for exp(sqrt()).
0: don't pad kernel to 4n,
1: do

    upsampfac :: Cdouble

upsampling ratio sigma: 2.0 std, 1.25 small FFT, 0.0 auto

    spread_thread :: Cint

(vectorized ntr>1 only):    0: auto, 1: seq multithreaded,
                            2: parallel single-thread spread

    maxbatchsize :: Cint

option for vectorized ntr>1 only:
max transform batch, 0 auto

    spread_nthr_atomic :: Cint

if >=0, threads above which spreader OMP critical goes atomic

    spread_max_sp_size :: Cint

if >0, overrides spreader (dir=1) max subproblem size

"""
mutable struct nufft_opts
    modeord             :: Cint
    chkbnds             :: Cint
    #
    debug               :: Cint
    spread_debug        :: Cint
    showwarn            :: Cint
    #
    nthreads            :: Cint
    fftw                :: Cint
    spread_sort         :: Cint
    spread_kerevalmeth  :: Cint
    spread_kerpad       :: Cint
    upsampfac           :: Cdouble
    spread_thread       :: Cint
    maxbatchsize        :: Cint
    spread_nthr_atomic  :: Cint
    spread_max_sp_size  :: Cint
end

const nufft_c_opts = nufft_opts # backward compability

"""
    finufft_default_opts()

Return a [`nufft_opts`](@ref) struct with the default FINUFFT settings.\\
See: <https://finufft.readthedocs.io/en/latest/usage.html#options>
"""
function finufft_default_opts()
    opts = nufft_opts(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0)
    ccall( (:finufft_default_opts, libfinufft),
           Nothing,
           (Ref{nufft_opts},),
           opts
           )
    # default to number of julia threads
    opts.nthreads = Threads.nthreads()
    return opts
end

### Error handling
const WARN_EPS_TOO_SMALL            = 1
const ERR_MAXNALLOC                 = 2
const ERR_SPREAD_BOX_SMALL          = 3
const ERR_SPREAD_PTS_OUT_RANGE      = 4
const ERR_SPREAD_ALLOC              = 5
const ERR_SPREAD_DIR                = 6
const ERR_UPSAMPFAC_TOO_SMALL       = 7
const HORNER_WRONG_BETA             = 8
const ERR_NDATA_NOTVALID            = 9
const ERR_TYPE_NOTVALID             = 10
# some generic internal allocation failure...
const ERR_ALLOC                     = 11
const ERR_DIM_NOTVALID              = 12
const ERR_SPREAD_THREAD_NOTVALID    = 13

struct FINUFFTError <: Exception
    errno::Cint
    msg::String
end
Base.showerror(io::IO, e::FINUFFTError) = print(io, "FINUFFT Error ($(e.errno)): ", e.msg)

function check_ret(ret)
    # Check return value and output error messages
    if ret==0
        return
    elseif ret==WARN_EPS_TOO_SMALL
        @warn "requested tolerance epsilon too small to achieve"
        return
    elseif ret==ERR_MAXNALLOC
        msg = "attemped to allocate internal array larger than MAX_NF (defined in defs.h)"
    elseif ret==ERR_SPREAD_BOX_SMALL
        msg = "spreader: fine grid too small compared to spread (kernel) width"
    elseif ret==ERR_SPREAD_PTS_OUT_RANGE
        msg = "spreader: if chkbnds=1, a nonuniform point coordinate is out of input range [-3pi,3pi]^d"
    elseif ret==ERR_SPREAD_ALLOC
        msg = "spreader: array allocation error"
    elseif ret==ERR_SPREAD_DIR
        msg = "spreader: illegal direction (should be 1 or 2)"
    elseif ret==ERR_UPSAMPFAC_TOO_SMALL
        msg = "upsampfac too small (should be >1.0)"
    elseif ret==HORNER_WRONG_BETA
        msg = "upsampfac not a value with known Horner poly eval rule (currently 2.0 or 1.25 only)"
    elseif ret==ERR_NDATA_NOTVALID
        msg = "ntrans not valid in many (vectorized) or guru interface (should be >= 1)"
    elseif ret==ERR_TYPE_NOTVALID
        msg = "transform type invalid"
    elseif ret==ERR_ALLOC
        msg = "general allocation failure"
    elseif ret==ERR_DIM_NOTVALID
        msg = "dimension invalid"
    elseif ret==ERR_SPREAD_THREAD_NOTVALID
        msg = "spread_thread option invalid"
    else
        msg = "unknown error"
    end
    throw(FINUFFTError(ret, msg))
end

### Simple Interfaces (allocate output)

## Type-1

"""
    nufft1d1(xj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             ms      :: Integer
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-1 1D complex nonuniform FFT.
"""
function nufft1d1(xj::Array{T},
                  cj::Array{Complex{T}},
                  iflag::Integer,
                  eps::T,
                  ms::Integer,
                  opts::nufft_opts=finufft_default_opts()) where T <: Union{Float32,Float64}
    fk = Array{Complex{T}}(undef, ms)
    nufft1d1!(xj, cj, iflag, eps, fk, opts)
    return fk
end

"""
    nufft2d1(xj      :: Array{Float64},
             yj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             ms      :: Integer,
             mt      :: Integer,
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-1 2D complex nonuniform FFT.
"""
function nufft2d1(xj      :: Array{T},
                  yj      :: Array{T},
                  cj      :: Array{Complex{T}},
                  iflag   :: Integer,
                  eps     :: T,
                  ms      :: Integer,
                  mt      :: Integer,
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    fk = Array{Complex{T}}(undef, ms, mt)
    nufft2d1!(xj, yj, cj, iflag, eps, fk, opts)
    return fk
end

"""
    nufft3d1(xj      :: Array{Float64},
             yj      :: Array{Float64},
             zj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             ms      :: Integer,
             mt      :: Integer,
             mu      :: Integer,
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-1 3D complex nonuniform FFT.
"""
function nufft3d1(xj      :: Array{T},
                  yj      :: Array{T},
                  zj      :: Array{T},
                  cj      :: Array{Complex{T}},
                  iflag   :: Integer,
                  eps     :: T,
                  ms      :: Integer,
                  mt      :: Integer,
                  mu      :: Integer,
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    fk = Array{Complex{T}}(undef, ms, mt, mu)
    nufft3d1!(xj, yj, zj, cj, iflag, eps, fk, opts)
    return fk
end

## Type-2

"""
    nufft1d2(xj      :: Array{Float64},
             iflag   :: Integer,
             eps     :: Float64,
             fk      :: Array{ComplexF64}
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-2 1D complex nonuniform FFT.
"""
function nufft1d2(xj      :: Array{T},
                  iflag   :: Integer,
                  eps     :: T,
                  fk      :: Array{Complex{T}},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    cj = Array{Complex{T}}(undef, nj)
    nufft1d2!(xj, cj, iflag, eps, fk, opts)
    return cj
end

"""
    nufft2d2(xj      :: Array{Float64},
             yj      :: Array{Float64},
             iflag   :: Integer,
             eps     :: Float64,
             fk      :: Array{ComplexF64}
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-2 2D complex nonuniform FFT.
"""
function nufft2d2(xj      :: Array{T},
                  yj      :: Array{T},
                  iflag   :: Integer,
                  eps     :: T,
                  fk      :: Array{Complex{T}},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    cj = Array{Complex{T}}(undef, nj)
    nufft2d2!(xj, yj, cj, iflag, eps, fk, opts)
    return cj
end

"""
    nufft3d2(xj      :: Array{Float64},
             yj      :: Array{Float64},
             zj      :: Array{Float64},
             iflag   :: Integer,
             eps     :: Float64,
             fk      :: Array{ComplexF64}
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-2 3D complex nonuniform FFT.
"""
function nufft3d2(xj      :: Array{T},
                  yj      :: Array{T},
                  zj      :: Array{T},
                  iflag   :: Integer,
                  eps     :: T,
                  fk      :: Array{Complex{T}},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    cj = Array{Complex{T}}(undef, nj)
    nufft3d2!(xj, yj, zj, cj, iflag, eps, fk, opts)
    return cj
end

## Type-3

"""
    nufft1d3(xj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             sk      :: Array{Float64},
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-3 1D complex nonuniform FFT.
"""
function nufft1d3(xj      :: Array{T},
                  cj      :: Array{Complex{T}},
                  iflag   :: Integer,
                  eps     :: T,
                  sk      :: Array{T},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    @assert length(cj)==nj
    nk = length(sk)
    fk = Array{Complex{T}}(undef, nk)
    nufft1d3!(xj, cj, iflag, eps, sk, fk, opts)
    return fk
end

"""
    nufft2d3(xj      :: Array{Float64},
             yj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             sk      :: Array{Float64},
             tk      :: Array{Float64}
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-3 2D complex nonuniform FFT.
"""
function nufft2d3(xj      :: Array{T},
                  yj      :: Array{T},
                  cj      :: Array{Complex{T}},
                  iflag   :: Integer,
                  eps     :: T,
                  sk      :: Array{T},
                  tk      :: Array{T},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    @assert length(cj)==nj
    nk = length(sk)
    fk = Array{Complex{T}}(undef, nk)
    nufft2d3!(xj, yj, cj, iflag, eps, sk, tk, fk, opts)
    return fk
end

"""
    nufft3d3(xj      :: Array{Float64},
             yj      :: Array{Float64},
             zj      :: Array{Float64},
             cj      :: Array{ComplexF64},
             iflag   :: Integer,
             eps     :: Float64,
             sk      :: Array{Float64},
             tk      :: Array{Float64}
             uk      :: Array{Float64}
             [, opts :: nufft_opts]
            ) -> Array{ComplexF64}

Compute type-3 3D complex nonuniform FFT.
"""
function nufft3d3(xj      :: Array{T},
                  yj      :: Array{T},
                  zj      :: Array{T},
                  cj      :: Array{Complex{T}},
                  iflag   :: Integer,
                  eps     :: T,
                  sk      :: Array{T},
                  tk      :: Array{T},
                  uk      :: Array{T},
                  opts    :: nufft_opts = finufft_default_opts()) where T <: Union{Float32,Float64}
    nj = length(xj)
    @assert length(cj)==nj
    nk = length(sk)
    fk = Array{Complex{T}}(undef, nk)
    nufft3d3!(xj, yj, zj, cj, iflag, eps, sk, tk, uk, fk, opts)
    return fk
end

### Direct interfaces (No allocation)

## 1D

"""
    nufft1d1!(xj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-1 1D complex nonuniform FFT. Output stored in fk.
"""
function nufft1d1!(xj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(cj)==nj
    ms = length(fk)
    # Calling interface
    # int finufft1d1(BIGINT nj,FLT* xj,CPX* cj,int iflag,FLT eps,BIGINT ms,
    # 	       CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft1d1, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, cj, iflag, eps, ms, fk, opts
                )
    check_ret(ret)
end

"""
    nufft1d2!(xj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-2 1D complex nonuniform FFT. Output stored in cj.
"""
function nufft1d2!(xj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(cj)==nj
    ms = length(fk)
    # Calling interface
    # int finufft1d2(BIGINT nj,FLT* xj,CPX* cj,int iflag,FLT eps,BIGINT ms,
    #                CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft1d2, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, cj, iflag, eps, ms, fk, opts
                )
    check_ret(ret)
end

"""
    nufft1d3!(xj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              sk      :: Array{Float64},
              fk      :: Array{ComplexF64},
              [, opts :: nufft_opts]
             )

Compute type-3 1D complex nonuniform FFT. Output stored in fk.
"""
function nufft1d3!(xj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   sk      :: Array{Float64},
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(cj)==nj
    nk = length(sk)
    @assert length(fk)==nk
    # Calling interface
    # int finufft1d3(BIGINT nj,FLT* x,CPX* c,int iflag,FLT eps,BIGINT nk, FLT* s, CPX* f, nufft_opts opts);
    ret = ccall( (:finufft1d3, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                Ref{Cdouble},
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, cj, iflag, eps, nk, sk, fk, opts
                )
    check_ret(ret)
end

## 2D

"""
    nufft2d1!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-1 2D complex nonuniform FFT. Output stored in fk.
"""
function nufft2d1!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(cj)==nj
    ms, mt = size(fk)
    # Calling interface
    # int finufft2d1(BIGINT nj,FLT* xj,FLT *yj,CPX* cj,int iflag,FLT eps,
    #                BIGINT ms, BIGINT mt, CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft2d1, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, cj, iflag, eps, ms, mt, fk, opts
                )
    check_ret(ret)
end

"""
    nufft2d2!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-2 2D complex nonuniform FFT. Output stored in cj.
"""
function nufft2d2!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(cj)==nj
    ms, mt = size(fk)
    # Calling interface
    # int finufft2d2(BIGINT nj,FLT* xj,FLT *yj,CPX* cj,int iflag,FLT eps,
    #                BIGINT ms, BIGINT mt, CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft2d2, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, cj, iflag, eps, ms, mt, fk, opts
                )
    check_ret(ret)
end

"""
    nufft2d3!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              sk      :: Array{Float64},
              tk      :: Array{Float64},
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
             )

Compute type-3 2D complex nonuniform FFT. Output stored in fk.
"""
function nufft2d3!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   sk      :: Array{Float64},
                   tk      :: Array{Float64},
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(cj)==nj
    nk = length(sk)
    @assert length(tk)==nk
    @assert length(fk)==nk
    # Calling interface
    # int finufft2d3(BIGINT nj,FLT* x,FLT *y,CPX* cj,int iflag,FLT eps,BIGINT nk, FLT* s, FLT* t, CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft2d3, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, cj, iflag, eps, nk, sk, tk, fk, opts
                )
    check_ret(ret)
end

## 3D

"""
    nufft3d1!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              zj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-1 3D complex nonuniform FFT. Output stored in fk.
"""
function nufft3d1!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   zj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(zj)==nj
    @assert length(cj)==nj
    ms, mt, mu = size(fk)
    # Calling interface
    # int finufft3d1(BIGINT nj,FLT* xj,FLT *yj,FLT *zj,CPX* cj,int iflag,FLT eps,
    # 	       BIGINT ms, BIGINT mt, BIGINT mu, CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft3d1, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                BIGINT,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, zj, cj, iflag, eps, ms, mt, mu, fk, opts
                )
    check_ret(ret)
end

"""
    nufft3d2!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              zj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
            )

Compute type-2 3D complex nonuniform FFT. Output stored in cj.
"""
function nufft3d2!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   zj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(zj)==nj
    @assert length(cj)==nj
    ms, mt, mu = size(fk)
    # Calling interface
    # int finufft3d2(BIGINT nj,FLT* xj,FLT *yj,FLT *zj,CPX* cj,int iflag,FLT eps,
    #                BIGINT ms, BIGINT mt, BIGINT mu, CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft3d2, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                BIGINT,
                BIGINT,
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, zj, cj, iflag, eps, ms, mt, mu, fk, opts
                )
    check_ret(ret)
end

"""
    nufft3d3!(xj      :: Array{Float64},
              yj      :: Array{Float64},
              zj      :: Array{Float64},
              cj      :: Array{ComplexF64},
              iflag   :: Integer,
              eps     :: Float64,
              sk      :: Array{Float64},
              tk      :: Array{Float64},
              uk      :: Array{Float64},
              fk      :: Array{ComplexF64}
              [, opts :: nufft_opts]
             )

Compute type-3 3D complex nonuniform FFT. Output stored in fk.
"""
function nufft3d3!(xj      :: Array{Float64},
                   yj      :: Array{Float64},
                   zj      :: Array{Float64},
                   cj      :: Array{ComplexF64},
                   iflag   :: Integer,
                   eps     :: Float64,
                   sk      :: Array{Float64},
                   tk      :: Array{Float64},
                   uk      :: Array{Float64},
                   fk      :: Array{ComplexF64},
                   opts    :: nufft_opts = finufft_default_opts())
    nj = length(xj)
    @assert length(yj)==nj
    @assert length(zj)==nj
    @assert length(cj)==nj
    nk = length(sk)
    @assert length(tk)==nk
    @assert length(uk)==nk
    @assert length(fk)==nk
    # Calling interface
    # int finufft3d3(BIGINT nj,FLT* x,FLT *y,FLT *z, CPX* cj,int iflag,
    #                FLT eps,BIGINT nk,FLT* s, FLT* t, FLT *u,
    #                CPX* fk, nufft_opts opts);
    ret = ccall( (:finufft3d3, libfinufft),
                Cint,
                (BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Cint,
                Cdouble,
                BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{ComplexF64},
                Ref{nufft_opts}),
                nj, xj, yj, zj, cj, iflag, eps, nk, sk, tk, uk, fk, opts
                )
    check_ret(ret)
end

# Load single precision interfaces
include("single.jl")
include("guru.jl")

end # module


b"""
Parameter Inference with ABC-SMC.
A thin wrapper around the GpABC pacakge.
"""

using DataFrames, DataStructures
using Distributions
using DocStringExtensions
using Pipe
using PyCall
using StatsBase, LsqFit

np = pyimport("numpy")
stats = pyimport("scipy.stats")
include("Utils.jl")

"""
Get mean and standard deviation for scaling purposes, accounting for `Inf`s and `0`s.
$(TYPEDSIGNATURES)
"""
function scaler(x)
    if length(x)>0
        s_mean = mean(x)
        s_std = std(x)

        if isfinite(s_mean)==false
            s_mean=0
        end

        if s_std==0
            s_std=1
        end
        return (s_mean, s_std)
    else
        return (0, 1)
    end
end

"""
Get absolute upper and lower limits for a parameter, based on the sampling distribution `d`.

For uniform distributions, the limits are the lower and upper limit of the distribution.

For other distributions, the 0.01st and 99.9 th percentile are calcualted.

$(TYPEDSIGNATURES)
"""
function get_limits(d)
    if typeof(d)==Uniform{Float64}
        return [d.a, d.b]
    else
        lower=quantile(d, 0.001)
        upper=quantile(d, 0.999)
        return [lower, upper]
    end
end

"""
Simple Sequential ABC, loosely based on the Toni et al. algorithm.
$(TYPEDSIGNATURES)
"""
function SMC_ABC(
        priors::OrderedDict{String,Any},
        reference_data::Union{Array{DataFrame,1}, Array{Float64,1}, DataFrame, Float64},
        simulator_function::Function,
        loss_function::Function;
        n_loops=10,
        n_samples=500,
        n_samples_init=1000,
        acceptance_quantile=0.2,
        perturbation_factor=0.05,
        progress_every=n_loops/5,
        boundaries="absorbing"
    )

    # infer parameter limits based on the marginal prior pdf

    limits = [get_limits(d[2]) for d in priors]
    lower = [l[1] for l in limits]
    upper = [l[2] for l in limits]

    # arrays to store accepted particles + distances
    accepted_particles = zeros(n_samples)
    accepted_distances = zeros(n_samples)

    trace_df = DataFrame()

    # how many samples per loop?
    n_samples_array = vcat([n_samples_init], repeat([n_samples], n_loops))

    # TODO:
    # - implement cross-validation to automatically determine when to stop the calibration?
    # - also check coefficient of variation of parameter values at the end of every loop;
    #   if variation in all parameters is very low (e.g. <1%), calibration can also be stopped
    for loop in 1:(n_loops+1)

        if loop % progress_every==0
            @info("Loop "*string(loop)*" of "*string(n_loops+1))
        end

        distances = []
        particles = []

        # evaluate the loss function for some number of random samples
        for i in 1:n_samples_array[loop]
            let particle = zeros(length(priors))
                # in the first loop,
                # sample from priors
                if loop==1
                    particle = [rand(x) for x in priors.vals]
                # in the 1+nth loop,
                # sample from the particles which have been accepted in the nth loop
                else
                    particle = StatsBase.sample(accepted_particles, Weights(1 ./ accepted_distances))
                    perturbation = exp.(rand(Normal(0, perturbation_factor), size(particle)))
                    particle = particle .* perturbation

                    if boundaries=="absorbing"
                        # absorbing boundaries
                        particle = max.(lower, particle)
                        particle = min.(upper, particle)
                    elseif boundaries=="none"
                        nothing
                    else
                        error("Unknwon boundary rule "*string(boundaries))
                    end
                end

                # generate predicted data
                predicted = simulator_function(particle)

                # evaluate the loss function
                let distance = Inf
                    try
                        distance = loss_function(predicted, reference_data)
                    catch
                        distance = Inf
                    end
                # store values
                append!(distances, distance)
                append!(particles, [particle])
                end
            end
        end

        # at the end of the loop:
        # filter out infinite distances/particles
        isfin = isfinite.(distances)

        distances = distances[isfin]
        particles = particles[isfin]

        if loop==1
            num_invalid=length(isfin)-sum(isfin)
            if num_invalid>1
                @info(string(num_invalid)*" of "*string(length(isfin))* " samples invalid during intialization")
            end
        end

        # get the best-fitting fraction of particle
        accepted_particles = particles[distances.<=quantile(distances, acceptance_quantile)]
        accepted_distances = distances[distances.<=quantile(distances, acceptance_quantile)]

        if length(accepted_particles)==0
            error("Did not accept any particles. Bug in loss_function or simulator_function?")
        end

        # record trace data
        minloc = argmin(accepted_distances)
        pt_est = accepted_particles[minloc]
        rho = accepted_distances[minloc]
        trace_df_step = DataFrame(hcat(hcat(pt_est...), rho))
        trace_df_step[!,:step] .= loop
        rename!(trace_df_step, vcat(priors.keys, "rho", "step"))
        append!(trace_df, trace_df_step)
    end

    accepted_df = DataFrame(vcat(accepted_particles'...))
    rename!(accepted_df, priors.keys)
    accepted_df[!,:rho] = accepted_distances

    return trace_df, accepted_df
end

"""
Retrieve point estimate (min distances) with marginal uncertainty (95% credible limits).
$(SIGNATURES)
"""
function get_estimate(accepted::DataFrame)
    # this is assuming that all but the last column are parameter values!
    estimate = accepted[accepted.rho.==minimum(accepted.rho),1:end-1] |> Array |> np.ravel
    CL5 = [quantile(x,.5) for x in eachcol(accepted[:,1:end-1])]
    CL95 = [quantile(x,.95) for x in eachcol(accepted[:,1:end-1])]
    return DataFrame(
        param = names(accepted)[1:end-1],
        estimate = estimate,
        CL5 = CL5,
        CL95 = CL95
    )
end

"""
Get random sample from a DataFrame of accepted particles, weighted by inverse loss.
$(SIGNATURES)
"""
function posterior_sample(accepted::DataFrame)
    weights = Weights(1 ./ accepted.rho)
    params = names(accepted)[1:end-1]
    cols = Symbol.(params)
    row_sample = sample(Array(1:nrow(accepted)), weights)
    par_sample = accepted[row_sample,cols]
    return Array(par_sample)
end

"""
Fit a dthree-parameter log logistic function to data. \n
$(SIGNATURES)
"""
function fit_log_logistic(
        df,
        x_var,
        y_var;
        model = @. model(x, p) = p[2]/(1+(x/p[1])^p[3])
    )

    # extract relevant columns
    df = @pipe df[:,[x_var, y_var]] |>
    # drop missing
    _[completecases(_),:]

    try
        # set boundaries for parameters estimation
        lower_limits = [0., 0., 1e-5]
        upper_limits = [maximum(df[:,x_var]), maximum(df[:,y_var]*1.1), 20]
        # estimate parameters + error
        fit = curve_fit(model, df[:,x_var], df[:,y_var], [mean(df[:,x_var]), mean(df[:,y_var]), 3], lower=lower_limits, upper=upper_limits)
        errors = stderror(fit)
        confidence = confidence_interval(fit, 0.05)
        # return as DataFrame
        return DataFrame(
            param = ["EC50", "upper_limit", "beta"],
            estimate = fit.param,
            std_err = errors,
            CI_05 = [max(x[1], 0) for x in confidence],
            CI_95 = [max(x[2], 0) for x in confidence]
        )
    catch
        # return missing if model can't be fitted
        return DataFrame(
            param = ["EC50", "upper_limit", "beta"],
            estimate = repeat([missing], 3),
            std_err = repeat([missing], 3),
            CI_05 = repeat([missing], 3),
            CI_95 = repeat([missing], 3)
        )
    end
end

"""
Get dose-rsponse parameters per time-point. Wrapper around ´fit_log_logistic´.
$(SIGNATURES)
"""
function dose_response_per_timepoint(
        df::DataFrame,
        x_var::Symbol,
        y_var::Symbol)
    params = combine(groupby(df, :tday), x -> fit_log_logistic(x, x_var, y_var))
    return params
end

"""
Assign toxicity parameters. Stressor is either one of "Cu", "Ni", "Zn".
$(SIGNATURES)
"""
function assign_params(fleas_dict, particle, stressor)

    i = findall(x->x==stressor, ["Cu", "Ni", "Zn"])[1]

    fleas_dict["toxicity_particle"]["k_d"][i]    = particle[1]
    fleas_dict["toxicity_particle"]["h_max"][i]  = particle[2]
    fleas_dict["toxicity_particle"]["h_ED50"][i] = particle[3]
    fleas_dict["toxicity_particle"]["h_beta"][i] = particle[4]
    fleas_dict["toxicity_particle"]["S_max"][:,i]  = particle[5:8]
    fleas_dict["toxicity_particle"]["S_ED50"][:,i] = particle[9:12]
    fleas_dict["toxicity_particle"]["S_beta"][:,i] = particle[13:16]

    return fleas_dict
end

"""
Assign toxicity parameters for a particular pMoA. `i` is the index of the stressor.
$(SIGNATURES)
"""
function assign_params(fleas_dict, particle, stressor, pmoa)

    i = findall(x->x==stressor, ["Cu", "Ni", "Zn"])[1]
    j = findall(x->x==pmoa, ["Growth", "Maintenance", "Assimilation", "Reproduction"])[1]

    fleas_dict["toxicity_particle"]["S_max"] .= 0.

    fleas_dict["toxicity_particle"]["k_d"][i]    = particle[1]
    fleas_dict["toxicity_particle"]["h_max"][i]  = particle[2]
    fleas_dict["toxicity_particle"]["h_ED50"][i] = particle[3]
    fleas_dict["toxicity_particle"]["h_beta"][i] = particle[4]
    fleas_dict["toxicity_particle"]["S_max"][j,i]  = particle[5]
    fleas_dict["toxicity_particle"]["S_ED50"][j,i] = particle[6]
    fleas_dict["toxicity_particle"]["S_beta"][j,i] = particle[7]

    return fleas_dict
end