#=
    Used from https://github.com/lanl-ansi/Juniper.jl
=#

"""
MOI functions, sets and, other type definitions
"""
# indices
const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

# sets
const BOUNDS = Union{
    MOI.EqualTo{Float64},
    MOI.GreaterThan{Float64},
    MOI.LessThan{Float64},
}

const VAR_TYPES = Union{
    MOI.ZeroOne,
    MOI.Integer
}

# other MOI types
const AFF_TERM = MOI.ScalarAffineTerm{Float64}
const QUAD_TERM = MOI.ScalarQuadraticTerm{Float64}

"""
Variable information struct definition
"""
mutable struct VariableInfo
    lower_bound::Float64  # May be -Inf even if has_lower_bound == true
    has_lower_bound::Bool # false implies lower_bound == -Inf
    upper_bound::Float64  # May be Inf even if has_upper_bound == true
    has_upper_bound::Bool # false implies upper_bound == Inf
    is_fixed::Bool        # Implies lower_bound == upper_bound and !has_lower_bound and !has_upper_bound
    is_binary::Bool       # Implies lower_bound == 0, upper_bound == 1 and is MOI.ZeroOne
    is_integer::Bool      # Implies variable is MOI.Integer
    start::Real           # Primal start
    name::String
end
VariableInfo() = VariableInfo(-Inf, false, Inf, false, false, false, false, 0.0, "")

"""
Optimizer struct
"""
mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::Union{OriginalProblem, Nothing}
    variable_info::Vector{VariableInfo}
    nlp_data::MOI.NLPBlockData
    sense::MOI.OptimizationSense
    objective::Union{SVF, SAF, SQF, Nothing}
    linear_le_constraints::Vector{Tuple{SAF, MOI.LessThan{Float64}}}
    linear_ge_constraints::Vector{Tuple{SAF, MOI.GreaterThan{Float64}}}
    linear_eq_constraints::Vector{Tuple{SAF, MOI.EqualTo{Float64}}}
    quadratic_le_constraints::Vector{Tuple{SQF, MOI.LessThan{Float64}}}
    quadratic_ge_constraints::Vector{Tuple{SQF, MOI.GreaterThan{Float64}}}
    quadratic_eq_constraints::Vector{Tuple{SQF, MOI.EqualTo{Float64}}}
    #one field needs to be added to keep track of nlp constraint types
    options::SolverOptions
end

MOI.get(::Optimizer, ::MOI.SolverName) = "DuranSolver"

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.supports(::Optimizer, ::MOI.NumberOfThreads) = true
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::Optimizer, ::MOI.RawParameter) = true

function MOI.set(model::Optimizer, ::MOI.Silent, value::Bool)
    if value
        model.options.log_levels = []
    end
    model.options.silent = value
    return
end

function MOI.set(model::Optimizer, ::MOI.NumberOfThreads, value::Union{Nothing,Int})
    if value === nothing
        model.options.processors = 1
    else
        model.options.processors = value
    end
    return
end

function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, value::Union{Nothing,Float64})
    if value === nothing
        model.options.time_limit = Inf
    else
        model.options.time_limit = value
    end
    return
end
#this function needs to be updated
function MOI.set(model::Optimizer, p::MOI.RawParameter, value)
    p_symbol = Symbol(p.name)
    if in(p_symbol, fieldnames(SolverOptions))
        type_of_param = fieldtype(SolverOptions, p_symbol)
        if hasmethod(convert, (Type{type_of_param}, typeof(value)))
            passed_checks = true
            if p_symbol == :traverse_strategy
                if !(value in [:BFS, :DFS, :DBFS])
                    passed_checks = false
                    @error "Traverse strategy $(value) is not supported. Use one of `[:BFS, :DFS, :DBFS]`."
                end
            end
            if p_symbol == :branch_strategy
                if !(value in [:StrongPseudoCost, :PseudoCost, :Reliability, :MostInfeasible])
                    passed_checks = false
                    @error "Branch strategy $(value) is not supported. Use one of `[:StrongPseudoCost, :PseudoCost, :Reliability, :MostInfeasible]`."
                end
            end
            passed_checks && setfield!(model.options, p_symbol, convert(type_of_param, value))
        else
            @error "The option $(p.name) has a different type ($(type_of_param))"
        end
    else
        @error "The option $(p.name) doesn't exist."
    end
    return
end

# Returns the number of solutions that can be retrieved
MOI.get(model::Optimizer, ::MOI.ResultCount) = length(model.inner.solutions)

MOI.get(model::Optimizer, ::MOI.NumberOfThreads) = model.options.processors

MOI.get(model::Optimizer, ::MOI.TimeLimitSec) = model.options.time_limit

MOI.get(model::Optimizer, ::MOI.Silent) = model.options.silent

function MOI.get(model::Optimizer, p::MOI.RawParameter)
    if in(p.name, fieldnames(SolverOptions))
        return getfield(model.options, p.name)
    end
    @error "The option $(p.name) doesn't exist."
end

"""
EmptyNLPEvaluator struct and associated functions
"""
struct EmptyNLPEvaluator <: MOI.AbstractNLPEvaluator end
MOI.features_available(::EmptyNLPEvaluator) = [:Grad, :Jac, :Hess]
MOI.initialize(::EmptyNLPEvaluator, features) = nothing
MOI.eval_objective(::EmptyNLPEvaluator, x) = NaN
function MOI.eval_constraint(::EmptyNLPEvaluator, g, x)
    @assert length(g) == 0
    return
end
function MOI.eval_objective_gradient(::EmptyNLPEvaluator, g, x)
    fill!(g, 0.0)
    return
end
MOI.jacobian_structure(::EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
MOI.hessian_lagrangian_structure(::EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
function MOI.eval_constraint_jacobian(::EmptyNLPEvaluator, J, x)
    @assert length(J) == 0
    return
end
function MOI.eval_hessian_lagrangian(::EmptyNLPEvaluator, H, x, σ, μ)
    @assert length(H) == 0
    return
end
empty_nlp_data() = MOI.NLPBlockData([], EmptyNLPEvaluator(), false)

function register(s::Symbol, dimension::Integer, f::Function; autodiff::Bool=false)
    return RegisteredFunction(s, dimension, f, nothing, nothing, autodiff)
end

function register(s::Symbol, dimension::Integer, f::Function, gradf::Function; autodiff::Bool=false)
    return RegisteredFunction(s, dimension, f, gradf, nothing, autodiff)
end

function register(s::Symbol, dimension::Integer, f::Function, gradf::Function, grad2f::Function; autodiff::Bool=false)
    return RegisteredFunction(s, dimension, f, gradf, grad2f, autodiff)
end

"""
Optimizer struct constructor
"""
function Optimizer(;options...)

    solver_options = combine_options(options)

    return Optimizer(
    nothing,
    [],
    empty_nlp_data(),
    MOI.FEASIBILITY_SENSE,
    nothing,
    [], [], [], # linear constraints
    [], [], [], # quadratic constraints
    solver_options)
end

function Optimizer(options::Vector{Pair{String,Any}})
    symbol_options = Dict{Symbol, Any}()
    for option in options
        symbol_options[Symbol(option.first)] = option.second
    end
    return Optimizer(symbol_options)
end

function Optimizer(options::Dict{Symbol,Any})

    solver_options = combine_options(options)

    return Optimizer(
    nothing,
    [],
    empty_nlp_data(),
    MOI.FEASIBILITY_SENSE,
    nothing,
    [], [], [], # linear constraints
    [], [], [], # quadratic constraints
    solver_options)
end

"""
Printing the optimizer
"""
function Base.show(io::IO, model::Optimizer)
    println("A HybridMINLPSolver MathOptInterface model with backend")
    return
end

"""
Copy constructor for the optimizer
"""
MOIU.supports_default_copy_to(model::Optimizer, copy_names::Bool) = true
function MOI.copy_to(model::Optimizer, src::MOI.ModelLike; kws...)
    return MOI.Utilities.automatic_copy_to(model, src; kws...)
end

"""
``MOI.is_empty(model::Optimizer)`` overload for Alpine.Optimizer
"""
function MOI.is_empty(model::Optimizer)
    return isempty(model.variable_info) &&
        model.nlp_data.evaluator isa EmptyNLPEvaluator &&
        model.sense == MOI.FEASIBILITY_SENSE &&
        isempty(model.linear_le_constraints) &&
        isempty(model.linear_ge_constraints) &&
        isempty(model.linear_eq_constraints) &&
        isempty(model.quadratic_le_constraints) &&
        isempty(model.quadratic_ge_constraints) &&
        isempty(model.quadratic_eq_constraints)
end

"""
``MOI.empty!(model::Optimizer)`` overload for Alpine.Optimizer
"""
function MOI.empty!(model::Optimizer)
    model.inner = nothing
    empty!(model.variable_info)
    model.nlp_data = empty_nlp_data()
    model.sense = MOI.FEASIBILITY_SENSE
    model.objective = nothing
    empty!(model.linear_le_constraints)
    empty!(model.linear_ge_constraints)
    empty!(model.linear_eq_constraints)
    empty!(model.quadratic_le_constraints)
    empty!(model.quadratic_ge_constraints)
    empty!(model.quadratic_eq_constraints)
end

"""
ordering of constraints provided to HybridMINLPSolver.jl
"""
linear_le_offset(model::Optimizer) = 0
linear_ge_offset(model::Optimizer) = length(model.linear_le_constraints)
linear_eq_offset(model::Optimizer) = linear_ge_offset(model) + length(model.linear_ge_constraints)
quadratic_le_offset(model::Optimizer) = linear_eq_offset(model) + length(model.linear_eq_constraints)
quadratic_ge_offset(model::Optimizer) = quadratic_le_offset(model) + length(model.quadratic_le_constraints)
quadratic_eq_offset(model::Optimizer) = quadratic_ge_offset(model) + length(model.quadratic_ge_constraints)
nlp_constraint_offset(model::Optimizer) = quadratic_eq_offset(model) + length(model.quadratic_eq_constraints)

function info_array_of_variables(variable_info::Vector{VariableInfo}, attr::Symbol)
    len_var_info = length(variable_info)
    type_dict = get_type_dict(variable_info[1])
    result = Array{type_dict[attr], 1}(undef, len_var_info)
    for i = 1:len_var_info
        result[i] = getfield(variable_info[i], attr)
        # if type is binary then set bounds correctly
        if result[i] < 0 && attr == :lower_bound && getfield(variable_info[i], :is_binary)
            result[i] = 0
        end
        if result[i] > 1 && attr == :upper_bound && getfield(variable_info[i], :is_binary)
            result[i] = 1
        end
    end
    return result
end

"""
``MOI.optimize!()`` for Duran Solver
"""
function MOI.optimize!(model::Optimizer)
    MOI.initialize(model.nlp_data.evaluator, [:Grad, :Jac, :ExprGraph])
    if ~isa(model.nlp_data.evaluator, EmptyNLPEvaluator)
    else
        @info "no explicit NLP constraints or objective provided using @NLconstraint or @NLobjective macros"
    end
    model.inner = OriginalProblem()
    @views op = model.inner
    #initiate to replicate original problem
    init_original_problem!(op, model)
    #complete original problem
    create_original_model!(model, op)
    print("This is the original problem \n", op.model)
    #perform outer approximation algorithm
    oa_algorithm(model, jp)
end

getnsolutions(m::OriginalProblem) = m.nsolutions
getsolutions(m::OriginalProblem) = m.solutions

include("moi_wrapper.jl")


using Test
using ChemistryFeaturization.Utils.GraphBuilding
using Xtals

@testset "GraphBuilding" begin
    path1 = abspath(@__DIR__, "..", "test_data", "strucs", "mp-195.cif")
    adj, els = build_graph(path1; use_voronoi = true)
    wm_true = [0.0 1.0 1.0 1.0; 1.0 0.0 1.0 1.0; 1.0 1.0 0.0 1.0; 1.0 1.0 1.0 0.0]
    els_true = ["Ho", "Pt", "Pt", "Pt"]

    @test adj == wm_true
    @test els == els_true

    adj, els = build_graph(path1; use_voronoi = false)
    @test adj == wm_true
    @test els == els_true

    # test that we get the same results building from a Crsytal object
    adjc, elsc = build_graph(Crystal(path1))
    @test adjc == wm_true
    @test elsc == els_true
end


function parser_cvrplib_unit_tests()
    parse_cvrplib_tests()
end

function parse_cvrplib_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)*
                              "/../../data/cvrplib/P/P\\P-n16-k8.vrp")

    @test typeof(data) == RVRP.RvrpInstance
    @test data.id == "P-n16-k8"
    @test data.travel_specifications[1].id == "unique_period_cat"
    travel_time_mat = data.travel_specifications[1].travel_time_matrix
    travel_distance_mat = data.travel_specifications[1].travel_distance_matrix
    energy_consum_mat = data.travel_specifications[1].energy_consumption_matrix
    @test size(travel_distance_mat) == (16, 16)
    for i in 1:16
        @test travel_distance_mat[i,i] == 0.0
        for j in i+1:16
            @test travel_distance_mat[i,j] == travel_distance_mat[j,i]
        end
    end
    @test travel_time_mat == Array{Float64,2}(undef, 0, 0)
    @test energy_consum_mat == Array{Float64,2}(undef, 0, 0)
    @test length(data.locations) == 16
    @test data.locations[1].id == "depot"
    @test length(data.location_groups) == length(data.locations)
    for i in 1:length(data.locations)
        @test(string(data.locations[i].id, "_loc_group") ==
                data.location_groups[i].id)
    end
    @test length(data.product_compatibility_classes) == 1
    @test data.product_compatibility_classes[1].id == "default_id"
    @test length(data.product_sharing_classes) == 1
    @test data.product_sharing_classes[1].id == "default_id"
    @test length(data.requests) == 15
    @test length(data.vehicle_categories) == 2
    @test length(data.vehicle_sets) == 1
    @test data.vehicle_categories[1].id == "unique_vehicle_category"
    @test haskey(data.vehicle_categories[1].capacities.of_vehicle, "unique_measure")
    @test(data.vehicle_categories[1].capacities.of_vehicle ==
            Dict{String,Float64}("unique_measure" => 35.0))
    @test length(data.vehicle_sets[1].work_periods) == 1
    @test length(data.vehicle_sets[1].cost_periods) == 1
    @test data.vehicle_sets[1].cost_periods[1].travel_distance_unit_cost == 1.0
    @test data.vehicle_sets[1].id == "unique_vehicle_set"
    @test data.vehicle_sets[1].vehicle_category_id == "unique_vehicle_category"
    @test(data.vehicle_sets[1].departure_location_group_id ==
            data.vehicle_sets[1].arrival_location_group_id == "depot_loc_group")
    @test data.vehicle_sets[1].nb_of_vehicles_range == RVRP.FlexibleRange(
        RVRP.Range(0, 15), RVRP.Range(0, 15), RVRP.Flexibility(0, 0.0, 0.0),
        RVRP.Flexibility(0, 0.0, 0.0)
    )
    for req_idx in 1:length(data.requests)
        req = data.requests[req_idx]
        @test req.product_quantity_range.lb == req.product_quantity_range.ub
        @test req.pickup_location_group_id == data.location_groups[req_idx+1].id
        @test(req.product_sharing_class_id ==
                req.product_compatibility_class_id == "default_id")
        @test req.product_specification_class_id == "unique_p_spec_c"
    end

    # apply checker on cvrp data, computed data
    computed_data = RVRP.build_computed_data(data)
    RVRP.check_instance(data, computed_data)
    RVRP.print_features(computed_data.features)
end


import Base: convert
import Base: ==

export CalcFactor

"""
$TYPEDEF

User factor interface method for computing the residual values of factors.

Notes
- Also see #467 on API consolidation

```juila
function (cf::CalcFactor{<:LinearRelative})(res::AbstractVector{<:Real}, z, xi, xj)
  cf.metadata.variablelist
  cf.metadata.targetvariable
  cf.metadata.usercache
  generic on-manifold residual function

  return distance(z, distance(xj, xi))
end
```

DevNotes
- Follow the Github project in IIF to better consolidate CCW FMD CPT CF CFM

Related

[`CalcFactorMahalanobis`](@ref), [`CommonConvWrapper`](@ref), [`FactorMetadata`](@ref), [`ConvPerThread`](@ref)
"""
struct CalcFactor{T <: AbstractFactor, M, P <: Union{<:Tuple,Nothing,AbstractVector}, X}
  # the interface compliant user object functor containing the data and logic
  factor::T
  # the metadata to be passed to the user residual function
  metadata::M
  # what is the sample (particle) id for which the residual is being calculated
  _sampleIdx::Int
  # legacy support when concerned with how many measurement tuple elements are used by user
  _measCount::Int
  # legacy suport for measurement sample values of old functor residual functions
  _legacyMeas::P
  # legacy support for variable values old functor residual functions
  _legacyParams::X
end

"""
    $TYPEDEF

Internal parametric extension to [`CalcFactor`](@ref) used for buffering measurement and calculating Mahalanobis distance

Related

[`CalcFactor`](@ref)
"""
struct CalcFactorMahalanobis{CF<:CalcFactor, S, N}
  calcfactor!::CF
  varOrder::Vector{Symbol}
  meas::NTuple{N, <:AbstractVector{Float64}}
  iΣ::NTuple{N, Matrix{Float64}}
  specialAlg::S
end

abstract type _AbstractThreadModel end

"""
$(TYPEDEF)
"""
struct SingleThreaded <: _AbstractThreadModel
end
"""
$(TYPEDEF)
"""
struct MultiThreaded <: _AbstractThreadModel
end

"""
$(TYPEDEF)

Notes
- type-stable `usercache`, see https://github.com/JuliaRobotics/IncrementalInference.jl/issues/783#issuecomment-665080114

DevNotes
- TODO why not just a NamedTuple? Perhaps part of #467
- Follow the Github project in IIF to better consolidate CCW FMD CPT CF CFM
- TODO standardize -- #927, #1025, #784, #692, #640
- TODO make immutable #825
"""
mutable struct FactorMetadata{FV<:AbstractVector{<:DFGVariable},
                              VL<:AbstractVector{Symbol},
                              AR, # <:AbstractVector{<:AbstractArray}
                              CD}
  # full list of Vector{DFGVariable} connected to the factor
  fullvariables::FV # Vector{<:DFGVariable}
  #TODO full variable can perhaps replace this
  variablelist::VL # Vector{Symbol}
  # TODO consolidate, same as ARR used in CCW,
  arrRef::AR # Vector{Matrix{Float64}}
  # label of which variable is being solved for
  solvefor::Symbol
  # for type specific user data, see (? #784)
  cachedata::CD
end

"""
$(TYPEDEF)

DevNotes
- FIXME consolidate with CalcFactor
- TODO consolidate with CCW, FMd, CalcFactor
- TODO consider renaming `.p` to `.decisionDims`
  - TODO `.decisionDims::DD where DD <: Union{<:AbstractVector{Int},Colon}` -- ensure type stability
- TODO figure out if we want static parameter THRID
- TODO make static params {XDIM, ZDIM, P}
- TODO make immutable
"""
mutable struct ConvPerThread{R,F<:FactorMetadata,P}
  thrid_::Int
  # the actual particle being solved at this moment
  particleidx::Int
  # additional/optional data passed to user function
  factormetadata::F
  # subsection indices to select which params should be used for this hypothesis evaluation
  activehypo::Vector{Int}
  # slight numerical perturbation for degenerate solver cases such as division by zero
  perturb::Vector{Float64}
  # working memory location for optimization routines on target decision variables
  X::Vector{P}
  # working memory to store residual for optimization routines
  res::R # was Vector{Float64}
end

"""
$(TYPEDEF)

Main factor memory container used during inference operations -- i.e. values specific to one complete convolution operation

Notes
- CCW does not get serialized / persisted
- At writing, the assumption is there is just one CCW per factor
- Any multithreaded design needs to happens as sub-constainers inside CCW or otherwise, to carry separate memory.
- Since #467, `CalcFactor` is the only type 'seen by user' during `getSample` or function residual calculations `(cf::CalcFactor{<:MyFactor})`, s.t. `MyFactor <: AbstractRelative___`
- There also exists a `CalcFactorMahalanobis` for parameteric computations using as much of the same mechanics as possible.

DevNotes
- Follow the Github project in IIF to better consolidate CCW FMD CPT CF CFM

Related

[`CalcFactor`](@ref), [`CalcFactorMahalanobis`](@ref), [`FactorMetadata`](@ref), [`ConvPerThread`](@ref)
"""
mutable struct CommonConvWrapper{ T<:AbstractFactor,
                                  H<:Union{Nothing, Distributions.Categorical},
                                  C<:Union{Nothing, Vector{Int}},
                                  NTP <: NamedTuple,
                                  G,
                                  MT} <: FactorOperationalMemory
  #
  ### Values consistent across all threads during approx convolution
  usrfnc!::T # user factor / function
  # general setup
  xDim::Int
  zDim::Int
  # is this a partial constraint as defined by the existance of factor field `.partial::Tuple`
  partial::Bool
  # multi hypothesis settings
  hypotheses::H
  # categorical to select which hypothesis is being considered during convolution operation
  certainhypo::C
  nullhypo::Float64
  # parameters passed to each hypothesis evaluation event on user function, #1321
  params::NTP
  # which index is being solved for in params?
  varidx::Int
  # FIXME make type stable, JT should now be type stable if rest works
  #   user defined measurement values for each approxConv operation
  measurement::Vector{MT}
  # TODO refactor and deprecate this old approach, Union{Type{SingleThreaded}, Type{MultiThreaded}}
  threadmodel::Type{<:_AbstractThreadModel}
  # FIXME, deprecate for only `readonly` and build CalcFactor objects on stack instead
    ## will be obsolete: particular convolution computation values per particle idx (varies by thread)
  cpt::Vector{<:ConvPerThread}
  # inflationSpread
  inflation::Float64
  # Which dimensions does this factor influence.  Sensitive (mutable) to both which 'solvefor index' variable and whether the factor is partial dimension
  partialDims::Vector{<:Integer}
  # variable types for points in params
  vartypes::Vector{DataType}
  # experimental feature to embed gradient calcs with ccw
  _gradients::G
end

#


# This file is part of Kpax3. License is MIT.

module Kpax3
  #################
  # Load packages #
  #################
  using Clustering
  using DelimitedFiles
  using Distances
  using FileIO
  using KernelDensity
  using Printf
  using RecipesBase
  using SpecialFunctions
  using StatsBase

  ####################
  # Export functions #
  ####################
  export

  # types
  KSettings,
  AminoAcidData,
  CategoricalData,
  AminoAcidState,

  # distances
  distaamtn84,

  # algorithms
  kpax3mcmc,
  kpax3ga,

  # estimation
  kpax3estimate,
  optimumstate,

  # i/o
  initializepartition,
  normalizepartition,
  readfasta,
  categorical2binary,
  save,
  loadaa,
  writeresults,
  readposteriork,
  readposteriorP,
  readposteriorC,

  # diagnostics
  traceR,
  traceC,

  # plot
  plotD,
  plotk,
  plotP,
  plotC,
  plottrace,
  plotdensity,
  plotjump,
  plotdgn

  ##############
  # Load Types #
  ##############
  include("types/types.jl")

  ########################
  # Load basic functions #
  ########################
  include("misc/misc.jl")

  #################
  # Load the rest #
  #################
  include("data_processing/data_processing.jl")
  include("distances/distances.jl")
  include("model/model.jl")
  include("optimizer/optimizer.jl")
  include("mcmc/mcmc.jl")
  include("estimate/estimate.jl")
  include("plots/plots.jl")
end


using Compat.SparseArrays

const ZeroCones = Union{MOI.EqualTo, MOI.Zeros}
const LPCones = Union{MOI.GreaterThan, MOI.LessThan, MOI.Nonnegatives, MOI.Nonpositives}

_dim(s::MOI.AbstractScalarSet) = 1
_dim(s::MOI.AbstractVectorSet) = MOI.dimension(s)

# Computes cone dimensions
constroffset(cone::Cone, ci::CI{<:MOI.AbstractFunction, <:ZeroCones}) = ci.value
function _allocateconstraint!(cone::Cone, f, s::ZeroCones)
    ci = cone.f
    cone.f += _dim(s)
    ci
end
constroffset(cone::Cone, ci::CI{<:MOI.AbstractFunction, <:LPCones}) = ci.value
function _allocateconstraint!(cone::Cone, f, s::LPCones)
    ci = cone.l
    cone.l += _dim(s)
    ci
end
constroffset(cone::Cone, ci::CI{<:MOI.AbstractFunction, <:MOI.SecondOrderCone}) = cone.l + ci.value
function _allocateconstraint!(cone::Cone, f, s::MOI.SecondOrderCone)
    push!(cone.qa, s.dimension)
    ci = cone.q
    cone.q += _dim(s)
    ci
end
constroffset(cone::Cone, ci::CI{<:MOI.AbstractFunction, <:MOI.ExponentialCone}) = cone.l + cone.q + ci.value
function _allocateconstraint!(cone::Cone, f, s::MOI.ExponentialCone)
    ci = 3cone.ep
    cone.ep += 1
    ci
end
constroffset(instance::ECOSOptimizer, ci::CI) = constroffset(instance.cone, ci::CI)
MOIU.canallocateconstraint(::ECOSOptimizer, ::Type{<:SF}, ::Type{<:SS}) = true
function MOIU.allocateconstraint!(instance::ECOSOptimizer, f::F, s::S) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    CI{F, S}(_allocateconstraint!(instance.cone, f, s))
end

# Build constraint matrix
scalecoef(rows, coef, minus, s) = minus ? -coef : coef
scalecoef(rows, coef, minus, s::Union{MOI.LessThan, Type{<:MOI.LessThan}, MOI.Nonpositives, Type{MOI.Nonpositives}}) = minus ? coef : -coef
_varmap(f) = map(vi -> vi.value, f.variables)
_constant(s::MOI.EqualTo) = s.value
_constant(s::MOI.GreaterThan) = s.lower
_constant(s::MOI.LessThan) = s.upper
constrrows(::MOI.AbstractScalarSet) = 1
constrrows(s::MOI.AbstractVectorSet) = 1:MOI.dimension(s)
constrrows(instance::ECOSOptimizer, ci::CI{<:MOI.AbstractScalarFunction, <:MOI.AbstractScalarSet}) = 1
constrrows(instance::ECOSOptimizer, ci::CI{<:MOI.AbstractVectorFunction, MOI.Zeros}) = 1:instance.cone.eqnrows[constroffset(instance, ci)]
constrrows(instance::ECOSOptimizer, ci::CI{<:MOI.AbstractVectorFunction, <:MOI.AbstractVectorSet}) = 1:instance.cone.ineqnrows[constroffset(instance, ci)]
matrix(data::Data, s::ZeroCones) = data.b, data.IA, data.JA, data.VA
matrix(data::Data, s::Union{LPCones, MOI.SecondOrderCone, MOI.ExponentialCone}) = data.h, data.IG, data.JG, data.VG
matrix(instance::ECOSOptimizer, s) = matrix(instance.data, s)
MOIU.canloadconstraint(::ECOSOptimizer, ::Type{<:SF}, ::Type{<:SS}) = true
MOIU.loadconstraint!(instance::ECOSOptimizer, ci, f::MOI.SingleVariable, s) = MOIU.loadconstraint!(instance, ci, MOI.ScalarAffineFunction{Float64}(f), s)
function MOIU.loadconstraint!(instance::ECOSOptimizer, ci, f::MOI.ScalarAffineFunction, s::MOI.AbstractScalarSet)
    a = sparsevec(_varmap(f), f.coefficients)
    # sparsevec combines duplicates with + but does not remove zeros created so we call dropzeros!
    dropzeros!(a)
    offset = constroffset(instance, ci)
    row = constrrows(s)
    i = offset + row
    # The ECOS format is b - Ax ∈ cone
    # so minus=false for b and minus=true for A
    setconstant = _constant(s)
    if s isa MOI.EqualTo
        instance.cone.eqsetconstant[offset] = setconstant
    else
        instance.cone.ineqsetconstant[offset] = setconstant
    end
    constant = f.constant - setconstant
    b, I, J, V = matrix(instance, s)
    b[i] = scalecoef(row, constant, false, s)
    append!(I, fill(i, length(a.nzind)))
    append!(J, a.nzind)
    append!(V, scalecoef(row, a.nzval, true, s))
end
MOIU.loadconstraint!(instance::ECOSOptimizer, ci, f::MOI.VectorOfVariables, s) = MOIU.loadconstraint!(instance, ci, MOI.VectorAffineFunction{Float64}(f), s)
# SCS orders differently than MOI the second and third dimension of the exponential cone
orderval(val, s) = val
function orderval(val, s::Union{MOI.ExponentialCone, Type{MOI.ExponentialCone}})
    val[[1, 3, 2]]
end
orderidx(idx, s) = idx
expmap(i) = (1, 3, 2)[i]
function orderidx(idx, s::MOI.ExponentialCone)
    expmap.(idx)
end
function MOIU.loadconstraint!(instance::ECOSOptimizer, ci, f::MOI.VectorAffineFunction, s::MOI.AbstractVectorSet)
    A = sparse(f.outputindex, _varmap(f), f.coefficients)
    # sparse combines duplicates with + but does not remove zeros created so we call dropzeros!
    dropzeros!(A)
    colval = zeros(Int, length(A.rowval))
    for col in 1:A.n
        colval[A.colptr[col]:(A.colptr[col+1]-1)] = col
    end
    @assert !any(iszero.(colval))
    offset = constroffset(instance, ci)
    rows = constrrows(s)
    if s isa MOI.Zeros
        instance.cone.eqnrows[offset] = length(rows)
    else
        instance.cone.ineqnrows[offset] = length(rows)
    end
    i = offset + rows
    # The ECOS format is b - Ax ∈ cone
    # so minus=false for b and minus=true for A
    b, I, J, V = matrix(instance, s)
    b[i] = scalecoef(rows, orderval(f.constant, s), false, s)
    append!(I, offset + orderidx(A.rowval, s))
    append!(J, colval)
    append!(V, scalecoef(A.rowval, A.nzval, true, s))
end

function MOIU.allocatevariables!(instance::ECOSOptimizer, nvars::Integer)
    instance.cone = Cone()
    VI.(1:nvars)
end

function MOIU.loadvariables!(instance::ECOSOptimizer, nvars::Integer)
    cone = instance.cone
    m = cone.l + cone.q + 3cone.ep
    IA = Int[]
    JA = Int[]
    VA = Float64[]
    b = zeros(cone.f)
    IG = Int[]
    JG = Int[]
    VG = Float64[]
    h = zeros(m)
    c = zeros(nvars)
    instance.data = Data(m, nvars, IA, JA, VA, b, IG, JG, VG, h, 0., c)
end

MOIU.canallocate(::ECOSOptimizer, ::MOI.ObjectiveSense) = true
function MOIU.allocate!(instance::ECOSOptimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    instance.maxsense = sense == MOI.MaxSense
end
MOIU.canallocate(::ECOSOptimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
function MOIU.allocate!(::ECOSOptimizer, ::MOI.ObjectiveFunction, ::MOI.ScalarAffineFunction) end

MOIU.canload(::ECOSOptimizer, ::MOI.ObjectiveSense) = true
function MOIU.load!(::ECOSOptimizer, ::MOI.ObjectiveSense, ::MOI.OptimizationSense) end
MOIU.canload(::ECOSOptimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
function MOIU.load!(instance::ECOSOptimizer, ::MOI.ObjectiveFunction, f::MOI.ScalarAffineFunction)
    c0 = full(sparsevec(_varmap(f), f.coefficients, instance.data.n))
    instance.data.objconstant = f.constant
    instance.data.c = instance.maxsense ? -c0 : c0
end

function MOI.optimize!(instance::ECOSOptimizer)
    cone = instance.cone
    m = instance.data.m
    n = instance.data.n
    A = ECOS.ECOSMatrix(sparse(instance.data.IA, instance.data.JA, instance.data.VA, cone.f, n))
    b = instance.data.b
    G = ECOS.ECOSMatrix(sparse(instance.data.IG, instance.data.JG, instance.data.VG, m, n))
    h = instance.data.h
    objconstant = instance.data.objconstant
    c = instance.data.c
    instance.data = nothing # Allows GC to free instance.data before A is loaded to ECOS
    ecos_prob_ptr = ECOS.setup(n, m, cone.f, cone.l, length(cone.qa), cone.qa, cone.ep, G, A, c, h, b)
    ret_val = ECOS.solve(ecos_prob_ptr)
    ecos_prob = unsafe_wrap(Array, ecos_prob_ptr, 1)[1]
    primal    = unsafe_wrap(Array, ecos_prob.x, n)[:]
    dual_eq   = unsafe_wrap(Array, ecos_prob.y, cone.f)[:]
    dual_ineq = unsafe_wrap(Array, ecos_prob.z, m)[:]
    slack     = unsafe_wrap(Array, ecos_prob.s, m)[:]
    ECOS.cleanup(ecos_prob_ptr, 0)
    objval = (instance.maxsense ? -1 : 1) * dot(c, primal) + objconstant
    instance.sol = Solution(ret_val, primal, dual_eq, dual_ineq, slack, objval)
end


using ArgParse

function parse_commandline()
	s = ArgParseSettings()
	@add_arg_table s begin
		"--dest_range"
			help = "Distance between the initial position and the destination of the ego agent"
			arg_type = Float64
			default = 10000.0 # [m]
		"--sensing_range"
			help = "Radius of the alerting perimeter"
			arg_type = Float64
			default = 1000.0 # [m]

		"--penalty_action" # max =  - 100 * pen
			help = "Reward function parameter: penalty on action"
			arg_type = Float64
			default = 2.0
		"--penalty_conflict" # max =  - pen
			help = "Reward function parameter: penalty on conflict"
			arg_type = Float64
			default = 50.0
		"--penalty_closeness" # max = - pen (on NMAC) or - pen * 2.71 (on collision)
			help = "Reward function parameter: penalty on closeness between agents"
			arg_type = Float64
			default = 500.0
		"--penalty_nmac" # max = - pen
			help = "Reward function parameter: penalty on near mid-air collisions"
			arg_type = Float64
			default = 7000.0

		"--penalty_deviation" # max = - 180 * pen
			help = "Reward function parameter: penalty on direction deviation from destination"
			arg_type = Float64
			default = 0.8
		"--penalty_digression" # max =
			help = "Reward function parameter: penalty on digression from destination"
			arg_type = Float64
			default = 50.0
		"--reward_destination" # max = rew
			help = "Reward function parameter: reward on reaching destination"
			arg_type = Float64
			default = 0.0
	end
	return parse_args(s)
end

parsed_args = parse_commandline()

const SPAWN_RATE_SET = [25, 35, 45]
const CONST_NUM_AC_SET = [1]
# SPAWN_RATE = rand(SPAWN_RATE_SET)
const BEARING_DIST =   [[0.1047 0.0489];
   						[0.3142 0.0474];
   						[0.5236 0.0474];
   						[0.7330 0.0459];
   						[0.9425 0.0425];
   						[1.1519 0.0403];
   						[1.3614 0.0365];
   						[1.5708 0.0333];
   						[1.7802 0.0304];
   						[1.9897 0.0269];
   						[2.1991 0.0237];
   						[2.4086 0.0210];
   						[2.6180 0.0194];
   						[2.8274 0.0183];
   						[3.0369 0.0181];
   						[3.2463 0.0181];
   						[3.4558 0.0183];
   						[3.6652 0.0194];
   						[3.8746 0.0210];
   						[4.0841 0.0237];
   						[4.2935 0.0269];
   						[4.5029 0.0304];
   						[4.7124 0.0333];
   						[4.9218 0.0365];
   						[5.1313 0.0403];
   						[5.3407 0.0425];
   						[5.5501 0.0459];
   						[5.7596 0.0474];
   						[5.9690 0.0474];
   						[6.1785 0.0489];]

const DT = 1.0 # [s]
const V_MIN = 50.0 # [m/s]
const V_MAX = 50.0 # [m/s], currently using single value for speed
const AIRSPACE_DIM = 1500  # [m]
const STATE_DIM = 5
const AUG_STATE_DIM = 3

const DEST_CRITERION = 120.0 # [m]
const NMAC_RANGE = 150.0 # [m]
const NUM_SECTIONS = 4
const RELAX_FACTOR = 1.2 # for train
# const RELAX_FACTOR = 1.0 # for resTime evaluation

const TERM_VAR = terminal_state_var # imported from CAS_QMDP

# state ind
# s = [ρ, θ, ϕ, v_own, v_int]
const IND_rho   = 1
const IND_theta = 2
const IND_phi   = 3
const IND_v_own = 4
const IND_v_int = 5
const IND_deviation = STATE_DIM * NUM_SECTIONS + 1
const IND_dist_to_dest = STATE_DIM * NUM_SECTIONS + 2
const IND_prev_dist_to_dest = STATE_DIM * NUM_SECTIONS + 3

const SENSING_RANGE = parsed_args["sensing_range"]
const DEST_RANGE = parsed_args["dest_range"]

# Reward weights for the correction network
const PEN_ACTION = parsed_args["penalty_action"]
const PEN_CLOSENESS = parsed_args["penalty_closeness"]
const PEN_NMAC = parsed_args["penalty_nmac"]
const PEN_CONFLICT = parsed_args["penalty_conflict"]
const PEN_DEVIATION = parsed_args["penalty_deviation"]
const PEN_DIGRESSION = parsed_args["penalty_digression"]
const REW_DESTINATION = parsed_args["reward_destination"]

# Leaving out the huge positive reward from reaching destination:
const POSSIBLE_MAX_REW = V_MAX * DT * PEN_DIGRESSION
# Keeping the huge positive reward from reaching destination:
# const POSSIBLE_MAX_REW = REW_DESTINATION + V_MAX * DT * PEN_DIGRESSION

# Leaving out the huge negative reward from NMAC:
# const POSSIBLE_MIN_REW = - PEN_CONFLICT - PEN_ACTION * 10^2 - PEN_CLOSENESS * e - 180 * PEN_DEVIATION - V_MAX * DT * PEN_DIGRESSION
# Keeping the huge negative reward from NMAC:
const POSSIBLE_MIN_REW = - PEN_CONFLICT - PEN_ACTION * 10^2 - PEN_NMAC - PEN_CLOSENESS * e - 180 * PEN_DEVIATION - V_MAX * DT * PEN_DIGRESSION

const LOW_FI_POLICY = "mid3_maxdist_2km_nmac_150m_Discount_0.95_ActWeight_0.5_IfSample_false_PenAction_3.0_PenConflict_50.0_PenCloseness_500.0_PenNmac_7000.0_Sigma_v_5.0_alert_5.0_theta_3.0_phi_7.0_coc_3.0"
const LOW_FI_POLICY_FILENAME =
	"../../VICAS/policies/" * LOW_FI_POLICY * ".jld"

# Step reward bounds from VICAS
# These values need to match those used to train VICAS
# Read them from LOW_FI_POLICY
const penalty_action_VICAS = 3.0
const penalty_conflict_VICAS = 50.0
const penalty_closeness_VICAS = 500.0
const penalty_nmac_VICAS = 7000.0

const POSSIBLE_MAX_REW_VICAS = 0.0
const POSSIBLE_MIN_REW_VICAS = - penalty_conflict_VICAS - penalty_action_VICAS * 10.^2 - penalty_nmac_VICAS - penalty_closeness_VICAS * e


"""
    tour = sortperm_msa(msa::AbstractMultipleSequenceAlignment)

Order the sequences in `msa` to minimize the "tour length" visiting each sequence once.
The length between sequences is defined at `100 - percentsimilarity(seq1, seq2)`.

This can be useful for graphical or alignment display by grouping obviously-similar
sequences near one another.
"""
function sortperm_msa(msa::AbstractMultipleSequenceAlignment)
    sim = percentsimilarity(msa)
    D = 100 .- Matrix(sim)
    tour, _ = solve_tsp(D; quality_factor=80)
    pop!(tour)  # returns a closed path, but that duplicates the first/last entry
    return tour
end

"""
    filter_species!(msa::AbstractMultipleSequenceAlignment, speciesname::AbstractString)

Remove all sequences from `msa` except those with [`species(sequencename)`](@ref) equal to `speciesname`.
"""
function filter_species!(msa::AbstractMultipleSequenceAlignment, speciesname::AbstractString)
    mask = map(x -> species(x) == speciesname, sequencenames(msa))
    filtersequences!(msa, mask)
end

"""
    filter_long!(msa::AbstractMultipleSequenceAlignment, minres::Real)

Remove all sequences from `msa` with fewer than `minres` matching residues.
"""
function filter_long!(msa::AbstractMultipleSequenceAlignment, minres::Real)
    # Get rid of short sequences
    nresidues = map(eachrow(msa)) do v
        sum(!=(Residue('-')), v)
    end
    mask = nresidues .> minres
    filtersequences!(msa, mask)
end

# Move this to MIToS?
if !hasmethod(getsequencemapping, Tuple{AnnotatedAlignedSequence})
    function MIToS.MSA.getsequencemapping(seq::AnnotatedAlignedSequence)
        getsequencemapping(seq, sequencenames(seq)[1])
    end
    function MIToS.MSA.getsequencemapping(msa::Union{AnnotatedAlignedSequence,AnnotatedMultipleSequenceAlignment}, seq_id::String)
        MIToS.MSA._str2int_mapping(getannotsequence(msa, seq_id, "SeqMap"))
    end
    function MIToS.MSA.getsequencemapping(msa::AnnotatedMultipleSequenceAlignment, seqid::Regex)
        id = findfirst(str -> occursin(seqid, str), sequencenames(msa))
        getsequencemapping(msa, id)
    end
end

struct SequenceMapping <: AbstractVector{Int}
    seqmap::Vector{Int}
end
SequenceMapping(seq::AnnotatedAlignedSequence) = SequenceMapping(getsequencemapping(seq))

Base.size(sm::SequenceMapping) = size(sm.seqmap)
Base.getindex(sm::SequenceMapping, i::Int) = sm.seqmap[i]
Base.setindex!(sm::SequenceMapping, v, i::Int) = sm.seqmap[i] = v
Base.similar(::SequenceMapping, ::Type{Int}, dims::Dims{1}) = SequenceMapping(Vector{Int}(undef, dims))

# We restrict the following to `Vector` to prevent ambiguities
Base.getindex(fullseqvec::Vector{T}, sm::SequenceMapping) where T = Union{T,Nothing}[i == 0 ? nothing : fullseqvec[i] for i in sm.seqmap]

function Base.setindex!(fullseqvec::Vector, colvalues::AbstractVector, sm::SequenceMapping)
    idxnz = sm.seqmap .> 0
    fullseqvec[sm.seqmap[idxnz]] = colvalues[idxnz]
end

"""
    chimerax_script(filename, uprot_list, msa::AnnotatedMultipleSequenceAlignment, colidxs;
                    dir=pwd(), chain_transparency=80, styles=Dict{Int,String}(), extras=String[])

Create a [chimerax]() visualization script with name `filename`. `uprot_list` is a list of UniProtX names that you
want to visualize.  The AlphaFold PDB files for these proteins should be `dir`. `msa` is a Multiple Sequence alignment
and `colidxs` specifies the column indices in `msa` that you'd like to visualize.

`chain_transparency` sets the transparency on the ribbon diagrams (0 = not transparent).
`styles` can be used to affect the display, e.g., `Dict(k => "@SD sphere")` would cause methionines at column index `k`
to be displayed with the sulfur in sphere mode. `extras` can be used to hand-specify a number of additional commands;
this can be useful if, for example, the `msa` has occasional misalignments.

# Examples

Suppose you have the `msa` for rhodopsin (mouse: P15409), then:

```julia
chimerax_script("myscript.cxc", ["P15409"], msa, [i1, i2, i3])
```

where `i1` through `i3` are column-indices in the `msa` that you'd like to view.
"""
function chimerax_script(filename, uprot_list, msa::AnnotatedMultipleSequenceAlignment, colidxs;
                         dir=pwd(), chain_transparency=80, styles=Dict{Int,String}(), extras=String[])
    open(filename, "w") do io
        for (i, p) in enumerate(uprot_list)
            afname = alphafoldfile(p)
            println(io, "open ", joinpath(dir, afname))
            seqidx = findfirst(str -> startswith(str, p), sequencenames(msa))
            sm = getsequencemapping(msa, seqidx)
            for c in colidxs
                ridx = sm[c]
                if iszero(ridx)
                    @warn "column $c not set in $p"
                    continue
                end
                println(io, "show #", i, " :", ridx)
                style = get(styles, c, nothing)
                if style !== nothing
                    println(io, "style #", i, " :", ridx, " ", style)
                end
            end
        end
        for ex in extras
            println(io, ex)
        end
        println(io, "transparency #1-", length(uprot_list), ' ', chain_transparency, " target c")
        if !any(str -> startswith(str, "align"), extras)
            println(io, "matchmaker #2-", length(uprot_list), " to #1")
        end
    end
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    MetropolisHastings

Metropolis-Hastings MCMC sampling algorithm.
!!!!!
"""
struct MetropolisHastings{
    Q<:ProposalDistSpec,
    W<:Real,
    WS<:AbstractWeightingScheme{W}
} <: MCMCAlgorithm
    proposalspec::Q
    weighting::WS

    MetropolisHastings(proposalspec::Q,weighting::WS) where {
        Q<:ProposalDistSpec, W<:Real, WS<:AbstractWeightingScheme{W}} =
        new{Q,W,WS}(proposalspec, weighting)
end

export MetropolisHastings

MetropolisHastings(proposalspec::ProposalDistSpec = MvTDistProposal()) =
    MetropolisHastings(proposalspec, RepetitionWeighting())

MetropolisHastings(weighting::AbstractWeightingScheme) =
    MetropolisHastings(MvTDistProposal(), weighting)

mcmc_compatible(::MetropolisHastings, ::AbstractProposalDist, ::NoParamBounds) = true

mcmc_compatible(::MetropolisHastings, proposaldist::AbstractProposalDist, bounds::HyperRectBounds) =
    issymmetric(proposaldist) || all(x -> x == hard_bounds, bounds.bt)

_sample_weight_type(::Type{MetropolisHastings{Q,W,WS}}) where {Q,W,WS} = W

mutable struct MHIterator{
    SP<:MCMCSpec,
    R<:AbstractRNG,
    Q<:AbstractProposalDist,
    SV<:DensitySampleVector
} <: MCMCIterator
    spec::SP
    rng::R
    info::MCMCIteratorInfo
    proposaldist::Q
    samples::SV
    nsamples::Int64
    stepno::Int64
end

function MHIterator(
    spec::MCMCSpec,
    info::MCMCIteratorInfo,
    initial_params::AbstractVector{P},
) where {P<:Real}
    stepno::Int64 = 0
    rng = spec.rngseed()
    reset_rng_counters!(rng, info.id, info.cycle, stepno)

    postr = spec.posterior
    npar = nparams(postr)
    alg = spec.algorithm

    params_vec = Vector{P}(undef, npar)
    if isempty(initial_params)
        initial_params!(params_vec, rng, postr, alg)
    else
        params_vec .= initial_params
    end
    !(params_vec in param_bounds(postr)) && throw(ArgumentError("Parameter(s) out of bounds"))

    proposaldist = alg.proposalspec(P, npar)

    # ToDo: Make numeric type configurable:

    log_posterior_value = apply_bounds_and_eval_posterior_logval_strict!(postr, params_vec)

    T = typeof(log_posterior_value)
    W = _sample_weight_type(typeof(alg))

    sample_info = MCMCSampleID(info.id, info.cycle, 1, CURRENT_SAMPLE)
    current_sample = DensitySample(params_vec, log_posterior_value, one(W), sample_info, nothing)

    samples = DensitySampleVector{P,T,W,MCMCSampleID,Nothing}(undef, 0, npar)
    push!(samples, current_sample)

    nsamples::Int64 = 0

    chain = MHIterator(
        spec,
        rng,
        info,
        proposaldist,
        samples,
        nsamples,
        stepno
    )

    chain
end

function reset_rng_counters!(chain::MHIterator)
    reset_rng_counters!(chain.rng, chain.info.id, chain.info.cycle, chain.stepno)
end

function (spec::MCMCSpec{<:MetropolisHastings})(
    chainid::Integer,
)
    P = float(eltype(param_bounds(spec.posterior)))

    cycle = 0
    tuned = false
    converged = false
    info = MCMCIteratorInfo(chainid, cycle, tuned, converged)

    MHIterator(spec, info, Vector{P}())
end

@inline _current_sample_idx(chain::MHIterator) = firstindex(chain.samples)
@inline _proposed_sample_idx(chain::MHIterator) = lastindex(chain.samples)

function _available_samples_idxs(chain::MHIterator)
    sampletype = chain.samples.info.sampletype
    @uviews sampletype begin
        from = firstindex(chain.samples)

        to = if samples_available(chain)
            lastidx = lastindex(chain.samples)
            @assert sampletype[from] == ACCEPTED_SAMPLE
            @assert sampletype[lastidx] == CURRENT_SAMPLE
            lastidx - 1
        else
            from - 1
        end

        r = from:to
        @assert all(x -> x > INVALID_SAMPLE, view(sampletype, r))
        r
    end
end

mcmc_spec(chain::MHIterator) = chain.spec

getrng(chain::MHIterator) = chain.rng

mcmc_info(chain::MHIterator) = chain.info

nsteps(chain::MHIterator) = chain.stepno

nsamples(chain::MHIterator) = chain.nsamples

current_sample(chain::MHIterator) = chain.samples[_current_sample_idx(chain)]

sample_type(chain::MHIterator) = eltype(chain.samples)

function samples_available(chain::MHIterator)
    i = _current_sample_idx(chain::MHIterator)
    chain.samples.info.sampletype[i] == ACCEPTED_SAMPLE
end

function get_samples!(appendable, chain::MHIterator, nonzero_weights::Bool)::typeof(appendable)
    if samples_available(chain)
        idxs = _available_samples_idxs(chain)
        samples = chain.samples

        @uviews samples begin
            # if nonzero_weights
                for i in idxs
                    if !nonzero_weights || samples.weight[i] > 0
                        push!(appendable, samples[i])
                    end
                end
            # else
            #     append!(appendable, view(samples, idxs))
            # end
        end
    end
    appendable
end

function next_cycle!(chain::MHIterator)
    chain.info = MCMCIteratorInfo(chain.info, cycle = chain.info.cycle + 1)
    chain.nsamples = 0
    chain.stepno = 0

    reset_rng_counters!(chain)

    resize!(chain.samples, 1)

    i = _current_sample_idx(chain)
    @assert chain.samples.info[i].sampletype == CURRENT_SAMPLE

    chain.samples.weight[i] = 1
    chain.samples.info[i] = MCMCSampleID(chain.info.id, chain.info.cycle, chain.stepno, CURRENT_SAMPLE)

    chain
end

function mcmc_step!(
    callback::AbstractMCMCCallback,
    chain::MHIterator
)
    alg = algorithm(chain)

    if !mcmc_compatible(alg, chain.proposaldist, param_bounds(getposterior(chain)))
        error("Implementation of algorithm $alg does not support current parameter bounds with current proposal distribution")
    end

    chain.stepno += 1
    reset_rng_counters!(chain)

    rng = getrng(chain)
    pstr = getposterior(chain)

    proposaldist = chain.proposaldist
    samples = chain.samples

    # Grow samples vector by one:
    resize!(samples, size(samples, 1) + 1)
    samples.info[lastindex(samples)] = MCMCSampleID(chain.info.id, chain.info.cycle, chain.stepno, PROPOSED_SAMPLE)

    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    @assert current != proposed

    accepted = @uviews samples begin
        current_params = samples.params[current]
        proposed_params = samples.params[proposed]

        # Propose new parameters:
        samples.weight[proposed] = 0
        proposal_rand!(rng, proposaldist, proposed_params, current_params)

        current_log_posterior = samples.logdensity[current]
        T = typeof(current_log_posterior)

        # Evaluate prior and likelihood with proposed parameters:
        proposed_log_posterior = apply_bounds_and_eval_posterior_logval!(T, pstr, proposed_params)

        samples.logdensity[proposed] = proposed_log_posterior

        p_accept = if proposed_log_posterior > -Inf
            # log of ratio of forward/reverse transition probability
            log_tpr = if issymmetric(proposaldist)
                T(0)
            else
                log_tp_fwd = distribution_logpdf(proposaldist, proposed_params, current_params)
                log_tp_rev = distribution_logpdf(proposaldist, current_params, proposed_params)
                T(log_tp_fwd - log_tp_rev)
            end

            p_accept_unclamped = exp(proposed_log_posterior - current_log_posterior - log_tpr)
            T(clamp(p_accept_unclamped, 0, 1))
        else
            zero(T)
        end

        @assert p_accept >= 0
        accepted = rand(chain.rng, float(typeof(p_accept))) < p_accept

        if accepted
            samples.info.sampletype[current] = ACCEPTED_SAMPLE
            samples.info.sampletype[proposed] = CURRENT_SAMPLE
            chain.nsamples += 1
        else
            samples.info.sampletype[proposed] = REJECTED_SAMPLE
        end

        delta_w_current, w_proposed = _mh_weights(alg, p_accept, accepted)
        samples.weight[current] += delta_w_current
        samples.weight[proposed] = w_proposed

        callback(1, chain)

        if accepted
            current_params .= proposed_params
            samples.logdensity[current] = samples.logdensity[proposed]
            samples.weight[current] = samples.weight[proposed]
            samples.info[current] = samples.info[proposed]
        end

        accepted
    end # @uviews

    if accepted
        resize!(samples, 1)
    end

    chain
end

function _mh_weights(
    algorithm::MetropolisHastings{Q,W,<:RepetitionWeighting},
    p_accept::Real,
    accepted::Bool
) where {Q,W}
    if accepted
        (0, 1)
    else
        (1, 0)
    end
end

function _mh_weights(
    algorithm::MetropolisHastings{Q,W,<:ARPWeighting},
    p_accept::Real,
    accepted::Bool
) where {Q,W}
    T = typeof(p_accept)
    if p_accept ≈ 1
        (zero(T), one(T))
    elseif p_accept ≈ 0
        (one(T), zero(T))
    else
        (T(1 - p_accept), p_accept)
    end
end

eff_acceptance_ratio(chain::MHIterator) = nsamples(chain) / nsteps(chain)


#=============================#
#== put needed modules here ==#
#=============================#
import JuliaChem
import Statistics
import HypothesisTests

#================================#
#== JuliaChem execution script ==#
#================================#
Base.@ccallable function julia_main(ARGS::Vector{String})::Cint
  #== initialize JuliaChem runtime ==#
  JuliaChem.initialize()

  #== read in input file ==#
  input_file = ARGS[1]

  input_time1_t1 = time_ns()/1e9
  molecule, driver, model, keywords = JuliaChem.JCInput.run(input_file)
  input_time1_t2 = time_ns()/1e9
  input_time1 = input_time1_t2 - input_time1_t1

  input_time2_t1 = time_ns()/1e9
  molecule, driver, model, keywords = JuliaChem.JCInput.run(input_file)
  input_time2_t2 = time_ns()/1e9
  input_time2 = input_time2_t2 - input_time2_t1

  input_jit = input_time1 - input_time2
  molecule, driver, model, keywords = JuliaChem.JCInput.run(input_file)

  #== generate basis set ==#
  basis_time1_t1 = time_ns()/1e9
  basis = JuliaChem.JCBasis.run(molecule, model)
  basis_time1_t2 = time_ns()/1e9
  basis_time1 = basis_time1_t2 - basis_time1_t1

  basis_time2_t1 = time_ns()/1e9
  basis = JuliaChem.JCBasis.run(molecule, model)
  basis_time2_t2 = time_ns()/1e9
  basis_time2 = basis_time2_t2 - basis_time2_t1

  basis_jit = basis_time1 - basis_time2
  basis = JuliaChem.JCBasis.run(molecule, model)

  #== perform scf benchmark ==#
  timeof = Vector{Float64}(undef,0)
  scf_time1 = 0.0
  if (driver == "energy")
    if (model["method"] == "RHF")
      scf_time1_t1 = time_ns()/1e9
      scf = JuliaChem.JCRHF.run(basis, molecule, keywords)
      scf_time1_t2 = time_ns()/1e9
      scf_time1 = scf_time1_t2 - scf_time1_t1

      JuliaChem.reset()

      for index in 1:10
        molecule, driver, model, keywords = JuliaChem.JCInput.run(input_file)
        basis = JuliaChem.JCBasis.run(molecule, model)

        scf_timeof_t1 = time_ns()/1e9
        scf = JuliaChem.JCRHF.run(basis, molecule, keywords) #initial run
        scf_timeof_t2 = time_ns()/1e9
        push!(timeof, scf_timeof_t2 - scf_timeof_t1)

        JuliaChem.reset()
      end
    end
  end
  scf_jit = scf_time1 - Statistics.mean(timeof)

  #== output relevant information ==#
  println("Input JIT: ", input_jit)
  println("Basis JIT: ", basis_jit)
  println("SCF JIT: ", scf_jit)
  println("Total JIT: ", input_jit + basis_jit + scf_jit)
  println("")

  #== perform t-test to compare to GAMESS ==#
  p = HypothesisTests.OneSampleTTest(timeof,2.94+0.19)
  println(p)

  #== finalize JuliaChem runtime ==#
  JuliaChem.finalize()

  #== we are done ==#
  return 0
end


#!/usr/bin/env julia

using MIToS.Utils.Scripts

Args = parse_commandline(
    # TO DO ----------------------------------------------------------------------
    ["--distance", "-d"],
    Dict(
        :help => "The distance to be calculated, options: All, Heavy, CA, CB",
        :arg_type => ASCIIString,
        :default => "All"
    ),
    ["--format", "-f"],
    Dict(
        :help => "Format of the PDB file: It should be PDBFile or PDBML",
        :arg_type => ASCIIString,
        :default => "PDBFile"
    ),
    ["--model", "-m"],
    Dict(
        :help => "The model to be used, * for all",
        :arg_type => ASCIIString,
        :default => "1"
    ),
    ["--chain", "-c"],
    Dict(
        :help => "The chain to be used, * for all",
        :arg_type => ASCIIString,
        :default => "*"
    ),
    ["--group", "-g"],
    Dict(
        :help => "Group of atoms to be used, should be ATOM, HETATM or * for all",
        :arg_type => ASCIIString,
        :default => "*"
    ),
    ["--inter", "-i"],
    Dict(
        :help => "Calculate inter chain distances",
        :action => :store_true
    ),
    # Keywords...
    description="""
    Calculates residues distance and writes them into a *.distances.csv.gz gzipped file.
    """,
    output=".distances.csv.gz"
    # ----------------------------------------------------------------------------
    )

set_parallel(Args["parallel"])

@everywhere begin

    const args = remotecall_fetch(1,()->Args)

    import MIToS.Utils.Scripts: script

    # TO DO ----------------------------------------------------------------------
    using MIToS.PDB
    # ----------------------------------------------------------------------------

    function script(input::Union{Base.LibuvStream,  AbstractString},
                    args,
                    fh_out::Union{Base.LibuvStream, IO})
        # TO DO ------------------------------------------------------------------
        println(fh_out, "# MIToS ", Pkg.installed("MIToS"), " Distances.jl ", now())
        println(fh_out, "# used arguments:")
        for (key, value) in args
            println(fh_out, "# \t", key, "\t\t", value)
        end
        println(fh_out, "model_i,chain_i,group_i,pdbe_i,number_i,name_i,model_j,chain_j,group_j,pdbe_j,number_j,name_j,distance")
        dtyp = ascii(args["distance"])
        form = ascii(args["format"])
        if form == "PDBFile"
            res = readorparse(input, PDBFile)
        elseif form == "PDBML"
            res = readorparse(input, PDBML)
        else
            throw(ErrorException("--format should be PDBFile or PDBML."))
        end
        res = residues(res, ascii(args["model"]), ascii(args["chain"]), ascii(args["group"]), "*")
        N = length(res)
        inter = !Bool(args["inter"])
        for i in 1:(N-1)
            for j in (i+1):N
                @inbounds res1 = res[i]
                @inbounds res2 = res[j]
                if inter && res1.id.chain != res2.id.chain
                    continue
                end
                dist = distance(res1, res2, criteria=dtyp)
                println(fh_out, res1.id.model, ",", res1.id.chain, ",", res1.id.group, ",", res1.id.PDBe_number, ",", res1.id.number, ",", res1.id.name, ",",
                        res2.id.model, ",", res2.id.chain, ",", res2.id.group, ",", res2.id.PDBe_number, ",", res2.id.number, ",", res2.id.name, ",",
                        dist )
            end
        end
        # ------------------------------------------------------------------------
    end

end

runscript(args)


"""
    struct CachedPool{O,T}

        Use this pool to host the objects that are heavy to evaluate so that one wants to cache their status.
        The user should defines a compare

# Members
- name::Symbol : name of the pool
- object::O    : object
- current::T      : current status
- new::T       : the new status wants to assign later
- version::Int128 : the current version
- excited::Bool   : if set to excited, then the current status needs to be replaced with the new status
"""
mutable struct CachedPool{O,T}
    name::Symbol
    object::Vector{O}
    current::Vector{T}
    new::Vector{T}
    version::Vector{Int128}
    excited::Vector{Bool}

    function CachedPool(name::Symbol, objType::DataType, weightType::DataType)
        object = Vector{objType}(undef, 0)
        current = Vector{weightType}(undef, 0)
        _new = Vector{weightType}(undef, 0)
        version = Vector{Int128}(undef, 0)
        excited = Vector{Bool}(undef, 0)
        return new{objType,weightType}(name, object, current, _new, version, excited)
    end
    # function CachedPool{T}(obj::Vector{O}) where {O,T}
    #     weight = zeros(5, length(obj))
    #     return new{O,T}(obj, weight)
    # end
end

Base.length(pool::CachedPool) = length(pool.object)
Base.size(pool::CachedPool) = size(pool.object)
Base.show(io::IO, pool::CachedPool) = print(io, pool.object)
# Base.view(pool::Pool, inds...) = Base.view(pool.pool, inds...)

#index interface for Pool
Base.getindex(pool::CachedPool, i) = pool.object[i]
Base.setindex!(pool::CachedPool, v, i) = setindex!(pool.object, v, i)
Base.firstindex(pool::CachedPool) = 1
Base.lastindex(pool::CachedPool) = length(pool)

# iterator interface
function Base.iterate(pool::CachedPool)
    if length(pool) == 0
        return nothing
    else
        return (pool.object[1], 1)
    end
end

function Base.iterate(pool::CachedPool, state)
    if state >= length(pool) || length(pool) == 0
        return nothing
    else
        return (pool.object[state+1], state + 1)
    end
end

function append(pool::CachedPool, object)
    # @assert para isa eltype(pool.pool)
    # for (oi, o) in enumerate(pool.object)
    #     if o == object
    #         return oi #existing obj
    #     end
    # end

    id = length(pool.object) + 1
    push!(pool.object, object)

    return id #new momentum
end

function initialize!(pool::CachedPool{O,T}) where {O,T}
    N = length(pool)
    pool.current = zeros(T, N)
    pool.new = zeros(T, N)
    pool.version = ones(T, N)
    pool.excited = Vector{Bool}(zeros(Int, N))
end

function updateAll(pool::CachedPool, ignoreCache::Bool, eval::Function; kwargs...)
    N = length(pool.object)
    if ignoreCache
        T = eltype(pool.current)
        for (idx, o) in enumerate(pool.object)
            pool.current[idx] = T(eval(obj; kwargs...))
        end
    else
        error("not implemented!")
        # pool.version .= 1
        # pool.excited .= false
    end
end

"""
    struct LoopPool{T}

    Pool of loop basis. Each loop basis corresponds to a loop variable.
    A loop variable is a linear combination of N independent loops. The combination coefficients is what we call a loop basis.
    For example, if a loop is a momentum K, then

    varibale_i = K_1*basis[1, i] + K_2*basis[2, i] + K_3*basis[3, i] + ...

# Members
- name::Symbol : name of the pool
- dim::Int     : dimension of a loop variable (for example, the dimension of a momentum-frequency loop variable is (d+1) where d is the spatial dimension)
- N::Int       : number of independent loops (dimension of loop basis)
- basis::Matrix{T}    : Matrix of (N x Nb) that stores the loop basis, where Nb is the number of loop basis (or number of loop variables).
- current::Matrix{T}  : Matrix of (dim x Nb) that stores the loop variables, where Nb is the number of loop basis (or number of loop variables).
"""
mutable struct LoopPool{T}
    name::Symbol
    dim::Int #dimension
    N::Int #number of basis
    basis::Matrix{T}
    current::Matrix{T}

    function LoopPool(name::Symbol, dim::Int, N::Int, type::DataType = Float64)
        basis = Matrix{type}(undef, N, 0) # Nx0 matrix
        current = Matrix{type}(undef, dim, 0) # dimx0 matrix
        return new{type}(name, dim, N, basis, current)
    end
end

Base.length(pool::LoopPool) = size(pool.basis)[2]
Base.size(pool::LoopPool) = length(pool)
Base.show(io::IO, pool::LoopPool) = print(io, pool.basis)
# Base.view(pool::Pool, inds...) = Base.view(pool.pool, inds...)

#index interface for Pool
Base.getindex(pool::LoopPool, i) = pool.basis[:, i]
Base.setindex!(pool::LoopPool, v, i) = setindex!(pool.basis, v, i)
Base.firstindex(pool::LoopPool) = 1
Base.lastindex(pool::LoopPool) = length(pool)

# iterator interface
function Base.iterate(pool::LoopPool)
    if length(pool) == 0
        return nothing
    else
        return (pool.basis[:, 1], 1)
    end
end

function Base.iterate(pool::LoopPool, state)
    if state >= length(pool) || length(pool) == 0
        return nothing
    else
        return (pool.basis[:, state+1], state + 1)
    end
end

function update(pool::LoopPool, variable = rand(eltype(pool.current), pool.dim, pool.N))
    loopNum = size(pool.basis)[1]
    pool.current = view(variable, :, 1:loopNum) * pool.basis
    # B = view(pool.basis, 1:loopNum, :)
    # B = view(pool.basis, 1:loopNum, :)
    # C = pool.current[:, 1:length(pool)]
    # LinearAlgebra.mul!(pool.current, A, pool.basis)
    # LinearAlgebra.BLAS.gemm!('N', 'N', false, variable[:, 1:loopNum], pool.basis[1:loopNum, 1:length(pool)], false, pool.current[:, 1:length(pool)])
end

# current(pool::LoopPool, idx) = pool.current[:, idx]
current(pool::LoopPool, idx) = view(pool.current, :, idx)

function append(pool::LoopPool, basis::AbstractVector)
    for bi in 1:length(pool)
        if pool.basis[:, bi] ≈ basis
            return bi
        end
    end

    pool.basis = hcat(pool.basis, basis)
    pool.current = hcat(pool.current, rand(eltype(pool.current), pool.dim))
    # pool.size += 1
    # @assert pool.size <= size(pool.basis)[2] "Too many loop basis!. Increase maxSize when creates the LoopPool!"
    # pool.basis[:, pool.size] = basis
    return length(pool)
end


export rollout

using StatsBase: wsample

function rollout(policy; epsilon = 0.1, kwargs...)
    rewards = Float32[]
    state = Vector{Float32}[]
    actions = Int[]

    env = CartPoleEnv(;kwargs...)

    action_space = get_action_space(env)

    action_spacec = collect(action_space)

    obs = observe(env)

    while true
        push!(state, obs.state)
        if rand() < epsilon
            action = rand(action_space)
        else
            pos = policy(obs.state)
            if any(isnan, pos)
                error("wtf")
            end
            action = wsample(action_spacec, pos)
        end
        push!(actions, action)
        env(action)
        obs = observe(env)
        push!(rewards, obs.reward)

        get_terminal(obs) && break
    end
    state, rewards, actions
end


using Statistics, StatsBase
using PyCall
xr = pyimport("xarray")
using Blosc
using Elefridge
using JLD
using ZfpCompression
using TranscodingStreams,CodecZstd

ZstdCompressorL3 = ZstdCompressor(level=3)
TranscodingStreams.initialize(ZstdCompressorL3)

ZstdCompressorL10 = ZstdCompressor(level=10)
TranscodingStreams.initialize(ZstdCompressorL10)

ZstdCompressorL22 = ZstdCompressor(level=22)
TranscodingStreams.initialize(ZstdCompressorL22)

path = "/Users/milan/cams/unstructured/"
filelist = filter(x->endswith(x,".grib"),readdir(path))

n = length(filelist)
varnames = fill("",n)
methods = ["ZstdL3","ZstdL10","ZstdL22"]
stats = ["cf","meanerror","absmedian","abs90%","absmax","decmedian","dec90%","decmax"]

E = load("/Users/milan/cams/error/round_zstd_all_unstr.jld")
E = E["E"]

D = load("/Users/milan/cams/entropy/keepbits_98.jld")

## RUN
function rs_zfp(r::Int)
    if r > 6
        rd = 6
    elseif r > 2
        rd = 5
    else
        rd = 4
    end
    return r+rd
end

function statscalc(X0::Array{T,N},Xa::Array{T,N}) where {T,N}
    Xamean = mean(Xa)
    X0mean = mean(X0)
    meanerror = abs(X0mean-Xamean)/X0mean       # normalised

    X0amean = mean(abs.(X0))

    abserr = abs.(X0-Xa)/X0amean                # normalised
    abserr = sort(vec(abserr))
    abs50 = T(quantile(abserr,.5,sorted=true))
    abs90 = T(quantile(abserr,.9,sorted=true))
    absmax = abserr[end]

    # exception cases for decimal error
    Xa0 = Xa .== zero(T)
    X00 = X0 .== zero(T)

    Xas = signbit.(Xa)
    X0s = signbit.(X0)

    decerr = abs.(log10.(abs.(X0./Xa)))

    decerr[Xas .⊻ X0s] .= Inf        # signs don't match
    decerr[Xa0 .⊻ X00] .= Inf        # one is zero the other one not
    decerr[Xa0 .& X00] .= zero(T)    # both zero, set decimal error to zero

    decerr = sort(vec(decerr))
    dec50 = T(quantile(decerr,.5,sorted=true))
    dec90 = T(quantile(decerr,.9,sorted=true))
    decmax = decerr[end]

    return [meanerror,abs50,abs90,absmax,dec50,dec90,decmax]
end

for (i,file) in enumerate(filelist)
    varname = split(split(file,"0_")[end],".")[1]
    varnames[i] = varname

    println("---")
    println("Reading $varname")
    X = xr.open_dataarray(joinpath(path,file),engine="cfgrib").data

    # transpose array
    # use only half the globe
    X = permutedims(X[:,1:(348528÷2)],[2,1])  # unravels as longitude, latitude, level

    # number of bits of information
    r = D["infbits"][D["varnames"] .== varname][1]-9
    r = max(r,2)    # don't use less than two

    println("$r real information bits.")

    Xr = round(X,r)

    ## Round+Zstd
    E[i,1,2:end] = statscalc(X,Xr)
    E[i,2,2:end] = E[i,1,2:end]         # copy across for Zstd
    E[i,3,2:end] = E[i,1,2:end]

    Xr8 = unsafe_wrap(Array, Ptr{UInt8}(pointer(Xr)), sizeof(Xr))

    E[i,1,1] = sizeof(X)/sizeof(transcode(ZstdCompressorL3,Xr8))
    println("ZstdL3 finished.")
    E[i,2,1] = sizeof(X)/sizeof(transcode(ZstdCompressorL10,Xr8))
    println("ZstdL10 finished.")
    E[i,3,1] = sizeof(X)/sizeof(transcode(ZstdCompressorL22,Xr8))
    println("ZstdL22 finished.")

    println(E[i,:,1])
    @save "/Users/milan/cams/error/round_zstd_all_unstr.jld" varnames methods stats E
end


using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities
const MOIB = MOI.Bridges

# Iterations:
# linear5 : > 1000, < 2000
# linear9 : > 3000, < 4000
# linear15: > 20000, Don't know if ever converges so we exclude it
import CDCS
const OPTIMIZER_CONSTRUCTOR = MOI.OptimizerWithAttributes(CDCS.Optimizer, MOI.Silent() => true, "maxIter" => 4000)
const OPTIMIZER = MOI.instantiate(OPTIMIZER_CONSTRUCTOR)

@testset "SolverName" begin
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "CDCS"
end

@testset "supports_allocate_load" begin
    @test MOIU.supports_allocate_load(OPTIMIZER, false)
    @test !MOIU.supports_allocate_load(OPTIMIZER, true)
end

const BRIDGED = MOI.instantiate(OPTIMIZER_CONSTRUCTOR, with_bridge_type=Float64)
const CONFIG = MOIT.TestConfig(atol=3e-2, rtol=3e-2)

@testset "Unit" begin
    MOIT.unittest(BRIDGED, CONFIG, [
        # `NumberOfThreads` not supported.
        "number_threads",
        # `TimeLimitSec` not supported.
        "time_limit_sec",
        # Need to investigate...
        "solve_with_lowerbound", "solve_affine_deletion_edge_cases", "solve_blank_obj", "solve_single_variable_dual_min", "solve_single_variable_dual_max",
        # Need https://github.com/JuliaOpt/MathOptInterface.jl/issues/529
        "solve_qp_edge_cases",
        # Error using cdcs_hsde.preprocess (line 14)
        # No variables in your problem?
        "solve_unbounded_model",
        # Integer and ZeroOne sets are not supported
        "solve_integer_edge_cases", "solve_objbound_edge_cases",
        "solve_zero_one_with_bounds_1",
        "solve_zero_one_with_bounds_2",
        "solve_zero_one_with_bounds_3"])
end

@testset "Continuous linear problems" begin
    MOIT.contlineartest(BRIDGED, CONFIG, [
        # Need to investigate...
        "linear12", "linear15"])
end

@testset "Continuous conic problems" begin
    MOIT.contconictest(BRIDGED, CONFIG, [
        # ITERATION_LIMIT
        "normone1v", "normone1f",
        # rotatedsoc2: Returns Inf and -Inf instead of infeasibility certificate
        "rotatedsoc2",
        # Need to investigate...
        "psdt3", "psds3",
        # Unsupported cones
        "pow", "dualpow", "rootdets", "exp", "dualexp", "logdet", "normspec", "normnuc", "relentr"])
end


#!/usr/bin/env julia
#=
This is a general test about the imputation rate from LD-7k to MD-17k.
I randomly sampled 50 ID to mask loci that are on 17k, but not on 7k.
Then I used the resot of the 4,796 ID as refererence.
=#

using DelimitedFiles, Plots, Statistics
dat = readdlm("num.txt")
dat = reshape(dat, 2, :)'
n = length(dat[:, 1])
ave = mean(dat[:, 1])

plot(dat[:, 1], label="Error rates", dpi=300)
plot!([1, n], [ave, ave], label="Mean error rates", ylabel="Error rate", dpi=300)

savefig("fix-2-fix.png")


import Base.Sort: QuickSort, MergeSort, InsertionSort

Pkg.add("SortingAlgorithms")
using SortingAlgorithms #Provides the other sorting algorithms

include("../perfutil.jl")

sorts = [InsertionSort, QuickSort, MergeSort, HeapSort, RadixSort, TimSort]

randstr_fn!(str_len::Int) = d -> (for i = 1:length(d); d[i] = randstring(str_len); end; d)
randint_fn!(m::Int) = d -> rand!(1:m,d)

# If we're reporting to codespeed, only do a few tests.
if codespeed
    for (T, typename, randfn!) in [(Int, string(Int), randint_fn!(10)),
                                    (String, "String_10", randstr_fn!(10))]
        for size in [2^6,2^16]
            for s in sorts
                if s == InsertionSort && size != 2^6; continue; end
                data = Array(T, size)
                gc()

                ## Random
                name = "$(typename)_$(size)_$(string(s)[1:end-5])_random"
                desc = "$(string(s)) run on $(size) $(typename) elements in random order"
                @timeit_init(sort!(data, alg=s), randfn!(data), name, "")

                name = "$(typename)_$(size)_$(string(s)[1:end-5])_append"
                desc = "$(string(s)) run on $(size) $(typename) elements pre-sorted 10 random elements appended"
                @timeit_init(sort!(data, alg=s), begin data[end-9:end]=randfn!(Array(T,10)) end, name, "")
            end
        end
    end
else
    for (T, typename, randfn!) in [(Int, string(Int), randint_fn!(10)),
                                    (Float64, string(Float64), rand!),
                                    (String, "String_05", randstr_fn!(5)),
                                    (String, "String_10", randstr_fn!(10))]
        for logsize = 6:2:18
            size = 2^logsize
            for s in sorts
                if s == RadixSort && T == String continue end      #Radix sort not implemented
                if s == InsertionSort && logsize >=14 continue end #Too slow
                println(s, s==RadixSort, s, typename, typename==String, logsize)
                data = Array(T, size)
                gc()

                ## Random
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_random"
                @timeit_init(sort!(data, alg=s), randfn!(data), name, "")

                ## Sorted
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_sorted"
                @timeit(sort!(data, alg=s), name, "")

                ## Reverse sorted
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_reversed"
                @timeit_init(sort!(data, alg=s), reverse!(data), name, "")

                ## Sorted with 3 exchanges
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_3exchanges"
                @timeit_init(sort!(data, alg=s),
                             begin
                                 for i = 1:3
                                     n1 = rand(1:size)
                                     n2 = rand(1:size)
                                     data[n1], data[n2] = data[n2], data[n1]
                                 end
                             end,
                             name, "")

                ## Sorted with 10 unsorted values appended
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_append"
                @timeit_init(sort!(data, alg=s), begin data[end-9:end]=randfn!(Array(T,10)) end, name, "")

                ## Random data with 4 unique values
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_4unique"
                @timeit_init(sort!(data4, alg=s), begin data4=data[rand(1:4,size)] end, name, "")

                ## All values equal
                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_allequal"
                data1 = data[ones(Int, size)]
                @timeit(sort!(data1, alg=s), name, "")

                ## QuickSort median killer
                if s == QuickSort && logsize > 16; continue; end  # too slow!

                name = "$(typename)_$(logsize)_$(string(s)[1:end-5])_qsortkiller"
                data = data[1:size>>1]
                data = sort!(data, alg=s)
                data = vcat(reverse(data), data)
                @timeit_init(sort!(qdata, alg=s), begin qdata=copy(data) end, name, "")
            end
        end
    end
end


"""
    struct LagrangianRefFE{D} <: ReferenceFE{D}
      data::GenericRefFE{D}
      face_own_nodeids::Vector{Vector{Int}}
      own_nodes_permutations::Vector{Vector{Int}},
    end

Type representing a Lagrangian finite element. In addition to all the information
provided by a `ReferenceFE`, this type also provides "node-based" information in the following
fields:

- `face_own_nodeids::Vector{Vector{Int}}`: nodes owned by each face
- `own_nodes_permutations::Vector{Vector{Int}}`: permutations of the nodes when the vertices are permuted

For this type

-  `get_dofs(reffe)` returns a `LagrangianDofBasis`
-  `get_prebasis(reffe)` returns a `MonomialBasis`
-  `ReferenceFE{N}(reffe,faceid) where N` returns a `LagrangianRefFE{N}`
"""
struct LagrangianRefFE{D} <: ReferenceFE{D}
  data::GenericRefFE{D}
  face_own_nodeids::Vector{Vector{Int}}
  own_nodes_permutations::Vector{Vector{Int}}
  @doc """
      LagrangianRefFE(
        polytope::Polytope{D},
        prebasis::MonomialBasis,
        dofs::LagrangianDofBasis,
        face_own_nodeids::Vector{Vector{Int}},
        own_nodes_permutations::Vector{Vector{Int}},
        reffaces...) where D

  Low level (inner) constructor of `LagrangianRefFE`.
  """
  function LagrangianRefFE(
    polytope::Polytope{D},
    prebasis::MonomialBasis,
    dofs::LagrangianDofBasis,
    face_own_nodeids::Vector{Vector{Int}},
    own_nodes_permutations::Vector{Vector{Int}},
    reffaces...) where D

    face_own_dofids = _generate_nfacedofs(face_own_nodeids,dofs.node_and_comp_to_dof)
    own_dofs_permutations = _find_own_dof_permutaions(own_nodes_permutations,dofs.node_and_comp_to_dof,face_own_nodeids,face_own_dofids)

    data = GenericRefFE(
      polytope,prebasis,dofs,face_own_dofids;
      own_dofs_permutations = own_dofs_permutations,
      reffaces = reffaces)

    new{D}(data,face_own_nodeids,own_nodes_permutations)
  end
end

num_dofs(reffe::LagrangianRefFE) = reffe.data.ndofs

get_polytope(reffe::LagrangianRefFE) = reffe.data.polytope

get_prebasis(reffe::LagrangianRefFE) = reffe.data.prebasis

get_dofs(reffe::LagrangianRefFE) = reffe.data.dofs

get_face_own_dofids(reffe::LagrangianRefFE) = reffe.data.face_own_dofids

get_own_dofs_permutations(reffe::LagrangianRefFE) = reffe.data.own_dofs_permutations

get_shapefuns(reffe::LagrangianRefFE) = reffe.data.shapefuns

"""
    get_node_coordinates(reffe::LagrangianRefFE) -> Vector{Point{D,Float64}}

Returns the nodal coordinates of the underlying `LagrangianDofBasis`.
"""
get_node_coordinates(reffe::LagrangianRefFE) = reffe.data.dofs.nodes

"""
    get_dof_to_node(reffe::LagrangianRefFE) -> Vector{Int}

Returns the field `dof_to_node` of the underlying `LagrangianDofBasis`.
"""
get_dof_to_node(reffe::LagrangianRefFE) = reffe.data.dofs.dof_to_node

"""
    get_dof_to_comp(reffe::LagrangianRefFE) -> Vector{Int}

Returns the field `dof_to_comp` of the underlying `LagrangianDofBasis`.
"""
get_dof_to_comp(reffe::LagrangianRefFE) = reffe.data.dofs.dof_to_comp

"""
    get_node_and_comp_to_dof(reffe::LagrangianRefFE) -> Vector

Returns the field `node_and_comp_to_dof` of the underlying `LagrangianDofBasis`.
"""
get_node_and_comp_to_dof(reffe::LagrangianRefFE) = reffe.data.dofs.node_and_comp_to_dof

function ReferenceFE{N}(reffe::LagrangianRefFE,iface::Integer) where N
  ReferenceFE{N}(reffe.data,iface)
end

function ReferenceFE{D}(reffe::LagrangianRefFE{D},iface::Integer) where D
  @assert iface==1 "Only one D-face"
  reffe
end

# Helpers for LagrangianRefFE

function _generate_nfacedofs(nfacenodes,node_and_comp_to_dof)
  faces = 1:length(nfacenodes)
  T = eltype(node_and_comp_to_dof)
  comps = 1:n_components(T)
  nfacedofs = [Int[] for i in faces]
  for face in faces
    nodes = nfacenodes[face]
    # Node major
    for comp in comps
      for node in nodes
        comp_to_dofs = node_and_comp_to_dof[node]
        dof = comp_to_dofs[comp]
        push!(nfacedofs[face],dof)
      end
    end
  end
  nfacedofs
end

function _find_own_dof_permutaions(node_perms,node_and_comp_to_dof,nfacenodeids,nfacedofsids)
  dof_perms = Vector{Int}[]
  T = eltype(node_and_comp_to_dof)
  ncomps = n_components(T)
  idof_to_dof = nfacedofsids[end]
  inode_to_node = nfacenodeids[end]
  for inode_to_pinode in node_perms
    ninodes = length(inode_to_pinode)
    nidofs = ncomps*ninodes
    idof_to_pidof = fill(INVALID_PERM,nidofs)
    for (inode,ipnode) in enumerate(inode_to_pinode)
      if ipnode == INVALID_PERM
        continue
      end
      node = inode_to_node[inode]
      pnode = inode_to_node[ipnode]
      comp_to_pdof = node_and_comp_to_dof[pnode]
      comp_to_dof = node_and_comp_to_dof[node]
      for comp in 1:ncomps
        dof = comp_to_dof[comp]
        pdof = comp_to_pdof[comp]
        idof = findfirst(i->i==dof,idof_to_dof)
        ipdof = findfirst(i->i==pdof,idof_to_dof)
        idof_to_pidof[idof] = ipdof
      end
    end
    push!(dof_perms,idof_to_pidof)
  end
  dof_perms
end

# Construction of LagrangianRefFE from Polytopes

"""
    LagrangianRefFE(::Type{T},p::Polytope,orders) where T
    LagrangianRefFE(::Type{T},p::Polytope,order::Int) where T

Builds a `LagrangianRefFE` object on top of the given polytope. `T` is the type of
the value of the approximation space (e.g., `T=Float64` for scalar-valued problems,
`T=VectorValue{N,Float64}` for vector-valued problems with `N` components). The arguments `order` or `orders`
are for the polynomial order of the resulting space, which allows isotropic or anisotropic orders respectively
(provided that the cell topology allows the given anisotropic order). The argument `orders` should be an
indexable collection of `D` integers (e.g., a tuple or a vector), being `D` the number of space dimensions.

In order to be able to use this function, the type of the provided polytope `p` has to implement the
following additional methods. They have been implemented for `ExtrusionPolytope` in the library. They
need to be implemented for new polytope types in order to build Lagangian reference elements on top of them.

- [`compute_monomial_basis(::Type{T},p::Polytope,orders) where T`](@ref)
- [`compute_own_nodes(p::Polytope,orders)`](@ref)
- [`compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)`](@ref)

The following methods are also used in the construction of the `LagrangianRefFE` object. A default implementation
of them is available in terms of the three previous methods. However, the user can also implement them for
new polytope types increasing customization possibilities.

- [`compute_nodes(p::Polytope,orders)`](@ref)
- [`compute_own_nodes_permutations(p::Polytope, interior_nodes)`](@ref)
- [`compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T`](@ref)

"""
function LagrangianRefFE(::Type{T},p::Polytope{D},orders) where {T,D}
  prebasis = compute_monomial_basis(T,p,orders)
  nodes, face_own_nodeids = compute_nodes(p,orders)
  dofs = LagrangianDofBasis(T,nodes)
  interior_nodes = dofs.nodes[face_own_nodeids[end]]
  own_nodes_permutations = compute_own_nodes_permutations(p, interior_nodes)
  reffaces = compute_lagrangian_reffaces(T,p,orders)
  LagrangianRefFE(p,prebasis,dofs,face_own_nodeids,own_nodes_permutations,reffaces...)
end

function MonomialBasis(::Type{T},p::Polytope,orders) where T
  compute_monomial_basis(T,p,orders)
end

function LagrangianDofBasis(::Type{T},p::Polytope,orders) where T
  nodes, _ = compute_nodes(p,orders)
  LagrangianDofBasis(T,nodes)
end

# Constructors taking Int

function LagrangianRefFE(::Type{T},p::Polytope{D},order::Int) where {T,D}
  orders = tfill(order,Val{D}())
  LagrangianRefFE(T,p,orders)
end

function MonomialBasis(::Type{T},p::Polytope{D},order::Int) where {D,T}
  orders = tfill(order,Val{D}())
  MonomialBasis(T,p,orders)
end

function LagrangianDofBasis(::Type{T},p::Polytope{D},order::Int) where {T,D}
  orders = tfill(order,Val{D}())
  LagrangianDofBasis(T,p,orders)
end

# Queries needed to be implemented for polytopes in order to use them
# for building LagrangianRefFEs in a seamless way

"""
    compute_monomial_basis(::Type{T},p::Polytope,orders) where T -> MonomialBasis

Returns the monomial basis of value type `T` and order per direction described by `orders`
on top of the polytope `p`.
"""
function compute_monomial_basis(::Type{T},p::Polytope,orders) where T
  @abstractmethod
end

"""
    compute_own_nodes(p::Polytope{D},orders) where D -> Vector{Point{D,Float64}}

Returns the coordinates of the nodes owned by the interior of the polytope
associated with a Lagrangian space with the order per direction described by `orders`.
"""
function compute_own_nodes(p::Polytope,orders)
  @abstractmethod
end

"""
    compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)

Returns a vector or a tuple with the order per direction at the face `face`
of the polytope `p` when restricting the order per direction `orders` to this face.
`iface` is the face id of `face` in the numeration restricted to the face dimension.
"""
function compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)
  @abstractmethod
end

"""
    compute_nodes(p::Polytope,orders)

When called

    node_coords, face_own_nodeids = compute_nodes(p,orders)

Returns `node_coords`, the nodal coordinates of all the Lagrangian nodes associated with the order per direction
`orders`, and `face_own_nodeids`, being a vector of vectors indicating which nodes are owned by each of
the faces of the polytope `p`.
"""
function compute_nodes(p::Polytope,orders)
  _compute_nodes(p,orders)
end

"""
    compute_own_nodes_permutations(
      p::Polytope, own_nodes_coordinates) -> Vector{Vector{Int}}

Returns a vector of vectors with the permutations of the nodes owned by the interior of the
polytope.
"""
function compute_own_nodes_permutations(p::Polytope, interior_nodes)
  _compute_node_permutations(p, interior_nodes)
end

"""
    compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T

Returns a tuple of length `D` being the number of space dimensions.
The entry `d+1` of this tuple contains a vector of `LagrangianRefFE`
one for each face of dimension `d` on the boundary of the polytope.
"""
function compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T
  _compute_lagrangian_reffaces(T,p,orders)
end

# Default implementations

function _compute_nodes(p,orders)
  if any( orders .== 0)
    _compute_constant_nodes(p,orders)
  elseif all(orders .== 1)
    _compute_linear_nodes(p)
  else
    _compute_high_order_nodes(p,orders)
  end
end

function _compute_constant_nodes(p,orders)
  @assert all( orders .== 0) "If an order is 0 in some direction, it should be 0 also in the others"
  x = compute_own_nodes(p,orders)
  facenodes = [Int[] for i in 1:num_faces(p)]
  push!(facenodes[end],1)
  x, facenodes
end

function _compute_linear_nodes(p)
  x = get_vertex_coordinates(p)
  facenodes = [Int[] for i in 1:num_faces(p)]
  for i in 1:num_vertices(p)
    push!(facenodes[i],i)
  end
  x, facenodes
end

function _compute_high_order_nodes(p::Polytope{D},orders) where D
  nodes = Point{D,Float64}[]
  facenodes = [Int[] for i in 1:num_faces(p)]
  _compute_high_order_nodes_dim_0!(nodes,facenodes,p)
  for d in 1:(num_dims(p)-1)
    _compute_high_order_nodes_dim_d!(nodes,facenodes,p,orders,Val{d}())
  end
  _compute_high_order_nodes_dim_D!(nodes,facenodes,p,orders)
  (nodes, facenodes)
end

function _compute_high_order_nodes_dim_0!(nodes,facenodes,p)
  x = get_vertex_coordinates(p)
  k = 1
  for vertex in 1:num_vertices(p)
    push!(nodes,x[vertex])
    push!(facenodes[vertex],k)
    k += 1
  end
end

@noinline function _compute_high_order_nodes_dim_d!(nodes,facenodes,p,orders,::Val{d}) where d
  x = get_vertex_coordinates(p)
  offset = get_offset(p,d)
  k = length(nodes)+1
  for iface in 1:num_faces(p,d)
    face = Polytope{d}(p,iface)
    face_ref_x = get_vertex_coordinates(face)
    face_prebasis = MonomialBasis(Float64,face,1)
    change = inv(evaluate(face_prebasis,face_ref_x))
    face_shapefuns = change_basis(face_prebasis,change)
    face_vertex_ids = get_faces(p,d,0)[iface]
    face_x = x[face_vertex_ids]
    face_orders = compute_face_orders(p,face,iface,orders)
    face_interior_nodes = compute_own_nodes(face,face_orders)
    face_high_x = evaluate(face_shapefuns,face_interior_nodes)*face_x
    for xi in 1:length(face_high_x)
      push!(nodes,face_high_x[xi])
      push!(facenodes[iface+offset],k)
      k += 1
    end
  end
end

function _compute_high_order_nodes_dim_D!(nodes,facenodes,p,orders)
  k = length(nodes)+1
  p_high_x = compute_own_nodes(p,orders)
  for xi in 1:length(p_high_x)
    push!(nodes,p_high_x[xi])
    push!(facenodes[end],k)
    k += 1
  end
end

_compute_node_permutations(::Polytope{0}, interior_nodes) = [[1]]

function _compute_node_permutations(p, interior_nodes)
  vertex_to_coord = get_vertex_coordinates(p)
  lbasis = MonomialBasis(Float64,p,1)
  change = inv(evaluate(lbasis,vertex_to_coord))
  lshapefuns = change_basis(lbasis,change)
  perms = get_vertex_permutations(p)
  map = evaluate(lshapefuns,interior_nodes)
  pvertex_to_coord = similar(vertex_to_coord)
  node_perms = Vector{Int}[]
  tol = 1.0e-10
  for vertex_to_pvertex in perms
    node_to_pnode = fill(INVALID_PERM,length(interior_nodes))
    pvertex_to_coord[vertex_to_pvertex] = vertex_to_coord
    pinterior_nodes = map*pvertex_to_coord
    for node in 1:length(interior_nodes)
      x = interior_nodes[node]
      pnode = findfirst(i->norm(i-x)<tol,pinterior_nodes)
      if pnode != nothing
         node_to_pnode[node] = pnode
      end
    end
    push!(node_perms,node_to_pnode)
  end
  node_perms
end

_compute_lagrangian_reffaces(::Type{T},p::Polytope{0},orders) where T = ()

function _compute_lagrangian_reffaces(::Type{T},p::Polytope{D},orders) where {T,D}
  reffaces = [ LagrangianRefFE{d}[]  for d in 0:D ]
  p0 = Polytope{0}(p,1)
  reffe0 = LagrangianRefFE(T,p0,())
  for vertex in 1:num_vertices(p)
    push!(reffaces[0+1],reffe0)
  end
  offsets = get_offsets(p)
  for d in 1:(num_dims(p)-1)
    offset = offsets[d+1]
    for iface in 1:num_faces(p,d)
      face = Polytope{d}(p,iface)
      face_orders = compute_face_orders(p,face,iface,orders)
      refface = LagrangianRefFE(T,face,face_orders)
      push!(reffaces[d+1],refface)
    end
  end
  tuple(reffaces...)
end

# Particular implementation for ExtrusionPolytope

function compute_monomial_basis(::Type{T},p::ExtrusionPolytope{D},orders) where {D,T}
  extrusion = Tuple(p.extrusion.array)
  terms = _monomial_terms(extrusion,orders)
  MonomialBasis{D}(T,orders,terms)
end

function compute_own_nodes(p::ExtrusionPolytope{D},orders) where D
  extrusion = Tuple(p.extrusion.array)
  if all(orders .== 0)
    _interior_nodes_order_0(p)
  else
    _interior_nodes(extrusion,orders)
  end
end

function _interior_nodes_order_0(p)
  x = get_vertex_coordinates(p)
  x0 = sum(x) / length(x)
  [x0,]
end

function compute_face_orders(p::ExtrusionPolytope,face::ExtrusionPolytope{D},iface::Int,orders) where D
  d = num_dims(face)
  offset = get_offset(p,d)
  nface = p.nfaces[iface+offset]
  face_orders = _extract_nonzeros(nface.extrusion,orders)
  face_orders
end

function compute_nodes(p::ExtrusionPolytope{D},orders) where D
  _nodes, facenodes = _compute_nodes(p,orders)
  if any( orders .== 0)
    return (_nodes, facenodes)
  end
  terms = _coords_to_terms(_nodes,orders)
  nodes = _terms_to_coords(terms,orders)
  (nodes, facenodes)
end

# Helpers for the ExtrusionPolytope-related implementation

function _monomial_terms(extrusion::NTuple{D,Int},orders) where D
  terms = CartesianIndex{D}[]
  if D == 0
    push!(terms,CartesianIndex(()))
    return terms
  end
  _check_orders(extrusion,orders)
  M = mutable(VectorValue{D,Int})
  term = zero(M)
  _orders = M(orders)
  k = 0
  _add_terms!(terms,term,extrusion,_orders,D,k)
  terms
end

function _interior_nodes(extrusion::NTuple{D,Int},orders) where D
  _check_orders(extrusion,orders)
  terms = CartesianIndex{D}[]
  M = mutable(VectorValue{D,Int})
  term = zero(M)
  _orders = M(orders)
  k = 1
  _add_terms!(terms,term,extrusion,_orders,D,k)
  _terms_to_coords(terms,orders)
end

function _check_orders(extrusion,orders)
  D = length(extrusion)
  @assert length(orders) == D "container of orders not long enough"
  _orders = collect(orders)
  if extrusion[D] == HEX_AXIS
    _orders[D] = 0
  end
  for d in (D-1):-1:1
    if (extrusion[d] == HEX_AXIS || d == 1) && _orders[d+1] == 0
      _orders[d] = 0
    end
  end
  nz = _orders[_orders .!= 0]
  if length(nz) > 1
    @assert all(nz .== nz[1]) "The provided anisotropic order is not compatible with polytope topology"
  end
  nothing
end

function _add_terms!(terms,term,extrusion,orders,dim,k)
  _term = copy(term)
  _orders = copy(orders)
  indexbase = 1
  for i in k:(_orders[dim]-k)
    _term[dim] = i + indexbase
    if dim > 1
      if (extrusion[dim] == TET_AXIS) && i != 0
        _orders .-= 1
      end
      _add_terms!(terms,_term,extrusion,_orders,dim-1,k)
    else
      push!(terms,CartesianIndex(Tuple(_term)))
    end
  end
end

function _coords_to_terms(coords::Vector{<:Point{D}},orders) where D
  indexbase = 1
  terms = CartesianIndex{D}[]
  P = Point{D,Int}
  t = zero(mutable(P))
  for x in coords
    for d in 1:D
      t[d] = round(x[d]*orders[d]) + indexbase
    end
    term = CartesianIndex(Tuple(t))
    push!(terms,term)
  end
  terms
end

function  _terms_to_coords(terms::Vector{CartesianIndex{D}},orders) where D
  P = Point{D,Float64}
  indexbase = 1
  nodes = P[]
  x = zero(mutable(P))
  for t in terms
    for d in 1:D
      x[d] = (t[d] - indexbase) / orders[d]
    end
    node = P(x)
    push!(nodes,node)
  end
  nodes
end

function _extract_nonzeros(mask,values)
  b = Int[]
  for (m,n) in zip(mask,values)
    if (m != 0)
      push!(b, n)
    end
  end
  return Tuple(b)
end


VERSION >= v"0.4.0-dev+6521" && __precompile__()
module WeightedNorms

@static if VERSION < v"0.7-"
    using Compat: AbstractRange
end

include("norms.jl")

include("traces.jl")

end # module


###
### Interface Implementations for running full sampling loops
###

using Printf

StatsBase.sample(rng::AbstractRNG, model::AbstractModel, sampler::Nested; kwargs...) =
    mcmcsample(rng, model, sampler, nested_isdone; progressname="Nested Sampling", chain_type=Chains, kwargs...)

StatsBase.sample(model::AbstractModel, sampler::Nested; kwargs...) =
    StatsBase.sample(GLOBAL_RNG, model, sampler; kwargs...)

function nested_isdone(rng, model, smapler, samples, state, i; progress=true, maxiter=Inf, maxcall=Inf, dlogz=0.5, maxlogl=Inf, kwargs...)
    # 1) iterations exceeds maxiter
    done_sampling = state.it > maxiter
    # 2) number of loglike calls has been exceeded
    done_sampling |= state.ncall > maxcall
    # 3) remaining fractional log-evidence below threshold
    logz_remain = maximum(state.logl) + state.logvol
    delta_logz = logaddexp(state.logz, logz_remain) - state.logz
    done_sampling |= delta_logz < dlogz
    # 4) last dead point loglikelihood exceeds threshold
    done_sampling |= state.logl_dead > maxlogl
    # 5) number of effective samples
    # TODO

    if progress
        str = @sprintf "iter=%d\tncall=%d\tΔlogz=%.2g\tlogl=%.2g\tlogz=%.2g" i state.ncall delta_logz state.logl_dead state.logz
        print("\r\33[2K", str)
    end

    return done_sampling
end


module EMLPIDs
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...types.EMLtypes
	import EMLtypes.length
	include("PIDs/PID.jl")
	include("PIDs/FirstOrder.jl")
	include("PIDs/StepSignal.jl")
	include("PIDs/ConstantSignal.jl")
end


using Base: splat
using DataStructures: BinaryHeap

test = """
    #############
    #...........#
    ###B#C#B#D###
      #A#D#C#A#
      #########
    """

# solution for Q1 test
s1 = [(:C, :AB, 1), (:B, :C, 1), (:B, :BC, 1), (:AB, :B, 2), (:A, :B, 1),
      (:D, :CD, 1), (:D, :right, 1), (:CD, :D, 2), (:BC, :D, 1), (:right, :A, 1)]

# solution for Q2 test
s2 = [(:D, :right, 2), (:D, :left, 2), (:C, :right, 1), (:C, :CD, 1),
      (:C, :left, 1), (:B, :C, 3), (:B, :C, 2), (:B, :BC, 1), (:B, :AB, 1),
      (:BC, :B, 4), (:CD, :B, 3), (:right, :B, 2), (:D, :C, 1), (:D, :right, 1),
      (:AB, :D, 4), (:A, :B, 1), (:A, :D, 3), (:A, :AB, 1), (:left, :A, 3),
      (:left, :A, 2), (:AB, :D, 2), (:right, :A, 1), (:right, :D, 1)]

puzzle = """
    #############
    #...........#
    ###D#B#A#C###
      #B#D#A#C#
      #########
    """

extra_lines = """
          #D#C#B#A#
          #D#B#A#C#
        """

function parse_input(input::AbstractString; extra::Bool = false)
    lines = split(input, '\n'; keepempty = false)
    @assert lines[1] == "#############"
    @assert lines[2] == "#...........#"
    @assert lines[5] == "  #########"
    if extra
        lines = [lines[1:3]
                 split(extra_lines, '\n'; keepempty = false)
                 lines[4:5]]
    end
    lines = [split(line[3:11], "") for line in lines[3:end - 1]] |>
        splat(hcat) |> permutedims
    @assert all(==("#"), lines[1:end, 1:2:9])
    lines = lines[1:end, 2:2:8] .|> Symbol |> Matrix{Any}

    return (left  = Any[nothing, nothing],
            A     = lines[:, 1],
            AB    = Any[nothing],
            B     = lines[:, 2],
            BC    = Any[nothing],
            C     = lines[:, 3],
            CD    = Any[nothing],
            D     = lines[:, 4],
            right = Any[nothing, nothing])
end

const ROOMS = [:A, :B, :C, :D]
const HALLS = [:left, :AB, :BC, :CD, :right]
const COSTS = Dict(:A => 1, :B => 10, :C => 100, :D => 1000)
const PATHS = let paths = Dict(),
    places = [:left, :A, :AB, :B, :BC, :C, :CD, :D, :right]
    for from in places, to in places
        (a, b) = string.(places[clamp.(sort(indexin([from, to], places)), 2, 8)])
        path = []
        while a != b && last(a) != first(b)
            a = "$(last(a))$(last(a) + 1)"
            push!(path, Symbol(a))
        end
        paths[from => to] = path
    end
    paths
end

distance(from, to) = 2 * length(PATHS[from => to])
iscomfy(b, t) = all(v -> v === nothing || v == t, b[t])
Base.isempty(b, t) = all(isnothing, b[t])
Base.isempty(b, from, to) = all(h -> isempty(b, h), PATHS[from => to])
top(b, t) = b[t][findfirst(!isnothing, b[t])]

function moves(b)::Vector
    ms = []
    for room in ROOMS
        !iscomfy(b, room) || continue
        # from room to hallway
        for hall in HALLS
            isempty(b, room, hall) && append!(ms, moves(b, room, hall))
        end
        # from room to room - not strictly necessary, but this way we can
        # verify the solver with the official solutions
        t = top(b, room)
        iscomfy(b, t) && isempty(b, room, t) && push!(ms, move(b, room, t))
    end
    for hall in HALLS
        !isempty(b, hall) || continue
        # from hallway to room
        t = top(b, hall)
        iscomfy(b, t) && isempty(b, hall, t) && push!(ms, move(b, hall, t))
    end
    return ms
end

function moves(b, from, to)
    ms = []
    to ∈ [:left, :right] && b[to][1] == nothing &&
        push!(ms, move(b, from, to, 1))
    isempty(b, to) && push!(ms, move(b, from, to))
    return ms
end

function move(b, from, to, j = nothing)
    i = findfirst(!isnothing, b[from])
    @assert i !== nothing "Cannot move from $from to $to:\n$b"
    if j === nothing
        j = findfirst(!isnothing, [b[to]; 1]) - 1
    end
    t = b[from][i]
    ch = (; from => replace(b[from], i, nothing),
          to => replace(b[to], j, t))
    m = merge(b, ch)
    return ((distance(from, to) + i + j) * COSTS[t], m, (from, to, j))
end

function Base.replace(v::AbstractVector{<: T}, i::Integer, x::T) where T
    return [v[1:i - 1]; x; v[i + 1:end]]
end

function isfinal(b)
    return all(hall -> isempty(b, hall), HALLS) &&
        all(room -> iscomfy(b, room), ROOMS)
end

function search(b)
    h = BinaryHeap(Base.By(first), [(0, b, [])])
    closed = Set{UInt64}()
    while true
        (c, b, p) = pop!(h)
        hb = hash(b)
        hb ∈ closed && continue
        isfinal(b) && return (cost = c, path = p)
        push!(closed, hb)
        for (cm, bm, m) in moves(b)
            push!(h, (c + cm, bm, [p; m]))
        end
    end
end

function check(b, path, cost)
    for (step, (from, to, j)) in enumerate(path)
        m = move(b, from, to, j)
        @assert m ∈ moves(b) "#$step $b $from $to[$j]"
        (c, b, _) = m
        cost -= c
    end
    @assert cost == 0 "residual cost"
end

let b1 = parse_input(test),
    b2 = parse_input(test; extra = true)

    check(b1, s1, 12521)
    check(b2, s2, 44169)
    @assert search(b1).cost == 12521
    @assert search(b2).cost == 44169
end

b1 = parse_input(puzzle)
b2 = parse_input(puzzle; extra = true)
println("Q1: ", search(b1).cost)
println("Q2: ", search(b2).cost)


# this is a script to convert the representations to an easily readable format
using PdfFraud
using JLD2
using FileIO
using ProgressMeter

# first gather data and their labels
data_path = joinpath(PdfFraud.proto_path, "protos")
protos = joinpath.(data_path, readdir(data_path))

function load_representation(file::String)
	data = PdfFraud.read_pb(file)
	y = PdfFraud.visual_representation(data)
	(length(y) == 0) ? nothing : PdfFraud.reshape_representation(y)
end

outfile = joinpath(PdfFraud.proto_path, "extracted_representations/data.jld2")

sha256 = []
representations = []
N = length(protos)
p = Progress(N,1)
for (i, proto) in enumerate(protos)
	x = load_representation(proto)
	if x != nothing
		push!(representations, x)
		push!(sha256, split(basename(proto), ".")[1])
	end
	next!(p)
end

# also get labels
indexf = joinpath(PdfFraud.proto_path, "index.json")
labelf = "/home/vit/vyzkum/bulletproof/pdf_experiment/gt-clustering-pdfs.txt"
labels, labeled_sha256, filenames = PdfFraud.get_labels(sha256, indexf, labelf);

@save outfile sha256 representations labels