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


text:
using LDrawParser
using LightGraphs, GraphUtils
using Logging
using GeometryBasics, CoordinateTransformations, Rotations

# global_logger(SimpleLogger(stderr, Logging.Debug))

filename = joinpath(dirname(pathof(LDrawParser)),"..","assets","ATTEWalker.mpd")
# filename = joinpath(dirname(pathof(LDrawParser)),"..","assets","Millennium Falcon.mpd")
model = parse_ldraw_file(filename)
# load geometry
LDrawParser.populate_part_geometry!(model)

# part_keys = sort(collect(keys(model.parts))[1:10])
# parts = Dict(k=>model.parts[k] for k in part_keys)
# transport_model = (
#     robot_radius = 15.0,
#     max_area_per_robot = 10000.0, #3000.0,
#     max_volume_per_robot = 1000000.0 #20000.0,
# )
# for (k,part) in parts
#     points = map(T, LDrawParser.extract_points(part))
#     if isempty(points)
#         println("Part $k has no geometry!")
#         continue
#     end
#     geom=map(SVector,points)
#     try
#         support_pts = HierarchicalGeometry.select_support_locations(
#             geom,transport_model)
#         polygon = VPolygon(convex_hull(map(p->p[1:2],geom)))
#
#         plt = plot(polygon,aspectratio=1,alpha=0.4)
#         scatter!(plt,map(p->p[1],geom),map(p->p[2],geom),label=k) #,legend=false)
#         plot!(plt,map(p->Ball2(p,transport_model.robot_radius),support_pts), aspectratio=1, alpha=0.4)
#         display(plt)
#     catch e
#         bt = catch_backtrace()
#         showerror(stdout,e,bt)
#     end
# end

vis = Visualizer()
render(vis)
vis_root = vis["root"]

T_base = CoordinateTransformations.Translation(0.0,0.0,0.0) ∘ CoordinateTransformations.LinearMap(LDrawParser.LDRAW_BASE_FRAME)
SCALE = 0.01
LDrawParser.change_coordinate_system!(model,T_base,SCALE)

m = model.models["20009 - AT-TE Walker.mpd"]
id_generator = DuplicateIDGenerator{String}()
for step in m.steps
    for ref in step.lines
        global vis_root
        @show ref
        if !LDrawParser.has_part(model,ref.file)
            continue
        end
        p = model.parts[ref.file]
        name = id_generator(p.name)
        vec = LDrawParser.extract_surface_geometry(p)
        M = GeometryBasics.Mesh(coordinates(vec),faces(vec))
        setobject!(vis_root[name], M)
        tr = LDrawParser.build_transform(ref)
        settransform!(vis_root[name], tr)
    end
end
delete!(vis)

# construct model graph
model_graph = construct_assembly_graph(model)
model_tree = convert(GraphUtils.CustomNTree{GraphUtils._node_type(model_graph),String},model_graph)
# @assert maximum(map(v->indegree(model_tree,v),vertices(model_tree))) == 1
print(model_tree,v->summary(v.val),"\t")

sched = LDrawParser.construct_model_schedule(model)
model_spec = LDrawParser.extract_single_model(sched,"20009 - AT-TE Walker.mpd")

GraphUtils.validate_graph(model_spec)

ext: jl


text:
using Distributed
using ProgressMeter
using Logging
using SharedArrays
using JLD2
using Dates

include("slurm.jl")
using .ClusterManagers

function IN_SLURM()
    return "SLURM_JOBID" in keys(ENV)
end

"""
job
Interface into running a job.
"""

function job(experiment_file::AbstractString,
             exp_dir::AbstractString,
             args_iter;
             exp_module_name::Union{String, Symbol}=:Main,
             exp_func_name::Union{String, Symbol}=:main_experiment,
             num_workers::Integer=5,
             expand_args::Bool=false,
             extra_args = [],
             store_exceptions=true,
             exception_dir="except",
             job_file_dir="")
    if "SLURM_ARRAY_TASK_ID" in keys(ENV)
        @info "This is an array Job! Time to get task and start job."
        task_id = parse(Int64, ENV["SLURM_ARRAY_TASK_ID"])
        @time task_job(experiment_file, exp_dir, args_iter, task_id;
                       exp_module_name=exp_module_name,
                       exp_func_name=exp_func_name,
                       expand_args=expand_args,
                       extra_args=extra_args,
                       store_exceptions=store_exceptions,
                       exception_dir=exception_dir)
    else
        if IN_SLURM()
            @time slurm_parallel_job(experiment_file, exp_dir, args_iter;
                                     exp_module_name=exp_module_name,
                                     exp_func_name=exp_func_name,
                                     expand_args=expand_args,
                                     extra_args=extra_args,
                                     store_exceptions=store_exceptions,
                                     exception_dir=exception_dir,
                                     job_file_dir=job_file_dir)
        else
            @time parallel_job(experiment_file, exp_dir, args_iter;
                               exp_module_name=exp_module_name,
                               exp_func_name=exp_func_name,
                               num_workers=num_workers,
                               expand_args=expand_args,
                               extra_args=extra_args,
                               store_exceptions=store_exceptions,
                               exception_dir=exception_dir)
        end
    end

end

function job(experiment_file::AbstractString,
             exp_dir::AbstractString,
             args_iter,
             task_id::Integer;
             exp_module_name::Union{String, Symbol}=:Main,
             exp_func_name::Union{String, Symbol}=:main_experiment,
             num_workers::Integer=5,
             expand_args::Bool=false,
             extra_args = [],
             store_exceptions=true,
             exception_dir="except")
    @info "This is a task job! ID is $(task_id)"
    @time task_job(experiment_file, exp_dir, args_iter, task_id;
                   exp_module_name=exp_module_name,
                   exp_func_name=exp_func_name,
                   expand_args=expand_args,
                   extra_args=extra_args,
                   store_exceptions=store_exceptions,
                   exception_dir=exception_dir)
end

job(exp::Experiment; exception_dir="except", kwargs...) =
    job(exp.file, exp.dir, exp.args_iter;
        exp_module_name=exp.module_name,
        exp_func_name=exp.func_name,
        exception_dir="$(exception_dir)/exp_0x$(string(exp.hash, base=16))", kwargs...)

job(exp::Experiment, job_id::Integer; exception_dir="except", kwargs...) =
    job(exp.file, exp.dir, exp.args_iter, job_id;
        exp_module_name=exp.module_name,
        exp_func_name=exp.func_name,
        exception_dir="$(exception_dir)/exp_0x$(string(exp.hash, base=16))", kwargs...)

"""
parallel_job

Run a parallel job over the arguments presented by args_iter. `args_iter` can be a enumeration OR ArgIterator. Each job will be dedicated to a specific
task. The experiment *must* save its own data! As this is not handled by this function (although could be added in the future.)
"""

function parallel_job(experiment_file::AbstractString,
                      exp_dir::AbstractString,
                      args_iter;
                      exp_module_name::Union{String, Symbol}=:Main,
                      exp_func_name::Union{String, Symbol}=:main_experiment,
                      num_workers=1,
                      expand_args=false,
                      project=".",
                      extra_args=[],
                      store_exceptions=true,
                      exception_dir="except", verbose=false)

    #######
    #
    # Preamble - Add processes, initialized shared memory
    #
    ######

    num_add_workers = num_workers - 1
    pids = Array{Int64, 1}

    args_list = collect(args_iter)
    exc_opts = Base.JLOptions()
    color_opt = "no"
    if exc_opts.color == 1
        color_opt = "yes"
    end
    if num_add_workers != 0
        println(num_workers, " ", nworkers())
        if nworkers() == 1
            pids = addprocs(num_workers; exeflags=["--project=$(project)", "--color=$(color_opt)"])
        elseif nworkers() < num_workers
            pids = addprocs((num_workers) - nworkers(); exeflags=["--project=$(project)", "--color=$(color_opt)"])
        else
            pids = procs()
        end
    end

    println(nworkers(), " ", pids)

    n = length(args_iter)
    job_ids = SharedArray{Int64, 1}(n)
    finished_jobs = SharedArray(fill(false, n))

    #########
    #
    # Meaty middle: Compiling code, running jobs, managing which jobs fail.
    #
    ########

    try

        channel = RemoteChannel(()->Channel{Bool}(length(args_iter)), 1)

        mod_str = string(exp_module_name)
        func_str = string(exp_func_name)

        try
            @everywhere const global exp_file=$experiment_file
            @everywhere const global expand_args=$expand_args
            @everywhere const global extra_args=$extra_args
            @everywhere const global store_exceptions=$store_exceptions
            @everywhere const global exception_loc = joinpath($exp_dir, $exception_dir)

            @everywhere begin
                eval(:(using Reproduce))
                eval(:(using Distributed))
                eval(:(using SharedArrays))
                include(exp_file)
                @info "$(exp_file) included on process $(myid())"
                mod = $mod_str=="Main" ? Main : getfield(Main, Symbol($mod_str))
                const global exp_func = getfield(mod, Symbol($func_str))
                experiment(args) = exp_func(args)
                @info "Experiment built on process $(myid())"
            end
        catch ex
            println(ex)
        end

        @info "Number of Jobs: $(n)"
        exception_loc = joinpath(exp_dir, exception_dir)
        if store_exceptions && !isdir(exception_loc)
            mkpath(exception_loc)
        end

        ProgressMeter.@showprogress pmap(args_iter) do (job_id, args)
            try
                if expand_args
                    Main.exp_func(args..., extra_args...)
                else
                    Main.exp_func(args, extra_args...)
                end
                finished_jobs[job_id] = true
            catch ex
                if isa(ex, InterruptException)
                    throw(InterruptException())
                end
                if verbose
                    @warn "Exception encountered for job: $(job_id)"
                end
                if store_exceptions
                    exception_file(
                        joinpath(exception_loc, "job_$(job_id).exc"),
                        job_id, ex, stacktrace(catch_backtrace()))
                end
            end
            Distributed.put!(channel, true)
            job_ids[job_id] = myid()
        end

    catch ex
        println(ex)
        Distributed.interrupt()
    end

    ########
    #
    # Finished. Return which jobs were unsuccessful.
    #
    ########

    println(job_ids)

    return findall((x)->x==false, finished_jobs)

end

function slurm_parallel_job(experiment_file::AbstractString,
                            exp_dir::AbstractString,
                            args_iter;
                            exp_module_name::Union{String, Symbol}=:Main,
                            exp_func_name::Union{String, Symbol}=:main_experiment,
                            expand_args=false,
                            project=".",
                            extra_args=[],
                            store_exceptions=true,
                            exception_dir="except", verbose=false, job_file_dir="")

    #######
    #
    # Preamble - Add processes, initialized shared memory
    #
    ######

    println("SLURM PARALLEL JOB")
    num_add_workers = parse(Int64, ENV["SLURM_NTASKS"])
    pids = Array{Int64, 1}

    args_list = collect(args_iter)
    exc_opts = Base.JLOptions()
    color_opt = "no"
    if exc_opts.color == 1
        color_opt = "yes"
    end
    if num_add_workers != 0
        # assume started fresh julia instance...
	println("Adding Slurm Jobs!!!")
        if job_file_dir == ""
            job_file_dir = joinpath(exp_dir, "jobs")
        end
        pids = addprocs(SlurmManager(num_add_workers); exeflags=["--project=$(project)", "--color=$(color_opt)"], job_file_loc=job_file_dir)
        print("\n")
    end

    println(nworkers(), " ", pids)
    n = length(args_iter)

    #########
    #
    # Meaty middle: Compiling code, running jobs, managing which jobs fail.
    #
    ########

    finished_jobs = RemoteChannel(()->Channel{Int}(n), 1)

    try

        channel = RemoteChannel(()->Channel{Bool}(length(args_iter)), 1)

        mod_str = string(exp_module_name)
        func_str = string(exp_func_name)
        try
            @everywhere const global exp_file=$experiment_file
            @everywhere const global expand_args=$expand_args
            @everywhere const global extra_args=$extra_args
            @everywhere const global store_exceptions=$store_exceptions
            @everywhere const global exception_loc = joinpath($exp_dir, $exception_dir)

            @everywhere begin
                eval(:(using Reproduce))
                eval(:(using Distributed))
                eval(:(using SharedArrays))
                include(exp_file)
                @info "$(exp_file) included on process $(myid())"
                mod = $mod_str=="Main" ? Main : getfield(Main, Symbol($mod_str))
                const global exp_func = getfield(mod, Symbol($func_str))
                experiment(args) = exp_func(args)
                @info "Experiment built on process $(myid())"
            end
        catch ex
            println(ex)
        end

        @info "Number of Jobs: $(n)"
        exception_loc = joinpath(exp_dir, exception_dir)
        if store_exceptions && !isdir(exception_loc)
            mkpath(exception_loc)
        end

        ProgressMeter.@showprogress pmap(args_iter) do (job_id, args)
            try
                if expand_args
                    Main.exp_func(args..., extra_args...)
                else
                    Main.exp_func(args, extra_args...)
                end
                Distributed.put!(finished_jobs, job_id)
            catch ex
                if isa(ex, InterruptException)
                    throw(InterruptException())
                end
                if verbose
                    @warn "Exception encountered for job: $(job_id)"
                end
                if store_exceptions
                    exception_file(
                        joinpath(exception_loc, "job_$(job_id).exc"),
                        job_id, ex, stacktrace(catch_backtrace()))
                end
            end
            Distributed.put!(channel, true)
        end

    catch ex
        println(ex)
        Distributed.interrupt()
    end

    #########
    #
    # Finished: get from channel which jobs succeeded.
    #
    ########

    finished_jobs_bool = fill(false, n)
    while isready(finished_jobs)
        job_id = take!(finished_jobs)
        finished_jobs_bool[job_id] = true
    end
    return findall((x)->x==false, finished_jobs_bool)
end

function task_job(experiment_file::AbstractString, exp_dir::AbstractString,
                  args_iter, task_id::Integer;
                  exp_module_name::Union{String, Symbol}=:Main,
                  exp_func_name::Union{String, Symbol}=:main_experiment,
                  expand_args::Bool=false,
                  extra_args=[],
                  store_exceptions=true,
                  exception_dir="except")

    mod_str = string(exp_module_name)
    func_str = string(exp_func_name)

    @everywhere begin
        include($experiment_file)
        mod = $mod_str=="Main" ? Main : getfield(Main, Symbol($mod_str))
        const global exp_func = getfield(mod, Symbol($func_str))
    end

    args = collect(args_iter)[task_id][2]
    @sync @async begin
        try
            if expand_args
                Main.exp_func(args..., extra_args...)
            else
                Main.exp_func(args, extra_args...)
            end
        catch ex
            @warn "Exception encountered for job: $(task_id)"
            if store_exceptions
                exception_loc = joinpath(exp_dir, exception_dir)
                if !isdir(exception_loc)
                    try
                        mkpath(exception_loc)
                    catch
                        sleep(1)
                    end
                end
                exception_file(
                    joinpath(exception_loc, "job_$(task_id).exc"),
                    task_id, ex, stacktrace(catch_backtrace()))
                return false
            end
            throw(ex)
        end
    end
    return true
end


text:
fft_frequencies(fs::Real, nfft::Int) = LinRange(0f0, fs / 2f0, (nfft >> 1) + 1)

"""
    MelSpectrogram{T, F, Ti} <: DSP.Periodograms.TFR{T}

# Arguments:
- `power::Matrix{T}`
- `mels::F`
- `time::Ti`
"""
struct MelSpectrogram{T, F,Ti} <: DSP.Periodograms.TFR{T}
    power::Matrix{T}
    mels::F
    time::Ti
end
DSP.freq(tfr::MelSpectrogram) = tfr.mels
Base.time(tfr::MelSpectrogram) = tfr.time

"""
    MFCC{T, F, Ti} <: DSP.Periodograms.TFR{T}

# Arguments:
- `mfcc::Matrix{T}`
- `number::F`
- `time::Ti`
"""
struct MFCC{T, F, Ti} <: DSP.Periodograms.TFR{T}
    mfcc::Matrix{T}
    number::F
    time::Ti
end
DSP.freq(tfr::MFCC) = tfr.number
Base.time(tfr::MFCC) = tfr.time

function hz_to_mel(frequencies)
    f_min = 0f0
    f_sp = 200f0 / 3

    mels = collect((frequencies .- f_min) ./ f_sp)
    min_log_hz = 1000f0
    min_log_mel = (min_log_hz .- f_min) ./ f_sp
    logstep = log(6.4f0) / 27f0

    @inbounds for i = 1:length(mels)
        if frequencies[i] >= min_log_hz
            mels[i] = min_log_mel + log(frequencies[i] / min_log_hz) / logstep
        end
    end

    mels
end

function mel_to_hz(mels)
    f_min = 0f0
    f_sp = 200f0 / 3
    frequencies = collect(f_min .+ f_sp .* mels)

    min_log_hz = 1000f0
    min_log_mel = (min_log_hz .- f_min) ./ f_sp
    logstep = log(6.4f0) / 27f0

    @inbounds for i = 1:length(frequencies)
        if mels[i] >= min_log_mel
            frequencies[i] = min_log_hz * exp(logstep * (mels[i] - min_log_mel))
        end
    end

    frequencies
end

function mel_frequencies(nmels::Int = 128, fmin::Real = 0.0f0, fmax::Real = 11025f0)
    min_mel = hz_to_mel(fmin)[1]
    max_mel = hz_to_mel(fmax)[1]

    mels = LinRange(min_mel, max_mel, nmels)
    mel_to_hz(mels)
end

"""
    M = mel(fs::Real, nfft::Int; nmels::Int = 128, fmin::Real = 0f0, fmax::Real = fs/2f0)

Returns a Mel matrix `M` such that `M*f` is a mel spectrogram if `f` is a vector of spectrogram powers, e.g.,
```julia
M*abs2.(rfft(sound))
```
"""
function mel(fs::Real, nfft::Int; nmels::Int = 128, fmin::Real = 0f0, fmax::Real = fs/2f0)
    weights = zeros(Float32, nmels, (nfft >> 1) + 1)
    fftfreqs = fft_frequencies(fs, nfft)
    melfreqs = mel_frequencies(nmels + 2, fmin, fmax)
    enorm = 2f0 ./ (melfreqs[3:end] - melfreqs[1:nmels])

    for i in 1:nmels
        lower = (fftfreqs .- melfreqs[i]) ./ (melfreqs[i+1] - melfreqs[i])
        upper = (melfreqs[i+2] .- fftfreqs) ./ (melfreqs[i+2] - melfreqs[i+1])

        weights[i, :] = max.(0, min.(lower, upper)) * enorm[i]
    end

    weights
end

"""
    melspectrogram(s, n=div(length(s), 8), args...; fs=1, nmels::Int=128, fmin::Real=0.0f0, fmax::Real=fs / 2.0f0, window=hanning, kwargs...)

DOCSTRING

#Arguments:
- `s`: signal
- `n`: number of points in each window
- `args`: are sent to `spectrogram`
- `fs`: sample frequency
- `nmels`: number of mel frequencies
- `fmin`: minimum freq
- `fmax`: maximum freq
- `window`: window function, defaults to hanning
- `kwargs`: are sent to `spectrogram`
"""
function melspectrogram(S::DSP.Periodograms.Spectrogram; fs=1, nmels::Int = 128, fmin::Real = 0f0, fmax::Real = fs / 2f0)
    n = 2size(S.power,1)-1
    data = mel(fs, n; nmels=nmels, fmin=fmin, fmax=fmax) * S.power
    nframes = size(data, 2)
    MelSpectrogram(data, LinRange(hz_to_mel(fmin)[1], hz_to_mel(fmax)[1], nmels), S.time)
end

function melspectrogram(s, n=div(length(s), 8), args...; fs=1, nmels::Int = 128, fmin::Real = 0f0, fmax::Real = fs / 2f0, window=hanning, kwargs...)
    S = DSP.spectrogram(s, n, args...; fs=fs, window=window, kwargs...)
    melspectrogram(S,fs=fs,nmels=nmels,fmin=fmin,fmax=fmax)
end

"""
    mfcc(s, args...; nmfcc::Int=20, nmels::Int=128, window=hanning, kwargs...)

DOCSTRING

#Arguments:
- `s`: signal
- `args`: are sent to `spectrogram`
- `nmfcc`: number of coeffs
- `nmels`: number of mel frequencies
- `window`: window function, defaults to hanning
- `kwargs`: are sent to `spectrogram`
"""
function mfcc(s, args...; nmfcc::Int = 20, nmels::Int = 128, window=hanning, kwargs...)
    if nmfcc >= nmels
        error("number of mfcc components should be less than the number of mel frequency bins")
    end
    M = melspectrogram(s, args...; nmels=nmels, window=window, kwargs...)
    mfcc = dct_matrix(nmfcc, nmels) * power(M)

    for frame in 1:size(mfcc, 2)
        mfcc[:, frame] /= norm(mfcc[:, frame])
    end

    MFCC(mfcc, 1:nmfcc, time(M))
end

"""returns the DCT filters"""
function dct_matrix(nfilters::Int, ninput::Int)
    basis = Array{Float32}(undef, nfilters, ninput)
    samples = (1f0:2f0:2ninput) * π / 2ninput
    for i = 1:nfilters
        basis[i, :] = cos.(i * samples)
    end

    basis *= sqrt(2f0/ninput)
    basis
end


text:
function sanitise_epcs(project)
    stockdata = project.dat["master"]

    # Find street names
    parseaddr(str) = split(str,", ")
    findstreet(str_arr) = str_arr[1+minimum([findfirst(length.(str_arr)[2:end].>3),length(str_arr)-3])]
    streetnames = findstreet.(parseaddr.(stockdata.Full_postal_address))
    streetlist = sort(unique(streetnames))

    # Find indices of all dwellings linked to each street in `streetlist`
    findall_arr(arr,val) = findall(arr.==val)
    i = findall_arr.((streetnames,),streetlist)

    # Util functions for propagating most common occurrences from `groupby`
    countoccur(arr,val) = sum(skipmissing(arr).==val)
    mostcommon_occur(arr,groupby) = length(arr)==sum(ismissing.(arr)) ? "tbc" : groupby[findmax(countoccur.((arr,),groupby))[2]]
    # Propagate based on most common occurrence on street
    propagate_mostcommon(newArr,refArr,groupby,i) = newArr[i].=mostcommon_occur(refArr[i],groupby)

    # Apply propagation method: Age bands
    stockdata[!,:ParaDwell_Age_Band] = fill!(Array{String}(undef,nrow(stockdata)),"")
    groupby = ["before 1919","1919-1929","1930-1949","1950-1964","1965-1975","1976-1983","1984-1991","1992-1998","1999-2002","2003-2007","2008 onwards"]
    propagate_mostcommon.((stockdata.ParaDwell_Age_Band,), (stockdata.EPC_Part_1_Construction_Age_Band,), (groupby,), i)

end

@time sanitise_epcs(project)


text:
# Thanks to Tamas Papp
@inline function view_slice_last(arr::AbstractArray{T,N}, dx::Int) where {T,N}
    view(arr, ntuple(_ -> Colon(), N - 1)..., dx)
end

@inline function view_slice_last(arr::AbstractArray{T,N}, dx::CartesianIndex{DX}) where {T,N,DX}
    view(arr, ntuple(_ -> Colon(), N - DX)..., dx)
end
@inline function view_slice_first(arr::AbstractArray{T,N}, dx::Int) where {T,N}
    view(arr, dx, ntuple(_ -> Colon(), N - 1)...)
end

@inline function view_slice_first(arr::AbstractArray{T,N}, dx::CartesianIndex{DX}) where {T,N,DX}
    view(arr, dx, ntuple(_ -> Colon(), N - DX)...)
end


text:
function end_round(cur_player::AbstractPlayer, is_restart::Bool=false)
  if is_repl
    println("""

        __ _   __ _  _ __ ___    ___    ___ __   __ ___  _ __
       / _` | / _` || '_ ` _ \\  / _ \\  / _ \\\\ \\ / // _ \\| '__|
      | (_| || (_| || | | | | ||  __/ | (_) |\\ V /|  __/| | _  _  _
       \\__, | \\__,_||_| |_| |_| \\___|  \\___/  \\_/  \\___||_|(_)(_)(_)
        __/ |
       |___/
    """)
  end

  cur_player.state.is_playing = false
  cur_player.state.is_done = true

  cur_js = """
    \$(".js-tetris-container").blur();
  """

  if is_restart
    cur_js *= """
      \$(".js-tetris-container td").removeClass();
      \$(".js-level-text").text("Level 1");
      \$(".js-score-text").text("$(lpad(0, 8, "0"))");
    """
  else
    cur_js *= """
      \$(".js-splash-text h1").text("game over.")
      \$(".js-splash-text div").text("try again soon.")
    """
  end

  is_new_high_score = !is_restart

  is_new_high_score = is_new_high_score &&
    cur_player.round.is_keeping_score

  is_new_high_score = is_new_high_score && (
    cur_player.score >
    cur_player.game.board.min_score
  )

  if !is_new_high_score
    set_default_values(cur_player)

    tetris_js(
      cur_player.game.scope,
      JSString(cur_js)
    )

    return false
  end

  cur_js *= """
    \$(".cs-scores-screen").addClass("hidden");
    \$(".js-new-score").removeClass("hidden");
    \$(".js-new-score h2").text("$(lpad(cur_player.score, 8, "0"))");
    \$(".js-new-score input").focus();
  """

  cur_func = function(cur_timer::Timer)
    cur_player.round.is_making_call && return
    send_logs(cur_player, false) || return

    tetris_js(
      cur_player.game.scope,
      JSString(cur_js)
    )

    close(cur_timer)
  end

  Timer(cur_func, 1, 10)

  false
end


text:
using JuMP

using SCIP
function scip_optimizer_factory()
    with_optimizer(SCIP.Optimizer)#, limits_time=2 * 60, heuristics_alns_priority=1000000, heuristics_subnlp_priority=1000000)
    # with_optimizer(SCIP.Optimizer,
    #     limits_gap=0.05, limits_time=10 * 60 * 60, display_verblevel=5,
    #     display_width=120, history_valuebased=true, lp_threads=10, branching_preferbinary=true, lp_scaling=false,
    #     branching_allfullstrong_priority=536870911, heuristics_multistart_freq=20, heuristics_multistart_onlynlps=false,
    #     heuristics_mpec_priority=536870911)#heuristics_subnlp_priority=536870911)#, nlp_solver="ipopt", heuristics_nlpdiving_priority=536870911)#,;
end

function scip_optimize_hook(model::JuMP.Model)
    caching_optimizer = backend(model)
    if caching_optimizer.mode == MOIU.AUTOMATIC && caching_optimizer.state == MOIU.EMPTY_OPTIMIZER
        MOIU.attach_optimizer(caching_optimizer)
    end
    mscip = caching_optimizer.optimizer.model.mscip
    SCIP.SCIPsetEmphasis(mscip, SCIP.SCIP_PARAMEMPHASIS_FEASIBILITY, true)#false); println()
    SCIP.SCIPsetPresolving(mscip, SCIP.SCIP_PARAMSETTING_AGGRESSIVE, true)#false); println()
    SCIP.SCIPsetHeuristics(mscip, SCIP.SCIP_PARAMSETTING_AGGRESSIVE, true)#false)
    # SCIP.set_parameter(mscip, "heuristics/alns/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/subnlp/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/mpec/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/feaspump/freq", 5)
    SCIP.set_parameter(mscip, "heuristics/rens/freq", -1)
    SCIP.set_parameter(mscip, "reoptimization/enable", true)
    SCIP.set_parameter(mscip, "limits/time", 2 * 60)
    # SCIP.set_parameter(mscip, "limits/time", 2 * 60 * 60)
    ret = MOI.optimize!(caching_optimizer)
    @assert caching_optimizer.optimizer.model.mscip === mscip
    return ret
end

using AmplNLWriter
AmplNLWriter.setdebug(true)

function bonmin_optimizer_factory()
    with_optimizer(AmplNLWriter.Optimizer, "/home/twan/code/bonmin/Bonmin-1.8.7/build/bin/bonmin")
end

function couenne_optimizer_factory()
    with_optimizer(AmplNLWriter.Optimizer, "/home/twan/code/couenne/couenne")
end

using Ipopt
# ENV["JULIA_IPOPT_LIBRARY_PATH"] = "/home/twan/code/ipopt/CoinIpopt/build/lib"
# ENV["JULIA_IPOPT_EXECUTABLE_PATH"] = "/home/twan/code/ipopt/CoinIpopt/build/bin"
# import Pkg; Pkg.build("Ipopt")

function ipopt_optimizer_factory()
    with_optimizer(Ipopt.Optimizer)
end

function complementarity_reformulate!(model::JuMP.Model)
    for var in JuMP.all_variables(problem.model)
        if JuMP.is_binary(var)
            @constraint model var * (1 - var) == 0
            JuMP.unset_binary(var)
        end
    end
end

# using Alpine
# function alpine_optimizer_factory()
#     with_optimizer(Alpine.Optimizer, nlp_optimizer=Ipopt.Optimizer(print_level=0), mip_optimizer=Gurobi.Optimizer(OutputFlag=0))
# end

# using Juniper
# function juniper_optimizer_factory()
#     params = Dict{Symbol,Any}()
#     params[:nl_solver] = with_optimizer(Ipopt.Optimizer, print_level=0)
#     # params[:mip_solver] = with_optimizer(Gurobi.Optimizer, OutputFlag=0)
#     params[:mip_solver] = with_optimizer(CPLEX.Optimizer, CPX_PARAM_SCRIND=0)
#     # Note to self: if you get "Cannot set bounds because variable is of type: BINARY", use LinQuadOptInterface master (needs #91)
#     params[:feasibility_pump] = false
#     with_optimizer(Juniper.Optimizer, params)
# end

using BARON
function baron_optimizer_factory()
    with_optimizer(BARON.Optimizer;
         CplexLibName="/opt/ibm/ILOG/CPLEX_Studio128/cplex/bin/x86-64_linux/libcplex1280.so",
         threads=Sys.CPU_THREADS ÷ 2, MaxTime=5 * 60 * 60, PrTimeFreq=5.,
         AllowFilterSD=1, AllowFilterSQP=1, AllowIpopt=1#=, NumLoc=20, LocRes=1=#)
end

using CPLEX
function cplex_optimizer_factory_continuous()
    # https://www.ibm.com/support/knowledgecenter/en/SSSA5P_12.9.0/ilog.odms.cplex.help/CPLEX/Parameters/topics/OptimalityTarget.html

    # Parameters tried:
    # CPX_PARAM_OPTIMALITYTARGET = 2: CPX_OPTIMALITYTARGET_FIRSTORDER - not available for MIP
    # CPX_PARAM_FPHEUR = 1 - not available for MIQCP
    # CPX_PARAM_DEPIND = 3 # dependency checking: turn on at beginning and at end of preprocessing
    # CPX_PARAM_MIQCPSTRAT = 1 - solve QCPs as subproblems
    # CPXPARAM_MIP_Cuts_LiftProj = 3 - Generate lift-and-project cuts very aggressively
    # CPX_PARAM_MIPEMPHASIS = CPLEX.CPX_MIPEMPHASIS_HIDDENFEAS - applies considerable additional effort toward finding high quality feasible solutions that are difficult to locate

    # Parameters to try:
    # CPX_PARAM_MIQCPSTRAT
    # CPX_MIPEMPHASIS_FEASIBILITY
    # barrier-related options
    with_optimizer(CPLEX.Optimizer,
        CPX_PARAM_OPTIMALITYTARGET = CPLEX.CPX_OPTIMALITYTARGET_AUTO,
        CPX_PARAM_DEPIND = 3,
        CPXPARAM_MIP_Cuts_LiftProj = 3)
end

function cplex_optimizer_factory()
    with_optimizer(CPLEX.Optimizer)
end

using Gurobi
function gurobi_optimizer_factory()
    with_optimizer(Gurobi.Optimizer)
end


text:
using PooksoftAlphaVantageDataStore
using Test
using JSON
using DataFrames

# -- User creation/low level download tests --------------------------------------------- #
function build_api_user_model_test()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    user_model = user_model_result.value
    user_email = user_model.alphavantage_api_email

    # is the user email the same as in Configuration.json?
    user_json_dictionary = JSON.parsefile(path_to_config_file)
    alpha_vantage_api_email = user_json_dictionary["user_data"]["alpha_vantage_api_email"]
    if (user_email == alpha_vantage_api_email)
        return true
    end

    # return -
    return false
end

function download_daily_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_daily_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end

function download_daily_adjusted_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_adjusted_daily_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end

function download_weekly_adjusted_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_adjusted_weekly_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end

function download_weekly_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_weekly_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end

function download_monthly_adjusted_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_adjusted_monthly_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end

function download_monthly_appl_sts_test_low_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # make an API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    api_call_result = execute_sts_monthly_api_call(user_model, stock_symbol; data_type = data_type, outputsize = outputsize, logger=nothing)

    # check -
    if (typeof(api_call_result.value) == DataFrame)
        return true
    end

    # return -
    return false
end
#---------------------------------------------------------------------------------------- #

# -- User creation/high level download tests -------------------------------------------- #

function download_daily_appl_sts_test_high_level()

    # initialize -
    my_current_dir = pwd()   # where am I?
    path_to_config_file = my_current_dir*"/configuration/Configuration.json"

    # build the api user model -
    user_model_result = build_api_user_model(path_to_config_file)
    if (typeof(user_model_result.value) == PSError)
        return false
    end

    # get the user model, we'll need this to make an API call -
    user_model = user_model_result.value

    # setup the API call -
    stock_symbol = "aapl"
    data_type = :json
    outputsize = :compact
    call = execute_sts_adjusted_daily_api_call
    apicall_model = build_datastore_apicall_model(call,stock_symbol; output=outputsize, datatype=data_type)
    if (isa(apicall_model.value, PSDataStoreAPICallModel) == false)
        return false
    end

    # make the apicall -
    api_call_result = execute_api_call(user_model, apicall_model.value)

    # check -
    if (isa(api_call_result.value,DataFrame) == true)
        return true
    end

    # return -
    return false
end
#---------------------------------------------------------------------------------------- #

@testset "user_test_set" begin
    # @test build_api_user_model_test_low_level() == true
    # @test download_daily_appl_sts_test_low_level() == true
    # @test download_daily_adjusted_appl_sts_test_low_level() == true
    @test download_weekly_appl_sts_test_low_level() == true
    @test download_weekly_adjusted_appl_sts_test_low_level() == true
    @test download_monthly_appl_sts_test_low_level() == true

    # sleep(60)   # wait 1 min -
    # @test download_monthly_adjusted_appl_sts_test_low_level() == true
    @test download_daily_appl_sts_test_high_level() == true
end


text:
using Base.Test

using ByteParsers: Skip, parsenext, parsedtype, MaybeParsed
using DiscreteRanges: DiscreteRange
using IndirectArrays

using AMDB:
    # errors
    FileError, FileErrors, log_error,
    # autoindexing
    AutoIndex, OrderedCounter,
    # dates
    AMDB_Date, DatePair,
    # utilities
    narrowest_Int, to_narrowest_Int,
    # tuples
    MultiSubs, get_positions

# write your own tests here
@testset "paths" begin
    mktempdir() do dir
        withenv("AMDB_FILES" => dir) do
            @test AMDB.data_directory() == dir
            @test AMDB.data_path("foo") == joinpath(dir, "foo")
            @test AMDB.data_file(2000) ==
                joinpath(dir, "mon_ew_xt_uni_bus_00.csv.gz")
            open(joinpath(dir, "mon_ew_xt_uni_bus.cols.txt"), "w") do io
                write(io, "A;B;C\n")
            end
            @test AMDB.data_colnames() == ["A","B","C"]
        end
    end
end

@testset "error logging and printing" begin
    fe = FileErrors("foo.gz")
    @test repr(fe) == "foo.gz: 0 errors\n"
    @test count(fe) == 0
    log_error(fe, 99, b"bad;bad line", 5)
    @test count(fe) == 1
    @test repr(fe) == "foo.gz: 1 error\nbad;bad line\n    ^ line 99, byte 5\n"
end

@testset "auto indexing" begin
    aa = b"aa"
    bb = b"bb"
    ai = AutoIndex(Vector{UInt8}, Int8)
    @test ai(@view(aa[1:2])) == 1
    @test ai(bb) == 2
    @test ai(aa) == 1
    @test ai(bb) == 2
    @test ai(b"cc") == 3
    @test ai(@view aa[:]) == 1
    @test length(ai) == 3
    @test keys(ai) == [b"aa", b"bb", b"cc"]
    for i in 4:typemax(Int8)
        @test ai(UInt8[1,i]) == i
    end
    @test length(ai) == 127
    # check that we cannot add an element above capacity
    @test_throws AssertionError ai(b"this will not fit")
end

struct AddOne end

@testset "integer narrowing" begin
    ≖(x, y) = false             # also compare types
    ≖(x::Vector{T}, y::Vector{T}) where {T} = isequal(x, y)
    @test to_narrowest_Int([1,2,3]) ≖ Int8[1, 2, 3]
    @test to_narrowest_Int([-129,2,3]) ≖ Int16[-129, 2, 3]
    @test to_narrowest_Int([99, 32768]) ≖ Int32[99, 32768]
end

@testset "multisubs" begin
    m = MultiSubs((1, 3), (x->x^2, x->x^3))
    @test @inferred(m((2, 3, 5))) ≡ (4, 3, 125)
    @test get_positions(m) ≡ (1, 3)
end

@testset "DatePair parsing" begin
    str = b"20010101;20020327;"
    @test parsenext(DatePair(), str, 1, UInt8(';')) ≡
        MaybeParsed(19, DiscreteRange(AMDB_Date(Date(2001, 1, 1)),
                                      AMDB_Date(Date(2002, 3, 27))))
end

@testset "ordered counter" begin
    oc = OrderedCounter{Int, Int32}()
    @test oc(1) == 1
    @test oc(2) == 2
    @test oc(2) == 2
    @test oc(3) == 3
    @test oc(1) == 1
    @test collect(keys(oc)) == Int32[1, 2, 3]
    @test collect(values(oc)) == Int[2, 2, 1]
end

@testset "interning keys" begin
    values = ["foo", "bar", "baz"]
    A = IndirectArray(1:3, values)
    keys = AMDB.intern_keys(A, values)
    all(keys .≡ values)
end


text:
using MESolve

include("test-solver.jl")


text:
# AUTO GENERATED FILE - DO NOT EDIT

export select

"""
    select(;kwargs...)

A Select component.
Custom searchable select. For more information, see: https://mantine.dev/core/select/
Keyword arguments:
- `id` (String; optional): The ID of this component, used to identify dash components in callbacks
- `className` (String; optional): Often used with CSS to style elements with common properties
- `clearable` (Bool; optional): Allow to clear value
- `data` (optional): Select options used to renderer items in dropdown. data has the following type: Array of lists containing elements 'label', 'value', 'disabled'.
Those elements have the following types:
  - `label` (String; required): The option's label
  - `value` (String; required): Option's value
  - `disabled` (Bool; optional): If true, this option is disabled and cannot be selecteds
- `description` (String; optional): Input description, displayed after label
- `disabled` (Bool; optional): The component can show it is currently unable to be interacted with
- `dropdownPosition` (a value equal to: "bottom", "top", "flip"; optional): Dropdown positioning behavior
- `error` (String; optional): Displays error message after input
- `initiallyOpened` (Bool; optional): Initial dropdown opened state
- `label` (String; optional): Input label, displayed before input
- `limit` (Real; optional): Limit amount of items displayed at a time for searchable select
- `maxDropdownHeight` (Real; optional): Maximum dropdown height in px
- `multiline` (Bool; optional): Will input have multiple lines?
- `nothingFound` (String; optional): Nothing found label
- `placeholder` (String; optional): Placeholder, displayed when date is not selected
- `radius` (a value equal to: "xs", "sm", "md", "lg", "xl" | Real; optional): Input border-radius from theme or number to set border-radius in px
- `required` (Bool; optional): Adds red asterisk on the right side of label
- `searchable` (Bool; optional): Set to true to enable search
- `size` (a value equal to: "xs", "sm", "md", "lg", "xl"; optional): Input size
- `style` (Dict; optional): Inline style override
- `value` (String; optional): Selected value
- `withinPortal` (Bool; optional): Whether to render the dropdown in a Portal
- `zIndex` (Real; optional): Dropdown z-index
"""
function select(; kwargs...)
        available_props = Symbol[:id, :className, :clearable, :data, :description, :disabled, :dropdownPosition, :error, :initiallyOpened, :label, :limit, :maxDropdownHeight, :multiline, :nothingFound, :placeholder, :radius, :required, :searchable, :size, :style, :value, :withinPortal, :zIndex]
        wild_props = Symbol[]
        return Component("select", "Select", "dash_mantine_components", available_props, wild_props; kwargs...)
end


text:
###########################
## Portfolio get moments ##
###########################

function getPMean(invs::Investments, mus::DataFrame)
    ## possibly multiple mu vectors
    (nMus, nAss) = size(mus)
    if nMus == 1
        pMean = weights(invs)*array(mus)'
    else                                # multiple mus
        pMean = sum(weights(invs).*array(mus), 2)
    end
    tm = Timematr(pMean, idx(invs))
    return tm
end

#################################
## Portfolio covariance matrix ##
#################################

function getPVar(wgts::Array{Float64, 2}, covMatr::Array{Float64, 2})
    nPfs = size(wgts, 1)
    pVars = Array(Float64, nPfs)
    for ii=1:nPfs
        pVars[ii, 1] = (wgts[ii, :]*covMatr*(wgts[ii, :])')[1]
    end
    return pVars
end

function getPVar(invs::Investments, covMatr::DataFrame)
    pVars = getPVar(core(invs), array(covMatr))
    return Timematr(pVars, idx(invs))
end


text:
# fbp2_back.jl

export fbp2_back

using LazyGrids: ndgrid
using ImageGeoms: embed

"""
    img = fbp2_back(sg, ig, sino; ia_skip)

2D backprojection for FBP.

in
- `sg::SinoGeom`
- `ig::ImageGeom`
- `sino::AbstractArray{<:Number}`      sinogram(s) (line integrals), usually ramp filtered

options
- `ia_skip::Int`                        downsample in angle to save time for quick tests (default: 1)

out
- `img::AbstractArray{<:Number}`       reconstructed image(s)
"""
function fbp2_back(
    sg::SinoPar,
    ig::ImageGeom,
    sino::AbstractMatrix{<:Number};
    ia_skip::Int = 1,
    do_r_mask::Bool = false
)

    # trick: extra zero column saves linear interpolation indexing within loop!

    nb = size(sino,1) # number of radial bins
    nb != sg.nb && throw("nb size")
    sino = [sino; zeros(eltype(sino),size(sino,2),2)']

    xc, yc = ndgrid(ig.x, ig.y)
    rr = sqrt.(abs2.(xc) + abs2.(yc)) # [nx ny]
    rmax = ((sg.nb - 1) / 2 - abs(sg.offset)) * sg.d
    mask = ig.mask
    if do_r_mask
        mask = mask .& (rr .< rmax)
    end
    xc = xc[vec(mask)] # [np] pixels within mask
    yc = yc[vec(mask)]

    sang = sin.(sg.ar)
    cang = cos.(sg.ar)

    img = 0
    for ia = 1:ia_skip:sg.na

        rr = @.(xc * cang[ia] + yc * sang[ia]) # [np]
        rr = @.(rr / sg.d + sg.w + 1) # unitless bin index, +1 because julia

        # nearest neighbor interpolation:
#=
    %    ib = round(bb);
    %    if any(ib < 1 | ib > nb), error 'bug', end
    %    % trick: make out-of-sinogram indices point to those extra zeros
    %    ib(ib < 1 | ib > nb) = nb+1;
    %    img = img + sino(ib, ia) ./ L2;
=#
        # linear interpolation:
        il = floor.(Int64, rr) # left bin

        if !do_r_mask
            il = min.(il,nb+1)
            il = max.(il,1)
        end
        # % (any(il .< 1) || any(il .>= nb)) && throw("bug")

        wr = rr - il # left weight
        #@show extrema(wr), rr[1000:1005], il[1000:1005]
        wl = 1 .- wr # right weight
        img = @.(img + wl * sino[il, ia] + wr * sino[il.+1, ia])
    end

    # img = (deg2rad(sg.orbit) / (sg.na/ia_skip)) * embed(img, mask);
    # img = pi / (sg.na/ia_skip) * embed(img, mask) % 2008-10-14
    return π / (sg.na/ia_skip) * embed(img, mask)
end

# fan-beam case
function fbp2_back(
    sg::SinoFan,
    ig::ImageGeom,
    sino::AbstractMatrix{<:Number};
    ia_skip::Int = 1,
)

    if sg.dfs == 0
        is_arc=true
    elseif isinf(sg.dfs)
        is_arc=false
    else
        throw("bad dsf")
    end

    return fbp2_back(
        sino, sg.orbit, sg.orbit_start,
        sg.dsd, sg.dso, sg.dfs, sg.ds, sg.offset,
        sg.source_offset,
        ig.nx, ig.ny, ig.dx, ig.dy, ig.offset_x, ig.offset_y,
        is_arc, ig.mask, ia_skip,
    )
end

function fbp2_back(
    sino::AbstractMatrix{<:Number},
    orbit::Union{Symbol,Real}, orbit_start::Real,
    dsd::RealU,
    dso::RealU,
    dfs::RealU,
    ds::RealU,
    offset::Real,
    source_offset::RealU,
    nx::Int,
    ny::Int,
    dx::RealU,
    dy::RealU,
    offset_x::Real,
    offset_y::Real,
    is_arc::Bool,
    mask::AbstractMatrix{Bool},
    ia_skip::Int,
)

    rmax = []

    nb,na = size(sino)

    # trick: extra zero column saves linear interpolation indexing within loop!

    sino = [sino; zeros(size(sino,2))']

    # precompute as much as possible
    wx = (nx+1)/2 - offset_x
    wy = (ny+1)/2 - offset_y
    xc, yc = ndgrid(dx * ((1:nx) .- wx), dy * ((1:ny) .- wy))
    rr = @.(sqrt(abs2(xc) + abs2(yc))) # [nx,ny]

    smax = ((nb-1)/2 - abs(offset)) * ds

    #todo: rmax possibly extractable from SinoGeom
    if is_arc
        gamma_max = smax / dsd
    else # flat
        gamma_max = atan(smax / dsd)
    end
    rmax = dso * sin(gamma_max)

    mask = mask .& (rr .< rmax)
    xc = xc[vec(mask)] # [np] pixels within mask
    yc = yc[vec(mask)]
    #clear wx wy rr smax

    betas = @.(deg2rad(orbit_start + orbit * (0:na-1) / na)) # [na]
    wb = (nb+1)/2 + offset

    img = 0

    for ia=1:ia_skip:na
    #ticker(mfilename, ia, na)

        beta = betas[ia]
        d_loop = @.(dso + xc * sin(beta) - yc * cos(beta)) # dso - y_beta
        r_loop = @.(xc * cos(beta) + yc * sin(beta) - source_offset) # x_beta-roff

        if is_arc
            sprime_ds = (dsd/ds) * atan.(r_loop, d_loop) # s' / ds
            w2 = @.(abs2(dsd) / (abs2(d_loop) + abs2(r_loop))) # [np] image weighting
        else # flat
            mag = dsd ./ d_loop
            sprime_ds = mag .* r_loop / ds
            w2 = abs2.(mag) # [np] image-domain weighting
        end

        bb = sprime_ds .+ wb # [np] bin "index"

        # nearest neighbor interpolation:
#=
    %    ib = round(bb);
    %    if any(ib < 1 | ib > nb), error 'bug', end
    %    % trick: make out-of-sinogram indices point to those extra zeros
    %%    ib(ib < 1 | ib > nb) = nb+1;
    %    img = img + sino(ib, ia) ./ L2;
=#

        # linear interpolation:
        il = floor.(Int, bb) # left bin
        ir = 1 .+ il # right bin

        # deal with truncated sinograms
        ig = (il .>= 1) .& (ir .<= nb)
        il[.!ig] .= nb+1
        ir[.!ig] .= nb+1
   #    if any(il < 1 | il >= nb), error 'bug', end

        wr = bb .- il # left weight
        wl = 1 .- wr # right weight

        img = @.(img + (wl * sino[il, ia] + wr * sino[ir, ia]) * w2)
    end

    return π / (na/ia_skip) * embed(img,mask)
end

# 3d support
function fbp2_back(sg::SinoGeom, ig::ImageGeom, sino::AbstractArray{<:Number}; kwargs...)
    return mapslices(sino -> fbp2_back(sg, ig, sino; kwargs...), sino, [1,2])
end


text:
export
    DriverModel,

    get_name,
    action_type,
    action_context,
    set_desired_speed!,
    observe!,
    reset_hidden_state!,
    prime_with_history!

abstract type DriverModel{DriveAction, ActionContext} end

get_name(::DriverModel) = "???"
action_type{A,C}(::DriverModel{A, C}) = A
action_context{A,C}(model::DriverModel{A, C}) = error("action_context not implemented for model $model")
set_desired_speed!(model::DriverModel, v_des::Float64) = model # do nothing by default
reset_hidden_state!(model::DriverModel) = model # do nothing by default
observe!(model::DriverModel, scene::Scene, roadway::Roadway, egoid::Int) = model  # do nothing by default
Base.rand{A}(model::DriverModel{A}) = error("rand not implemented for model $model")
Distributions.pdf{A}(model::DriverModel{A}, a::A) = error("pdf not implemented for model $model")
Distributions.logpdf{A}(model::DriverModel{A}, a::A) = error("logpdf not implemented for model $model")

function prime_with_history!(model::DriverModel, trajdata::Trajdata, roadway::Roadway, frame_start::Int, frame_end::Int, egoid::Int, scene::Scene=Scene())

    reset_hidden_state!(model)

    for frame in frame_start : frame_end
        get!(scene, trajdata, frame)
        observe!(model, scene, roadway, egoid)
    end

    model
end
function prime_with_history!(model::DriverModel, rec::SceneRecord, roadway::Roadway, egoid::Int;
    pastframe_start::Int=1-length(rec),
    pastframe_end::Int=0,
    )

    reset_hidden_state!(model)

    for pastframe in pastframe_start : pastframe_end
        scene = get_scene(rec, pastframe)
        observe!(model, scene, roadway, egoid)
    end

    model
end


text:
push!(LOAD_PATH, joinpath(homedir(),"GitHub/KenyaCoV/src"))
using Plots,Parameters,Distributions,DifferentialEquations,JLD2,DataFrames,StatsPlots,FileIO,MAT,RecursiveArrayTools
using Revise
import KenyaCoV
using LinearAlgebra:eigen
using Statistics: median, quantile

gr()#Plotting frontend
"""
IMPERFECT social distancing + spatial lockdown + decrease in effective contacts by 50% over 30 days
All households reduce contact outside household or workplace by 75%.
School contact rates set to zero, workplace contact rates reduced by 25%.
Household contact rates assumed to increase by 25%.
All spatial movements reduced to 0.1% of time spent in other areas
"""

"""
Load age structured data, define initial state and declare the KenyaCoV problem
"""

u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
T_normal = deepcopy(P.T)

@load "data/detection_rates_for_different_taus.jld2" d_0 d_01 d_025 d_05 d_1
@load "data/susceptibility_rates.jld2" σ

#Initial infecteds
"""
Load age mixing matrices (these are all in to (row) from (col) format)
"""

@load "data/agemixingmatrix_Kenya_all_types.jld2" M_Kenya M_Kenya_ho M_Kenya_other M_Kenya_school M_Kenya_work

"""
Declare the treatment/isolation rates considered, and a callback that limits intervention
"""
treatment_rates = [(0.,1),(0.,0.5),(1/7.,1.),(1/7,0.5),(1/3.5,1.),(1/3.5,0.5)]
reduced_treatment_rates = [(1/3.5,0.5)]
function ramp_down(t)
    if t < 30.
        return (1-t/30) + 0.5*t/30
    else
        return 0.5
    end
end

function isolation_limit(u,t,integrator) # Isolation can only continue until 1000 detected cases
  integrator.p.isolating_detecteds && sum(u[:,:,8]) > 1e3
end
function affect_isolation_limit!(integrator)
  integrator.p.τ = 0
  integrator.p.isolating_detecteds = false
end
cb_iso_limit = DiscreteCallback(isolation_limit,affect_isolation_limit!)

function social_distancing_limit(u,t,integrator) # Isolation can only continue until 1000 detected cases
  integrator.p.lockdown && t > 90.
end
function affect_social_distancing_limit!(integrator)
  integrator.p.M = M_Kenya
  integrator.p.T = T_normal
  integrator.p.lockdown = false
  integrator.p.c_t = t -> 1.
end

cb_SD_limit = DiscreteCallback(social_distancing_limit,affect_social_distancing_limit!)

both_cbs = CallbackSet(cb_iso_limit,cb_SD_limit)

"""
SCENARIO 1 --- Imperfect social distancing + spatial lockdown for 3 months:
* No age-specific susceptibilties --- its disease difference
* Asymptomatics are 0% as infectious as symptomatics
"""

u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
#Redistribute susceptibility JUST to rescale infectiousness so we get the correct R₀/r
P.χ = ones(KenyaCoV.n_a)
P.rel_detection_rate = d_0
P.dt = 0.25;
P.ext_inf_rate = 0.;
P.ϵ = 0.
P.γ = 1/2.5
R₀_scale = KenyaCoV.calculate_R₀_scale(P)
P.χ = ones(KenyaCoV.n_a)/R₀_scale

P.σ = 1/rand(KenyaCoV.d_incubation)
P.β = 2.5*P.γ

#Choosing initial conditions - scenario 1
rand_chosen_R,eq_case_distribution,fifth_gen_case_distribution = KenyaCoV.calculate_R₀(P)
num_initial_detected = ceil(Int64,sum(fifth_gen_case_distribution.*d_0))
detected_age_group = argmax(fifth_gen_case_distribution.*d_0)
undetected_age_profile = round.(Int64,fifth_gen_case_distribution.*(1 .- d_0))

#Implement social distancing + spatial lockdown
KenyaCoV.calculate_R₀(P)
P.M = M_Kenya_ho*(1.25) .+ M_Kenya_work*(0.75) .+ M_Kenya_other*(0.25)
KenyaCoV.calculate_R₀(P)
KenyaCoV.transportstructure_params!(P,[0.001 for i = 1:KenyaCoV.n_wa],P_dest)
P.c_t = ramp_down

u0[KenyaCoV.ind_nairobi_as,detected_age_group,4] = num_initial_detected
u0[KenyaCoV.ind_nairobi_as,:,3] = undetected_age_profile
prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,2*365.),P)

results_1SI_SL_DC_three_months = KenyaCoV.run_scenario(P,prob,1000,reduced_treatment_rates,both_cbs)
@save joinpath(homedir(),"Github/KenyaCoVOutputs/results_1SI_SL_DC_three_months.jld2") results_1SI_SL_DC_three_months

println("Finished 1 imperfect social distancing + spatial lockdown for 3 months")

"""
SCENARIO 2 --- Imperfect social distancing + spatial lockdown:
* No age-specific susceptibilties --- its disease difference
* Asymptomatics are 10% as infectious as symptomatics
"""

u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
#Redistribute susceptibility JUST to rescale infectiousness so we get the correct R₀/r
P.χ = ones(KenyaCoV.n_a)
P.rel_detection_rate = d_01
P.dt = 0.25;
P.ext_inf_rate = 0.;
P.ϵ = 0.1
P.γ = 1/2.5
R₀_scale = KenyaCoV.calculate_R₀_scale(P)
P.χ = ones(KenyaCoV.n_a)/R₀_scale

P.σ = 1/rand(KenyaCoV.d_incubation)
P.β = 2.5*P.γ

#Choosing initial conditions - scenario 2
rand_chosen_R,eq_case_distribution,fifth_gen_case_distribution = KenyaCoV.calculate_R₀(P)
num_initial_detected = ceil(Int64,sum(fifth_gen_case_distribution.*d_01))
detected_age_group = argmax(fifth_gen_case_distribution.*d_01)
undetected_age_profile = round.(Int64,fifth_gen_case_distribution.*(1 .- d_01))

#Implement social distancing + spatial lockdown
KenyaCoV.calculate_R₀(P)
P.M = M_Kenya_ho*(1.25) .+ M_Kenya_work*(0.75) .+ M_Kenya_other*(0.25)
KenyaCoV.calculate_R₀(P)
KenyaCoV.transportstructure_params!(P,[0.001 for i = 1:KenyaCoV.n_wa],P_dest)
P.c_t = ramp_down

u0[KenyaCoV.ind_nairobi_as,detected_age_group,4] = num_initial_detected
u0[KenyaCoV.ind_nairobi_as,:,3] = undetected_age_profile
prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,2*365.),P)

results_2SI_SL_DC_three_months = KenyaCoV.run_scenario(P,prob,1000,reduced_treatment_rates,both_cbs)
@save joinpath(homedir(),"Github/KenyaCoVOutputs/results_2SI_SL_DC_three_months.jld2") results_2SI_SL_DC_three_months

println("Finished 2 imperfect social distancing + spatial lockdown for 3 months")
"""
SCENARIO 3 --- social distancing + spatial lockdown:
* No age-specific susceptibilties --- its disease difference
* Asymptomatics are 25% as infectious as symptomatics
"""
u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
#Redistribute susceptibility JUST to rescale infectiousness so we get the correct R₀/r
P.χ = ones(KenyaCoV.n_a)
P.rel_detection_rate = d_025
P.dt = 0.25;
P.ext_inf_rate = 0.;
P.ϵ = 0.25
P.γ = 1/2.5
R₀_scale = KenyaCoV.calculate_R₀_scale(P)
P.χ = ones(KenyaCoV.n_a)/R₀_scale

P.σ = 1/rand(KenyaCoV.d_incubation)
P.β = 2.5*P.γ

#Choosing initial conditions - scenario 3
rand_chosen_R,eq_case_distribution,fifth_gen_case_distribution = KenyaCoV.calculate_R₀(P)
num_initial_detected = ceil(Int64,sum(fifth_gen_case_distribution.*d_025))
detected_age_group = argmax(fifth_gen_case_distribution.*d_025)
undetected_age_profile = round.(Int64,fifth_gen_case_distribution.*(1 .- d_025))

#Implement social distancing + spatial lockdown
KenyaCoV.calculate_R₀(P)
P.M = M_Kenya_ho*(1.25) .+ M_Kenya_work*(0.75) .+ M_Kenya_other*(0.25)
KenyaCoV.calculate_R₀(P)
KenyaCoV.transportstructure_params!(P,[0.001 for i = 1:KenyaCoV.n_wa],P_dest)
P.c_t = ramp_down

u0[KenyaCoV.ind_nairobi_as,detected_age_group,4] = num_initial_detected
u0[KenyaCoV.ind_nairobi_as,:,3] = undetected_age_profile
prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,2*365.),P)

results_3SI_SL_DC_three_months = KenyaCoV.run_scenario(P,prob,1000,reduced_treatment_rates,both_cbs)
@save joinpath(homedir(),"Github/KenyaCoVOutputs/results_3SI_SL_DC_three_months.jld2") results_3SI_SL_DC_three_months

println("Finished 3 imperfect social distancing + spatial lockdown for 3 months")

"""
SCENARIO 4 --- social distancing + spatial lockdown for 3 months:
* No age-specific susceptibilties --- its disease difference
* Asymptomatics are 50% as infectious as symptomatics
"""
u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
#Redistribute susceptibility JUST to rescale infectiousness so we get the correct R₀/r
P.χ = ones(KenyaCoV.n_a)
P.rel_detection_rate = d_05
P.dt = 0.25;
P.ext_inf_rate = 0.;
P.ϵ = 0.5
P.γ = 1/2.5
R₀_scale = KenyaCoV.calculate_R₀_scale(P)
P.χ = ones(KenyaCoV.n_a)/R₀_scale

P.σ = 1/rand(KenyaCoV.d_incubation)
P.β = 2.5*P.γ

#Choosing initial conditions - scenario 4
rand_chosen_R,eq_case_distribution,fifth_gen_case_distribution = KenyaCoV.calculate_R₀(P)
num_initial_detected = ceil(Int64,sum(fifth_gen_case_distribution.*d_05))
detected_age_group = argmax(fifth_gen_case_distribution.*d_05)
undetected_age_profile = round.(Int64,fifth_gen_case_distribution.*(1 .- d_05))

#Implement social distancing + spatial lockdown
KenyaCoV.calculate_R₀(P)
P.M = M_Kenya_ho*(1.25) .+ M_Kenya_work*(0.75) .+ M_Kenya_other*(0.25)
KenyaCoV.calculate_R₀(P)
KenyaCoV.transportstructure_params!(P,[0.001 for i = 1:KenyaCoV.n_wa],P_dest)
P.c_t = ramp_down

u0[KenyaCoV.ind_nairobi_as,detected_age_group,4] = num_initial_detected
u0[KenyaCoV.ind_nairobi_as,:,3] = undetected_age_profile
prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,2*365.),P)

results_4SI_SL_DC_three_months = KenyaCoV.run_scenario(P,prob,1000,reduced_treatment_rates,both_cbs)
@save joinpath(homedir(),"Github/KenyaCoVOutputs/results_4SI_SL_DC_three_months.jld2") results_4SI_SL_DC_three_months

println("Finished 4 imperfect social distancing + spatial lockdown for 3 months")

"""
SCENARIO 5 --- social distancing + spatial lockdown for 3 months:
* No age-specific susceptibilties --- its disease difference
* Asymptomatics are 100% as infectious as symptomatics
"""
u0,P,P_dest = KenyaCoV.model_ingredients_from_data("data/data_for_age_structuredmodel.jld2",
                                            "data/flight_numbers.csv",
                                            "data/projected_global_prevelance.csv")
#Redistribute susceptibility JUST to rescale infectiousness so we get the correct R₀/r
P.χ = ones(KenyaCoV.n_a)
P.rel_detection_rate = d_1
P.dt = 0.25;
P.ext_inf_rate = 0.;
P.ϵ = 1.
P.γ = 1/2.5
R₀_scale = KenyaCoV.calculate_R₀_scale(P)
P.χ = ones(KenyaCoV.n_a)/R₀_scale

P.σ = 1/rand(KenyaCoV.d_incubation)
P.β = 2.5*P.γ

#Choosing initial conditions - scenario 5
rand_chosen_R,eq_case_distribution,fifth_gen_case_distribution = KenyaCoV.calculate_R₀(P)
num_initial_detected = ceil(Int64,sum(fifth_gen_case_distribution.*d_1))
detected_age_group = argmax(fifth_gen_case_distribution.*d_1)
undetected_age_profile = round.(Int64,fifth_gen_case_distribution.*(1 .- d_1))

#Implement social distancing + spatial lockdown
KenyaCoV.calculate_R₀(P)
P.M = M_Kenya_ho*(1.25) .+ M_Kenya_work*(0.75) .+ M_Kenya_other*(0.25)
KenyaCoV.calculate_R₀(P)
KenyaCoV.transportstructure_params!(P,[0.001 for i = 1:KenyaCoV.n_wa],P_dest)
P.c_t = ramp_down

u0[KenyaCoV.ind_nairobi_as,detected_age_group,4] = num_initial_detected
u0[KenyaCoV.ind_nairobi_as,:,3] = undetected_age_profile
prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,2*365.),P)

results_5SI_SL_DC_three_months = KenyaCoV.run_scenario(P,prob,1000,reduced_treatment_rates,both_cbs)
@save joinpath(homedir(),"Github/KenyaCoVOutputs/results_5SI_SL_DC_three_months.jld2") results_5SI_SL_DC_three_months

println("Finished 5 imperfect social distancing + spatial lockdown for 3 months")


text:
module FESpacesTests2

using Test

@testset "FESpacesWithConstantFixed" begin include("FESpacesWithConstantFixedTests.jl") end

@testset "ZeroMeanFESpaces" begin include("ZeroMeanFESpacesTests.jl") end

@testset "CLagrangianFESpaces" begin include("CLagrangianFESpacesTests.jl") end

@testset "DirichletFESpaces" begin include("DirichletFESpacesTests.jl") end

@testset "ExtendedFESpaces" begin include("ExtendedFESpacesTests.jl") end

@testset "FESpacesWithLinearConstraints" begin include("FESpacesWithLinearConstraintsTests.jl") end

@testset "FEAutodiff" begin include("FEAutodiffTests.jl") end

@testset "FETermsWithAutodiff" begin include("FETermsWithAutodiffTests.jl") end

@testset "CellDofBases" begin include("CellDofBasesTests.jl") end

@testset "Law" begin include("LawTests.jl") end

@testset "CDLagrangianFESpaces" begin include("CDLagrangianFESpacesTests.jl") end

@testset "PhysicalBasesTests" begin include("PhysicalBasesTests.jl") end

@testset "PhysicalFESpaces" begin include("PhysicalFESpacesTests.jl") end

@testset "FESpaceFactories" begin include("FESpaceFactoriesTests.jl") end

@testset "AppendedTriangulations" begin include("AppendedTriangulationsTests.jl") end

end # module