text:
################################################################################
#                   CORE DISPATCHVARIABLEREF METHOD EXTENSIONS
################################################################################
# Extend dispatch_variable_ref
function dispatch_variable_ref(
    model::InfiniteModel,
    index::PointVariableIndex
    )::PointVariableRef
    return PointVariableRef(model, index)
end

# Extend _add_data_object
function _add_data_object(
    model::InfiniteModel,
    object::VariableData{<:PointVariable}
    )::PointVariableIndex
    return MOIUC.add_item(model.point_vars, object)
end

# Extend _data_dictionary (type based)
function _data_dictionary(
    model::InfiniteModel,
    ::Type{PointVariable}
    )::MOIUC.CleverDict{PointVariableIndex, VariableData{PointVariable{GeneralVariableRef}}}
    return model.point_vars
end

# Extend _data_dictionary (reference based)
function _data_dictionary(
    vref::PointVariableRef
    )::MOIUC.CleverDict{PointVariableIndex, VariableData{PointVariable{GeneralVariableRef}}}
    return JuMP.owner_model(vref).point_vars
end

# Extend _data_object
function _data_object(
    vref::PointVariableRef
    )::VariableData{PointVariable{GeneralVariableRef}}
    object = get(_data_dictionary(vref), JuMP.index(vref), nothing)
    if isnothing(object)
        error("Invalid point variable reference, cannot find ",
        "corresponding variable in the model. This is likely ",
        "caused by using the reference of a deleted variable.")
    end
    return object
end

# Extend _core_variable_object
function _core_variable_object(vref::PointVariableRef
    )::PointVariable{GeneralVariableRef}
    return _data_object(vref).variable
end

################################################################################
#                          DEFINTION HELPER METHODS
################################################################################
"""
    Point{V, VT <: VectorTuple} <: InfOptVariableType

A `DataType` to assist in making point variables. This can be passed as an
extra argument to `@variable` to make such a variable:
```julia
@variable(model, var_expr, Point(inf_var, parameter_values...), args...,
          kwargs...)
```
Here `parameter_values` must match the format of the infinite parameter
references associated with the infinite variable `inf_var` and can be comprised
of both real valued supports.

**Fields**
- `infinite_variable_ref::V`
- `parameter_values::VT`: The infinite parameter support values the
   variable will depend on.
"""
struct Point{V, VT <: Collections.VectorTuple} <: InfOptVariableType
    infinite_variable_ref::V
    parameter_values::VT
    function Point(vref::V, vt::VT) where {V, VT <: Collections.VectorTuple}
        return new{V, VT}(vref, vt)
    end
end
function Point(ivref, vals...)
    return Point(ivref, Collections.VectorTuple(vals))
end

# Ensure parameter values match shape of parameter reference tuple stored in the
# infinite variable reference
function _check_tuple_shape(
    _error::Function,
    ivref::Union{InfiniteVariableRef, DerivativeRef},
    values::Collections.VectorTuple
    )::Nothing
    prefs = raw_parameter_refs(ivref)
    if !Collections.same_structure(prefs, values)
        _error("The dimensions and array formatting of the infinite parameter ",
               "values must match those of the parameter references for the ",
               "infinite variable/derivative.")
    end
    return
end

## Dispatch methods for checking the supports of parameters in point variable
# IndependentParameterRefs
function _check_element_support(
    _error::Function,
    prefs::Vector{IndependentParameterRef},
    param_values::Vector{Float64},
    counter::Int
    )::Int
    for pref in prefs
        if !supports_in_domain(param_values[counter], infinite_domain(pref))
            _error("Parameter values violate parameter bounds.")
        end
        counter += 1
    end
    return counter
end

# DependentParameterRefs
function _check_element_support(
    _error::Function,
    prefs::Vector{DependentParameterRef},
    param_values::Vector{Float64},
    counter::Int
    )::Int
    len = length(prefs)
    supp = reshape(param_values[counter:counter+len-1], len, 1)
    if !supports_in_domain(supp, infinite_domain(prefs))
        _error("Parameter values violate parameter bounds.")
    end
    return counter += len
end

# Used to ensure values don't violate parameter bounds
function _check_tuple_values(
    _error::Function,
    ivref::Union{InfiniteVariableRef, DerivativeRef},
    param_values::Vector{Float64}
    )::Nothing
    raw_prefs = raw_parameter_refs(ivref)
    counter = 1
    for i in 1:size(raw_prefs, 1)
        prefs = map(e -> dispatch_variable_ref(e), raw_prefs[i, :])
        counter = _check_element_support(_error, prefs, param_values, counter)
    end
    return
end

# Update point variable info to consider the infinite variable
function _update_point_info(
    info::JuMP.VariableInfo,
    ivref::Union{InfiniteVariableRef, DerivativeRef},
    point::Vector{Float64}
    )::JuMP.VariableInfo
    if JuMP.has_lower_bound(ivref) && !info.has_fix && !info.has_lb
        info = JuMP.VariableInfo(true, JuMP.lower_bound(ivref),
                                 info.has_ub, info.upper_bound,
                                 info.has_fix, info.fixed_value,
                                 info.has_start, info.start,
                                 info.binary, info.integer)
    end
    if JuMP.has_upper_bound(ivref) && !info.has_fix && !info.has_ub
        info = JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                 true, JuMP.upper_bound(ivref),
                                 info.has_fix, info.fixed_value,
                                 info.has_start, info.start,
                                 info.binary, info.integer)
    end
    if JuMP.is_fixed(ivref) && !info.has_fix  && !info.has_lb  && !info.has_ub
        info = JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                 info.has_ub, info.upper_bound,
                                 true, JuMP.fix_value(ivref),
                                 info.has_start, info.start,
                                 info.binary, info.integer)
    end
    if !isnothing(start_value_function(ivref)) && !info.has_start
        if _is_vector_start(ivref)
            start = start_value_function(ivref)(point)
        else
            prefs = raw_parameter_refs(ivref)
            start = start_value_function(ivref)(Tuple(point, prefs)...)
        end
        info = JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                 info.has_ub, info.upper_bound,
                                 info.has_fix, info.fixed_value,
                                 true, Float64(start),
                                 info.binary, info.integer)
    end
    if JuMP.is_binary(ivref) && !info.integer
        info = JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                 info.has_ub, info.upper_bound,
                                 info.has_fix, info.fixed_value,
                                 info.has_start, info.start,
                                 true, info.integer)
    end
    if JuMP.is_integer(ivref) && !info.binary
        info = JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                 info.has_ub, info.upper_bound,
                                 info.has_fix, info.fixed_value,
                                 info.has_start, info.start,
                                 info.binary, true)
    end
    return info
end

"""
    JuMP.build_variable(_error::Function, info::JuMP.VariableInfo,
                        var_type::Point)::InfiniteVariable{GeneralVariableRef}

Build and return a point variable based on `info` and `var_type`. Errors
if the information stored in `var_type` is invalid. See [`Point`](@ref)
for more information.

**Example**
```julia-repl
julia> y
y(t)

julia> info = VariableInfo(false, 0, false, 0, false, 0, true, 0, false, false);

julia> pt_var = build_variable(error, info, SemiInfinite(y, 0));
```
"""
function JuMP.build_variable(
    _error::Function,
    info::JuMP.VariableInfo,
    var_type::Point;
    extra_kwargs...
    )::PointVariable{GeneralVariableRef}
    # check for unneeded keywords
    for (kwarg, _) in extra_kwargs
        _error("Keyword argument $kwarg is not for use with point variables.")
    end
    # check the infinite variable reference
    ivref = var_type.infinite_variable_ref
    if !(ivref isa GeneralVariableRef)
        _error("Expected an infinite variable/derivative reference dependency ",
               "of type `GeneralVariableRef`, but got an argument of type ",
               "$(typeof(ivref)).")
    end
    dispatch_ivref = dispatch_variable_ref(ivref)
    if !(dispatch_ivref isa Union{InfiniteVariableRef, DerivativeRef})
        _error("Expected an infinite variable/derivative reference dependency,",
               "but got a variable reference of type $(typeof(dispatch_ivref)).")
    end
    # check and format the values
    raw_vals = var_type.parameter_values
    if !(raw_vals.values isa Vector{<:Real})
        _error("Expected parameter values consisting of real numbers.")
    end
    pvalues = Vector{Float64}(raw_vals.values)
    _check_tuple_shape(_error, dispatch_ivref, raw_vals)
    _check_tuple_values(_error, dispatch_ivref, pvalues)
    info = _update_point_info(info, dispatch_ivref, pvalues)
    # enforce parameter significant digits on the values
    prefs = parameter_list(dispatch_ivref)
    for i in eachindex(pvalues)
        pvalues[i] = round(pvalues[i], sigdigits = significant_digits(prefs[i]))
    end
    # make variable and return
    return PointVariable(_make_float_info(info), ivref, pvalues)
end

## Dispatch methods for updating the supports of parameters in point variable
# IndependentParameterRefs
function _add_point_support(
    prefs::Vector{IndependentParameterRef},
    param_values::Vector{Float64},
    counter::Int
    )::Int
    for pref in prefs
        add_supports(pref, param_values[counter], check = false,
                     label = UserDefined)
        counter += 1
    end
    return counter
end

# DependentParameterRefs
function _add_point_support(
    prefs::Vector{DependentParameterRef},
    param_values::Vector{Float64},
    counter::Int
    )::Int
    len = length(prefs)
    supp = reshape(param_values[counter:counter+len-1], len, 1)
    add_supports(prefs, supp, check = false, label = UserDefined)
    return counter += len
end

# Used to add point variable support to parameter supports if necessary
function _update_param_supports(
    ivref::Union{InfiniteVariableRef, DerivativeRef},
    param_values::Vector{Float64}
    )::Nothing
    raw_prefs = raw_parameter_refs(ivref)
    counter = 1
    for i in 1:size(raw_prefs, 1)
        prefs = dispatch_variable_ref.(raw_prefs[i, :])
        counter = _add_point_support(prefs, param_values, counter)
    end
    return
end

# Used to update mapping infinite_to_points
function _update_infinite_point_mapping(
    pvref::PointVariableRef,
    ivref::Union{InfiniteVariableRef, DerivativeRef}
    )::Nothing
    push!(_point_variable_dependencies(ivref), JuMP.index(pvref))
    return
end

# Update the information constraints for a variable that was already created
function _update_info_constraints(info::JuMP.VariableInfo, gvref, vref)::Nothing
    # extract the preliminaries
    old_info = _variable_info(vref)
    old_info == info && return
    model = JuMP.owner_model(vref)

    # update the lower bound constraint as needed
    if info.has_lb && (!old_info.has_lb || info.lower_bound != old_info.lower_bound)
        constr = JuMP.ScalarConstraint(gvref, MOI.GreaterThan(info.lower_bound))
        if old_info.has_lb
            cindex = _lower_bound_index(vref)
            cref = _make_constraint_ref(model, cindex)
            _set_core_constraint_object(cref, constr)
        else
            cref = JuMP.add_constraint(model, constr, is_info_constr = true)
            _set_lower_bound_index(vref, JuMP.index(cref))
        end
    elseif old_info.has_lb && !info.has_lb
        JuMP.delete(model, JuMP.LowerBoundRef(vref))
        _set_lower_bound_index(vref, nothing)
    end

    # update the upper bound constraint as needed
    if info.has_ub && (!old_info.has_ub || info.upper_bound != old_info.upper_bound)
        constr = JuMP.ScalarConstraint(gvref, MOI.LessThan(info.upper_bound))
        if old_info.has_ub
            cindex = _upper_bound_index(vref)
            cref = _make_constraint_ref(model, cindex)
            _set_core_constraint_object(cref, constr)
        else
            cref = JuMP.add_constraint(model, constr, is_info_constr = true)
            _set_upper_bound_index(vref, JuMP.index(cref))
        end
    elseif old_info.has_ub && !info.has_ub
        JuMP.delete(model, JuMP.UpperBoundRef(vref))
        _set_upper_bound_index(vref, nothing)
    end

    # update the fix constraint as needed
    if info.has_fix && (!old_info.has_fix || info.fixed_value != old_info.fixed_value)
        constr = JuMP.ScalarConstraint(gvref, MOI.EqualTo(info.fixed_value))
        if old_info.has_fix
            cindex = _fix_index(vref)
            cref = _make_constraint_ref(model, cindex)
            _set_core_constraint_object(cref, constr)
        else
            cref = JuMP.add_constraint(model, constr, is_info_constr = true)
            _set_fix_index(vref, JuMP.index(cref))
        end
    elseif old_info.has_fix && !info.has_fix
        JuMP.delete(model, JuMP.FixRef(vref))
        _set_fix_index(vref, nothing)
    end

    # update the binary constraint as needed
    if info.binary && !old_info.binary
        constr = JuMP.ScalarConstraint(gvref, MOI.ZeroOne())
        cref = JuMP.add_constraint(model, constr, is_info_constr = true)
        _set_binary_index(vref, JuMP.index(cref))
    elseif old_info.binary && !info.binary
        JuMP.delete(model, JuMP.BinaryRef(vref))
        _set_binary_index(vref, nothing)
    end

    # update the integer constraint as needed
    if info.integer && !old_info.integer
        constr = JuMP.ScalarConstraint(gvref, MOI.Integer())
        cref = JuMP.add_constraint(model, constr, is_info_constr = true)
        _set_integer_index(vref, JuMP.index(cref))
    elseif old_info.integer && !info.integer
        JuMP.delete(model, JuMP.IntegerRef(vref))
        _set_integer_index(vref, nothing)
    end

    # finalize the update
    _update_variable_info(vref, info)
    set_optimizer_model_ready(model, false)
    return
end

"""
    JuMP.add_variable(model::InfiniteModel, var::PointVariable,
                      [name::String = ""])::GeneralVariableRef

Extend the `JuMP.add_variable` function to accomodate `PointVariable` variable
types. Adds a variable to an infinite model `model` and returns a
[`GeneralVariableRef`](@ref). Primarily intended to be an internal function of
the constructor macro `@variable`. However, it can be used in combination with
`JuMP.build_variable` to add variables to an infinite model object.
Errors if an invalid infinite variable reference is included in `var`.

**Example**
```julia-repl
julia> @infinite_parameter(m, t in [0, 10]);

julia> info = VariableInfo(false, 0, false, 0, false, 0, true, 0, false, false);

julia> inf_var = build_variable(error, info, Infinite(t));

julia> ivref = add_variable(m, inf_var, "var_name")
var_name(t)

julia> pt_var = build_variable(error, info, Point(ivref, 0.5));

julia> pvref = add_variable(m, pt_var, "var_alias")
var_alias
```
"""
function JuMP.add_variable(
    model::InfiniteModel,
    v::PointVariable,
    name::String = "";
    add_support = true,
    update_info = true
    )::GeneralVariableRef
    ivref = v.infinite_variable_ref
    divref = dispatch_variable_ref(ivref)
    JuMP.check_belongs_to_model(divref, model)
    existing_index = get(model.point_lookup, (ivref, v.parameter_values), nothing)
    if isnothing(existing_index)
        data_object = VariableData(v, name)
        vindex = _add_data_object(model, data_object)
        vref = PointVariableRef(model, vindex)
        if add_support
            _update_param_supports(divref, v.parameter_values)
        end
        _update_infinite_point_mapping(vref, divref)
        model.point_lookup[(ivref, v.parameter_values)] = vindex
        gvref = _make_variable_ref(model, vindex)
        _set_info_constraints(v.info, gvref, vref)
    else
        gvref = _make_variable_ref(model, existing_index)
        if update_info
            vref = PointVariableRef(model, existing_index)
            _update_info_constraints(v.info, gvref, vref)
        end
        if !isempty(name)
            JuMP.set_name(vref, name)
        end
    end
    model.name_to_var = nothing
    return gvref
end

################################################################################
#                         PARAMETER VALUE METHODS
################################################################################
"""
    infinite_variable_ref(vref::PointVariableRef)::GeneralVariableRef

Return the `InfiniteVariableRef` associated with the point variable `vref`.

**Example**
```julia-repl
julia> @variable(model, T, Infinite(t))
T(t)

julia> @variable(model, T0, Point(T, 0))
T0

julia> infinite_variable_ref(T0)
T(t)
```
"""
function infinite_variable_ref(vref::PointVariableRef)::GeneralVariableRef
    return _core_variable_object(vref).infinite_variable_ref
end

"""
    raw_parameter_values(vref::PointVariableRef)::Vector{Float64}

Return the raw support point values associated with the point variable `vref`.
"""
function raw_parameter_values(vref::PointVariableRef)::Vector{Float64}
    return _core_variable_object(vref).parameter_values
end

"""
    parameter_values(vref::PointVariableRef)::Tuple

Return the support point associated with the point variable `vref`.

**Example**
```julia-repl
julia> @variable(model, T, Infinite(t))
T(t)

julia> @variable(model, T0, Point(T, 0))
T0

julia> parameter_values(T0)
(0,)
```
"""
function parameter_values(vref::PointVariableRef)
    prefs = raw_parameter_refs(infinite_variable_ref(vref))
    return Tuple(raw_parameter_values(vref), prefs)
end

# Internal function used to change the parameter value tuple of a point variable
function _update_variable_param_values(
    vref::PointVariableRef,
    pref_vals::Vector{<:Real}
    )::Nothing
    info = _variable_info(vref)
    ivref = infinite_variable_ref(vref)
    new_var = PointVariable(info, ivref, pref_vals)
    _set_core_variable_object(vref, new_var)
    return
end

################################################################################
#                           VARIABLE INFO METHODS
################################################################################
# Set info for point variables
function _update_variable_info(
    vref::PointVariableRef,
    info::JuMP.VariableInfo
    )::Nothing
    ivref = infinite_variable_ref(vref)
    param_values = raw_parameter_values(vref)
    new_var = PointVariable(info, ivref, param_values)
    _set_core_variable_object(vref, new_var)
    return
end

################################################################################
#                                 DELETION
################################################################################
# Extend _delete_variable_dependencies (for use with JuMP.delete)
function _delete_variable_dependencies(vref::PointVariableRef)::Nothing
    # remove variable info constraints associated with vref
    _delete_info_constraints(vref)
    # remove the infinite variable dependency
    ivref = infinite_variable_ref(vref)
    filter!(e -> e != JuMP.index(vref), _point_variable_dependencies(ivref))
    # remove the lookup entry
    delete!(JuMP.owner_model(vref).point_lookup, (ivref, raw_parameter_values(vref)))
    return
end


text:
function nn_lstm(wordembeds::Matrix{T}, charembeds::Matrix{T}, ntags::Int) where T
    dims_w = Node(name=:dims_w)
    dims_c = Node(name=:dims_c)

    w = lookup(parameter(wordembeds), Node(name=:w))
    c = lookup(parameter(charembeds), Node(name=:c))
    csize = size(charembeds, 1)
    c = Conv1d(T,5,csize,5csize,padding=2)(c, dims_c)
    c = max(c, dims_c)

    h = concat(1, w, c)
    h = dropout(h, 0.5)
    hsize = size(wordembeds,1) + 5csize
    h = LSTM(T,hsize,hsize,1,0.0,true)(h, dims_w)
    h = dropout(h, 0.5)
    h = Linear(T,2hsize,ntags)(h)
    Graph(h)
end

function nn_cnn(wordembeds::Matrix{T}, charembeds::Matrix{T}, ntags::Int) where T
    dims_w = Node(name=:dims_w)
    dims_c = Node(name=:dims_c)

    w = lookup(parameter(wordembeds), Node(name=:w))
    c = lookup(parameter(charembeds), Node(name=:c))
    csize = size(charembeds, 1)
    #c = dropout(c, 0.5)
    c = Conv1d(T,5,csize,5csize,padding=2)(c, dims_c)
    c = max(c, dims_c)

    h = concat(1, w, c)
    #h = dropout(h, 0.5)
    hsize = size(wordembeds,1) + 5csize
    h = Conv1d(T,5,hsize,2hsize,padding=2)(h, dims_w)
    #h = dropout(h, 0.5)
    h = Linear(T,2hsize,ntags)(h)
    Graph(h)
end

struct SLSTM
    nlayers::Int
    W::Var
    b::Var
    h::Var
    c::Var
end

function SLSTM()
    W = parameter()
    b = parameter()
    h = parameter()
end

function (f::SLSTM)(x::Var, dims)
    g = f.h
    hxg = concat(1, f.h, x, g)
    h = linear(hxg, f.W, f.b)
    n = length(h) ÷ 7
    ilrfso = sigmoid(h[1:6n])
    u = tanh(h[6n+1:7n])
    h1 = ilrfso[1:n]
    h2 = ilrfso[n+1:2n]
    h3 = ilrfso[2n+1:3n]
    h4 = ilrfso[3n+1:4n]
    h5 = ilrfso[4n+1:5n]
    h6 = ilrfso[5n+1:6n]
    h2 .* cminus + h4 .*
    h1 .* u
end

function NN(wordembeds1::Var, wordembeds2::Vector{Var}, charembeds::Vector{Var}, ntags::Int)
    T = Float32
    w1 = Node()
    hw1 = lookup(wordembeds1, w1)
    w2 = Node()
    hw2 = lookup(wordembeds2, w2)

    c = Node()
    hc = lookup(charembeds, c)
    d = size(charembeds[1], 1)
    hc = Conv1D(T,5,d,5d,2,1)(hc)
    hc = max(hc, 2)
    hc = resize(hc, batchsize(w1))

    h = concat(1, hw1, hw2, hc)
    d = 200 + 5size(charembeds[1],1)
    dh = 300
    h = Conv1D(T,5,d,dh,2,1)(h)
    h = relu(h)
    #h = swish(h,zerograd([T(1)]))

    istrain = Node()
    for i = 1:2
        h = dropout(h, 0.3, istrain)
        h = Conv1D(T,5,dh,dh,2,1)(h)
        h = relu(h)
        #h = swish(h,zerograd([T(1)]))
    end
    h = Linear(T,dh,ntags)(h)
    g = Graph(input=(w1,w2,c,istrain), output=h)
    NN(g)
end
(f::SLSTM)(x::Node, dims) = Node(f, (x,dims))

function slstm_tstep(x::Var, dims)
    h = f.h
    c = f.c
    cg = f.cg
    g = h
    for l = 1:nlayers
        h = conv1d_index(h, dims, window=1)
        ci = conv1d_index(c, dims, window=1)
        cg = repeat(cg)
        h = concat(1, h, x, g)
        h = linear(h, W, b)
        n = length(h) ÷ 7
    end
end

ext: jl


text:
module NeuralQuantumState

export run_NQS, NETSETTINGS

using Distributed
using SharedArrays
using Parameters
using Random
using LinearAlgebra
using ProgressMeter
using DelimitedFiles
using StatsBase
using Distributions
using Plots
import GR
using Plots.PlotMeasures

include("structs.jl")
include("weights.jl")
include("network.jl")
include("statistics.jl")
include("plotting.jl")
include("conjugate_gradient.jl")
include("optimizer.jl")
include("driver.jl")

end # module


text:
"""
XML Utilities.

$(DocStringExtensions.IMPORTS)
$(DocStringExtensions.EXPORTS)
"""
module XmlUtils

using EzXML
using DocStringExtensions
using AbstractTrees
end # module XmlUtils


text:
ShaderArray = [
        (x⃗⃗::RayResults)->RGB(0.35,0.35,0.35),
        (x⃗::RayResults)->(RGB(0.8,0.8,0.8)*10*clamp01(1/norm(x⃗.Position))),
        (x⃗::RayResults)->(RGB(x⃗.Grad[1]/2+0.5,x⃗.Grad[2]/2+0.5,x⃗.Grad[3]/2+0.5)),
        (x⃗::RayResults)->RGB(0.22,0.545,0.133)*(abs(x⃗.Position[1]/(5*sqrt(2)))),
        (x⃗::RayResults)->RGB(0.584,0.345,0.689)*(abs(x⃗.Position[1]/(5*sqrt(2)))),
        (x⃗::RayResults)->RGB(0.796,0.235,0.2)*(abs(x⃗.Position[1]/(5*sqrt(2)))),
        (x⃗::RayResults)->RGB(0.35,0.35,0.55)]
## Max length of array is 2^16
## A sharder is (currenty) (subject to change due to unhandled reflections) a function that takes a RayResults struct and returns a RGB value


text:
module SlaterKoster

using Reexport

# Functionality we need - by order of
#  (2) AD derivatives
#        - w.r.t. U
#        - w.r.t. hop
#        all further derivatives are handles outside
#  (3) correct f and higher orbitals : implemented -> need to check!
#  (4) fast evaluation
#  (5) fast derivatives

# my old codes from the previous TB package
#include("oldsk.jl")

# basic SK machinery
include("basics.jl")

# evaluate the SK blocks
include("code_generation.jl")
include("skcore.jl")

# hamiltonian matrix assembly
include("matrixassembly.jl")
# exports: hamiltonian

# Some basic models - mostly for testing?
include("models.jl")
@reexport using ACEtb.SlaterKoster.Models
# This exports
#  - KwonTB

end # module


text:
#Note -- depends on test_syllablize.jl running first.
info("Testing syllable dominant frequency tracing...")

@test dftrace(a, syllables_boxed) == [[3, 3], [4, 2], [3]]
@test dftrace(a, syllables) == dftrace(a, syllables_boxed)


text:
module Test_pmtools

using LinearAlgebra
using MatrixPencils
using Polynomials
using Test

@testset "Polynomial Matrix Tools" begin

@testset "Conversions" begin

p = pm2poly(rand(2,3,3),:s);
@test p == pm2poly(poly2pm(p,grade=5),:s)

s = Polynomial([0, 1],:s)
v = [1+2*s;1+2*s+3*s^2]
@test v == pm2poly(poly2pm(v),:s)[:]

@test 2 == pmdeg(v)

@test [1,1] == pmeval(v,0)[:]

@test v == pm2poly(poly2pm(v),:s)[:]

s = Polynomial([0, 1],:s);
p = 1+2*s+3*s^2
@test p == pm2poly(poly2pm(p),:s)[1,1]

@test ([-2 -1; 1 0], [3 0; 0 1]) == pm2lpCF1(p)
@test ([-2 1; -1 0], [3 0; 0 1]) == pm2lpCF2(p)

@test reverse(p.coeffs) == pmreverse(p)[:]

@test 3im == pm2poly(poly2pm(3im))[1,1]

a = rand(2,3);
@test a == reshape(poly2pm(a),2,3)

a = rand(2,3);
@test [a zeros(size(a)...)] == reshape(poly2pm(a,grade=1),2,6)

a = rand(3);
@test a == reshape(poly2pm(a),3)

a = rand(3);
@test [a zeros(size(a)...)] == reshape(poly2pm(a,grade=1),3,2)

a = transpose(rand(3));
@test a == reshape(poly2pm(a),1,3)

a = transpose(rand(3));
@test [a zeros(size(a)...)] == reshape(poly2pm(a,grade=1),1,6)

a = rand(Complex{Float64},3);
@test a' == reshape(poly2pm(a'),1,3)

a = rand(Complex{Float64},3);
@test [a' zeros(size(a')...)] == reshape(poly2pm(a',grade=1),1,6)

end # conversions

@testset "Basic polynomial matrix manipulations" begin

P = rand(0,0,4);
@test pmeval(P,2) ≈ lseval(pm2ls(P)[1:5]...,2)

P = rand(0,0,4);
@test pmeval(P,2) ≈ lpseval(pm2lps(P)[1:8]...,2)

P = rand(2,3,4);
@test pmeval(P,2) ≈ lseval(pm2ls(P)[1:5]...,2)

P = rand(2,3,4);
@test pmeval(P,2) ≈ lpseval(pm2lps(P)[1:8]...,2)

P = rand(Complex{Float64},2,3,4);
@test pmeval(P,2im) ≈ lseval(pm2ls(P)[1:5]...,2im)

P = rand(Complex{Float64},2,3,4);
@test pmeval(P,2im) ≈ lpseval(pm2lps(P)[1:8]...,2im)

P = rand(0,0,4);
@test P[:,:,1:pmdeg(P)+1] ≈ ls2pm(pm2ls(P)[1:5]...)

P = rand(0,0,4);
@test P[:,:,1:pmdeg(P)+1] ≈ lps2pm(pm2lps(P)[1:8]...)

P = rand(2,3,4);
@test P ≈ ls2pm(pm2ls(P,contr = true)[1:5]...)

P = rand(2,3,4);
@test P ≈ ls2pm(pm2ls(P,obs = true)[1:5]...)

P = rand(2,3,4);
@test P ≈ ls2pm(pm2ls(P,contr = true, obs = true)[1:5]...)

P = rand(2,3,4);
@test P ≈ ls2pm(pm2ls(P)[1:5]...)

P = rand(3,2,4);
@test P ≈ ls2pm(pm2ls(P,contr = true)[1:5]...)

P = rand(3,2,4);
@test P ≈ ls2pm(pm2ls(P,obs = true)[1:5]...)

P = rand(3,2,4);
@test P ≈ ls2pm(pm2ls(P,contr = true, obs = true)[1:5]...)

P = rand(3,2,4);
@test P ≈ ls2pm(pm2ls(P)[1:5]...)

P = rand(2,3,4);
@test P ≈ ls2pm(pm2ls(P,minimal = true)[1:5]...)

P = rand(2,3,4);
@test P ≈ lps2pm(pm2lps(P)[1:8]...)

P = rand(Complex{Float64},2,3,4);
@test P ≈ ls2pm(pm2ls(P)[1:5]...)

P = rand(Complex{Float64},2,3,4);
@test P ≈ lps2pm(pm2lps(P)[1:8]...)

P = rand(2,3,4);
@test P ≈ pmreverse(pmreverse(P,6))[:,:,1:4]

P = rand(0,0,4);
@test P[:,:,1:pmdeg(P)+1] ≈ pmreverse(pmreverse(P,2))[:,:,1:pmdeg(P)+1]

end # Basic manipulations

@testset "Tests pmdivrem" begin

# Example Vlad Ionescu
s = Polynomial([0, 1],:s);
NUM = [s^2+3*s+3 1; -1 2*s^2+7*s+4];
DEN = [(s+1)^2 s+2; (s+1)^3 (s+1)*(s+2)]
X = divrem.(NUM,DEN)
Q = first.(X)
R = last.(X)
Q = [x[1] for x in X]
R = [x[2] for x in X]
Q1, R1 = pmdivrem(NUM,DEN)
@test Q1 == poly2pm(Q) && R1 == poly2pm(R)

@test reshape(Q1[:,1,:],2,1,1) == poly2pm(Q[:,1]) && reshape(R1[:,1,:],2,1,2) == poly2pm(R[:,1])

NUM = rand(1,1,4); DEN = rand(1,1,3);
num = Polynomial(NUM[1,1,1:4]); den = Polynomial(DEN[1,1,1:3]);
q, r = divrem(num,den)
@test num ≈ den*q+r
Q, R = pmdivrem(NUM,DEN)
@test num ≈ den*Polynomial(Q[1,1,1:size(Q,3)])+Polynomial(R[1,1,1:size(R,3)])

NUM = rand(0,0,4); DEN = rand(0,0,3);
Q, R = pmdivrem(NUM,DEN)
@test Q == Array{Float64}(undef,0,0,0) && R ==  Array{Float64}(undef,0,0,0)

NUM = rand(2,3,4); DEN = rand(2,3,3);
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

NUM = rand(2,3,4); DEN = ones(2,3,1);
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

NUM = rand(2,3,4); DEN = rand(2,3,3);
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

NUM = rand(2,3,4); DEN = ones(2,3,1);
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

NUM = zeros(2,3,4); DEN = ones(2,3,1);
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

NUM = rand(2,3,4); DEN = NUM;
Q, R = pmdivrem(NUM,DEN)
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

try
  NUM = zeros(2,3,4); DEN = zeros(2,3,1);
  Q, R = pmdivrem(NUM,DEN)
  @test false
catch
  @test true
end

Ty = Complex{Float64}
NUM = rand(Ty,2,3,4); DEN = rand(Ty,2,3,2);
Q, R = pmdivrem(NUM,DEN);
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

Ty = Int
NUM = rand(Ty,2,3,4); DEN = rand(Ty,2,3,2);
Q, R = pmdivrem(NUM,DEN);
@test all(pm2poly(NUM) .≈ pm2poly(DEN) .* pm2poly(Q) .+ pm2poly(R))

end # end pmdivrem

@testset "Structured polynomial matrix linearization tools" begin

A,E,B,C,D = spm2ls(1,1,1,1,atol = 1.e-7, minimal=true)
@test iszero(D)

A,E,B,C,D = spm2ls(1,1,1,1)
@test iszero(-C*inv(A)*B+D) && iszero(E)

A,E,B,F,C,G,D,H = spm2lps(1,1,1,1,atol = 1.e-7, minimal=true)
@test iszero(D) && iszero(H)

A,E,B,F,C,G,D,H = spm2lps(1,1,1,1)
@test iszero(-C*inv(A)*B+D) && iszero(E) && iszero(F) && iszero(G) && iszero(H)

#  simple test
D = rand(3,3);
W = zeros(3,3);
A2,E2,B2,C2,D2 = spm2ls(D,D,D,W,atol = 1.e-7, minimal=true)
@test D2 ≈ -D

D = rand(3,3);
W = zeros(3,3,1);
A2,E2,B2,C2,D2 = spm2ls(D,D,D,W,atol = 1.e-7, minimal=true)
@test D2 ≈ -D

A2,E2,B2,C2,D2 = spm2ls(D,D,D,W,atol = 1.e-7)
@test -C2*inv(A2)*B2+D2 ≈ -D && iszero(E2)

#  simple test
D = rand(3,3);
W = zeros(3,3);
sys = spm2lps(D,D,D,W,atol = 1.e-7, minimal=true)
@test sys[7] ≈ -D  && sys[8] ≈ 0*D

#  simple test
D = rand(3,3);
W = zeros(3,3,1);
sys = spm2lps(D,D,D,W,atol = 1.e-7, minimal=true)
@test sys[7] ≈ -D  && sys[8] ≈ 0*D

#
D = rand(3,3,2); V = zeros(3,3,1); V[:,:,1] = Matrix{eltype(V)}(I,3,3); W = zeros(3,3);
A2,E2,B2,C2,D2 = spm2ls(D,D,V,W,atol = 1.e-7, minimal=true)
@test D2 ≈ Matrix{eltype(V)}(I,3,3)

D = rand(3,3,2); V = zeros(3,3,1); V[:,:,1] = Matrix{eltype(V)}(I,3,3); W = zeros(3,3);
sys = spm2lps(D,D,V,W,atol = 1.e-7, minimal=true)
@test sys[7] ≈ Matrix{eltype(V)}(I,3,3) && sys[8] ≈ 0*W

D = rand(3,3,4); W = zeros(3,3);
sys2 = spm2ls(D,D,D,W,atol = 1.e-7, minimal=true)
sys1 = pm2ls(D)
@test lsequal(sys1[1:5]..., sys2...,atol1 = 1.e-7,atol2 = 1.e-7)

D = rand(3,3,4); W = zeros(3,3);
sys2 = spm2lps(D,D,D,W,atol = 1.e-7, minimal=true)
sys1 = pm2lps(D)
@test lpsequal(sys1..., sys2...,atol1 = 1.e-7,atol2 = 1.e-7)

T = rand(3,3,4); U = rand(3,3,2); V = rand(3,3,4); W = rand(3,3,3);
sys = spm2ls(T,U,V,W);
@test lseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

T = rand(3,3,4); U = rand(3,3,2); V = rand(3,3,4); W = rand(3,3,3);
sys = spm2lps(T,U,V,W);
@test lpseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

T = rand(Complex{Float64},3,3,4); U = rand(3,3,2); V = rand(3,3,4); W = rand(3,3,3);
sys = spm2ls(T,U,V,W);
@test lseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

T = rand(Complex{Float64},3,3,4); U = rand(3,3,2); V = rand(3,3,4); W = rand(3,3,3);
sys = spm2lps(T,U,V,W);
@test lpseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

T = rand(0,0,4); U = rand(0,3,2); V = rand(3,0,4); W = rand(3,3,3);
sys = spm2ls(T,U,V,W);
@test lseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

T = rand(0,0,4); U = rand(0,3,2); V = rand(3,0,4); W = rand(3,3,3);
sys = spm2lps(T,U,V,W);
@test lpseval(sys...,1) ≈ pmeval(V,1)*(pmeval(T,1)\pmeval(U,1))+pmeval(W,1)

# Example Rosenbrock
s = Polynomial([0, 1],:s);
T = [s+1 s^3+2s^2; s^2+3s+2 s^4+4s^3+4s^2+s+2];
U = [s^2+1;s^3+2s^2+s+3];
V = -[s^2+3s+1 s^4+4s^3+4s^2-1];
W = [s^3+2s^2+s+2];

# G = (2*s + 3)/(s^2 + 3*s + 2)
T1 = [s+1 0;0 s+2];
U1 = [1;1];
V1 = [1 1];
W1 = [0];

sys = spm2ls(T,U,V,W,atol = 1.e-7, minimal=true);
sys1 = spm2ls(T1,U1,V1,W1,atol = 1.e-7, minimal=true);
@test lsequal(sys1...,sys...,atol1=1.e-7,atol2=1.e-7)

syslps = spm2lps(T,U,V,W,atol = 1.e-7, minimal=true);
syslps1 = spm2lps(T1,U1,V1,W1,atol = 1.e-7, minimal=true);
@test lpsequal(syslps1...,syslps...,atol1=1.e-7,atol2=1.e-7)

#
s = Polynomial([0, 1],:s);
T = [s+1 s^3+2s^2; s^2+3s+2 s^4+4s^3+4s^2+s+2];
U = [s^2+1;s^3+2s^2+s+3];
V = -[s^2+3s+1 s^4+4s^3+4s^2-1];
W = [s^3+2s^2+s+2];

T1 = [s+1 0;0 s+2];
U1 = [1;1];
V1 = [1 1];
W1 = [0];
S1 = [-T1 U1; V1 W1];
A, E, B, F, C, G, D, H = pm2lps(S1,obs=true);
c = 3; v1 = lpseval(A, E, B, F, C, G, D, H,c)
A, E, B, F, C, G, D, H = pm2lps(S1,contr=true);
c = 3; v2 = lpseval(A, E, B, F, C, G, D, H,c)
@test v1 ≈ v2

S = [-T U; V W];

A, E, B, F, C, G, D, H = pm2lps(S,obs=true);
c = 3; v1 = lpseval(A, E, B, F, C, G, D, H,c)
A, E, B, F, C, G, D, H = pm2lps(S,contr=true);
c = 3; v2 = lpseval(A, E, B, F, C, G, D, H,c)
@test v1 ≈ v2

#  simple transfer function realization
D = reshape([-2,-1,2,1],1,1,4);
N = reshape([-1,1,0,1],1,1,4);
V = reshape([1],1,1);
W = reshape([0.],1,1);
sys = spm2ls(D,N,V,W,atol = 1.e-7, minimal=true);
sys_poles = eigvals(sys[1],sys[2])
@test sort(sys_poles) ≈ [-2., -1., 1.]
sys_zeros = spzeros(sys...)[1]
@test coeffs(fromroots(sys_zeros)) ≈ [-1,1,0,1]

D = reshape([-2,-1,2,1],1,1,4);
N = reshape([-1,1,0,1],1,1,4);
V = reshape([1],1,1);
W = reshape([0.],1,1);
sys = spm2lps(D,N,V,W,atol = 1.e-7, minimal=true);
sys_poles = eigvals(sys[1],sys[2])
@test sort(sys_poles) ≈ [-2., -1., 1.]
sys_zeros = pzeros([sys[1] sys[3]; sys[5] sys[7]],[sys[2] sys[4]; sys[6] sys[8]])[1]
@test coeffs(fromroots(sys_zeros)) ≈ [-1,1,0,1]

D1 = Polynomial([-2,-1,2,1]);
N1 = Polynomial([-1,1,0,1]);
V1 = 1;
W1 = 0.;
sys = spm2ls(D1,N1,V1,W1,atol = 1.e-7, minimal=true);
sys_poles = eigvals(sys[1],sys[2])
@test sort(sys_poles) ≈ [-2., -1., 1.]
sys_zeros = spzeros(sys...)[1]
@test coeffs(fromroots(sys_zeros)) ≈ [-1,1,0,1]

end # structured linearization

@testset "Linearization tools" begin

# Various simple examples

# P = empty

P = rand(0,0,3);
sys = pm2ls(P);
@test P[:,:,1:pmdeg(P)+1] ≈ ls2pm(sys...)

P = rand(0,0,3);
sys = pm2lps(P);
@test P[:,:,1:pmdeg(P)+1] ≈ lps2pm(sys...)

# P = 0

sys = pm2ls(0)
info = spkstruct(sys...)
@test (info.rki, info.lki,info.id, info.nf) == ([0], [0], [], 0)

sys=pm2lps(0)
info = pkstruct(sys[7],sys[8])
@test (info.rki, info.lki,info.id, info.nf) == ([0], [0], [], 0)

# P = 1

sys = pm2ls(1)
info = spkstruct(sys...)
@test (info.rki, info.lki,info.id, info.nf) == (Int64[], Int64[], [1], 0)

sys=pm2lps(1)
info = pkstruct(sys[7],sys[8])
@test (info.rki, info.lki,info.id, info.nf) == ([], [], [1], 0)

# P = λ

λ = Polynomial([0,1],:λ)

sys = pm2ls(λ)
info = spkstruct(sys...)
@test (info.rki, info.lki,info.id, info.nf) == (Int64[], Int64[], [1, 1], 1)

sys=pm2lps(λ)
info = pkstruct(sys[7],sys[8])
@test (info.rki, info.lki,info.id, info.nf) == ([], [], [], 1)

# P = [λ 1]
λ = Polynomial([0,1],:λ)
P = [λ one(λ)]

sys = pm2ls(P)
info = spkstruct(sys...)
@test (info.rki, info.lki,info.id, info.nf) == ([1], [], [1, 1], 0)

sys=pm2lps(P)
info = pkstruct(sys[7],sys[8])
@test (info.rki, info.lki,info.id, info.nf) == ([1], [], [], 0)

# Example 3: P = [λ^2 λ; λ 1] DeTeran, Dopico, Mackey, ELA 2009
λ = Polynomial([0,1],:λ)
P = [λ^2 λ; λ 1]
@test all(P .≈ pm2poly(ls2pm(pm2ls(P)...),:λ))
@test all(P .≈ pm2poly(lps2pm(pm2lps(P)...),:λ))

A, E, B, C, D = pm2ls(P,minimal = true)
M1 = [A B; C D]
N1 = [E zeros(size(B)...); zeros(size(C)...) zeros(size(D)...)]
val1, info1 = peigvals(M1,N1)
@test val1 ≈ [Inf, Inf] && (info1.rki, info1.lki, info1.id, info1.nf) == ([1], [1], [1, 1], 0)
val1, iz, info1 = pzeros(M1,N1)
@test val1 ≈ Float64[] && iz == [] && (info1.rki, info1.lki, info1.id, info1.nf) == ([1], [1], [1, 1], 0)

λ = Polynomial([0,1],:λ)
P1 = [λ^2 λ; λ 1; 0 0]
A, E, B, C, D = pm2ls(P1,minimal = true)
M1 = [A B; C D]
N1 = [E zeros(size(B)...); zeros(size(C)...) zeros(size(D)...)]
val1, info1 = peigvals(M1,N1)
@test val1 ≈ [Inf, Inf] && (info1.rki, info1.lki, info1.id, info1.nf) == ([1], [0,1], [1, 1], 0)

P = zeros(2,2,3);
P[:,:,1] = [0 0; 0 1.];
P[:,:,2] = [0 1.; 1. 0];
P[:,:,3] = [1. 0; 0 0];
@test P ≈ ls2pm(pm2ls(P)...)
@test P ≈ lps2pm(pm2lps(P)...)

val, info = pmeigvals(P)
@test val == Float64[] && (info.rki, info.lki,info.id, info.nf) == ([1], [1], [], 0)
val, iz, info = pmzeros(P)
@test val == Float64[] && iz == [] && (info.rki, info.lki,info.id, info.nf) == ([1], [1], [], 0)

# build a strong (least order) structured linearization which preserves the finite eigenvalues,
# infinite zeros (with multiplicities), the left and right Kronecker structure
A2,E2,B2,C2,D2 = pm2ls(P)
A = copy(A2); E = copy(E2); C = copy(C2); B = copy(B2); D = copy(D2);
@time A1, E1, B1, C1, D1, nuc, nuo, nse  = lsminreal(A,E,B,C,D)
@test lsequal(A,E,B,C,D,A1,E1,B1,C1,D1) &&
      nuc == 2 && nuo == 0 && nse == 1
@time val, iz, info = spzeros(A1,E1,B1,C1,D1)
@test val ≈ Float64[] && iz == [] &&  (info.rki, info.lki, info.id, info.nf) == ([1], [1], [1, 1], 0)
@time val, info = speigvals(A1,E1,B1,C1,D1)
@test val ≈ [Inf, Inf] &&  (info.rki, info.lki, info.id, info.nf) == ([1], [1], [1, 1], 0)

# build linearizations which preserve the finite eigenvalues and right structure;
# the infinite zeros are not preserved because P is singular and the left Kronecker indices are in excess with 1
#M2,N2 = pm2lp(P,left=false)
M2,N2 = pm2lpCF2(P)
M = copy(M2); N = copy(N2);
@time val, iz, info = pzeros(M,N)
@test val[1:info.nf] ≈ Float64[] && iz == [] &&  (info.rki, info.lki, info.nf) == ([1], [2], 0)
@time val, info = peigvals(M,N)
@test val ≈ Float64[] &&  (info.rki, info.lki, info.nf) == ([1], [2], 0)

# build a linearization which preserves the finite eigenvalues and left structure;
# the infinite zeros are not preserved and the right Kronecker indices are in excess with 1
# M2,N2 = pm2lp(P,left=true)
M2,N2 = pm2lpCF1(P)
M = copy(M2); N = copy(N2);
@time val, iz, info = pzeros(M,N)
@test val[1:info.nf] ≈ Float64[] && iz == [] &&  (info.rki, info.lki, info.nf) == ([2], [1], 0)
@time val, info = peigvals(M,N)
@test val ≈ Float64[] &&  (info.rki, info.lki, info.nf) == ([2], [1], 0)

# Strong reductions

# Example 4: Van Dooren, Dewilde, LAA, 1983
P = zeros(Int,3,3,3)
P[:,:,1] = [1 2 -2; 0 -1 -2; 0 0 0]
P[:,:,2] = [1 3 0; 1 4 2; 0 -1 -2]
P[:,:,3] = [1 4 2; 0 0 0; 1 4 2]
@test P ≈ ls2pm(pm2ls(P)...)
@test P ≈ lps2pm(pm2lps(P)...;atol1=1.e-7,atol2=1.e-7)

# computations as in paper
M, N = pm2lpCF1(P)
kinfo = pkstruct(M,N)
@test (kinfo.rki, kinfo.lki, kinfo.id, kinfo.nf) == ([1], [1], [2], 1)

val, kinfo = peigvals(M,N)
@test val ≈ [1, Inf, Inf] && (kinfo.rki, kinfo.lki, kinfo.id, kinfo.nf) == ([1], [1], [2], 1)

# this produces the correct structure at infinity, and shifted singular structure information
A, E, B, F, C, G, D, H = lpsminreal(pm2lps(P)...,contr=false);
M1 = [A B; C D];
N1 = [E F; G H];
val1, info1 = peigvals(M1,N1)
@test val1 ≈ [1, Inf, Inf] && (info1.rki, info1.lki,info1.id, info1.nf) == ([0], [2], [2], 1)

# this produces the wrong structure at infinity, and shifted singular structure information
A, E, B, F, C, G, D, H = lpsminreal(pm2lps(P)...,obs=false)
M1 = [A B; C D];
N1 = [E F; G H];
val1, info1 = peigvals(M1,N1)
@test val1 ≈ [1, Inf] && (info1.rki, info1.lki,info1.id, info1.nf) == ([0], [1], [1], 1)

# build controllable and observable pencil realizations
sys1  = pm2lps(P,contr=true)
sys2  = pm2lps(P,obs=true)
@test lpsequal(sys1..., sys2...)
@test lpseval(sys1...,2) ≈ lpseval(sys2...,2)

#
@test lps2pm(sys1...,atol1=1.e-7,atol2=1.e-7,val=2) ≈ lps2pm(sys2...,atol1=1.e-7,atol2=1.e-7,val=2)
@test ls2pm(pm2ls(P,contr=true)...,atol1=1.e-7,atol2=1.e-7,val=2) ≈ P[:,:,1:3]

(A, E, B, F, C, G, D, H) = copy.(sys2)
A1, E1, B1, F1, C1, G1, D1, H1, V, W, nuc,nuo = lpsminreal(A, E, B, F, C, G, D, H,atol1=1.e-7,atol2=1.e-7)
@test lpsequal(A1, E1, B1/W, F1/W, V'\C1, V'\G1, V'\D1/W, V'\H1/W, sys2...,atol1=1.e-7,atol2=1.e-7)

# compute zeros structure
M = [A1 B1; C1 D1]
N = [E1 F1; G1 H1]
val, info = peigvals(M,N)
@test val ≈ [1, Inf] && (info.rki, info.lki,info.id, info.nf) == ([0], [1], [1], 1)

# compute poles structure
p, m = size(D1)
n = size(B1,1)
M = [A1 B1 zeros(n,p); C1 D1 -V'; zeros(m,n) W zeros(m,p)]
N = [E1 F1 zeros(n,p); G1 H1 zeros(p,p); zeros(m,n+m+p)]
val, info = peigvals(M,N)
@test val ≈ [Inf, Inf, Inf, Inf, Inf, Inf, Inf] && (info.rki, info.lki,info.id, info.nf) == ([], [], [1, 1, 1, 1, 3], 0)
val, iz, info = pzeros(M,N)
@test val ≈ [Inf, Inf] && (info.rki, info.lki,info.id, info.nf) == ([], [], [1, 1, 1, 1, 3], 0)

(A, E, B, F, C, G, D, H) = copy.(sys1)
A1, E1, B1, F1, C1, G1, D1, H1, V, W, nuc,nuo = lpsminreal(A, E, B, F, C, G, D, H,atol1=1.e-7,atol2=1.e-7,contr=false)
@test lpsequal(A1, E1, B1/W, F1/W, V'\C1, V'\G1, V'\D1/W, V'\H1/W, sys1...,atol1=1.e-7,atol2=1.e-7)

M = [A1 B1; C1 D1]
N = [E1 F1; G1 H1]
val, info = peigvals(M,N)
@test val ≈ [1, Inf] && (info.rki, info.lki,info.id, info.nf) == ([0], [1], [1], 1)

# ensuring strong observability
# Example 4: Van Dooren, Dewilde, LAA 1983
P = zeros(3,3,3)
P[:,:,1] = [1 2 -2; 0 -1 -2; 0 0 0]
P[:,:,2] = [1 3 0; 1 4 2; 0 -1 -2]
P[:,:,3] = [1 4 2; 0 0 0; 1 4 2]

# build controllable and observable pencil realizations
sys1  = pm2lps(P,contr=true)
sys2  = pm2lps(P,obs=true)
@test lpsequal(sys1..., sys2...)
@test lpseval(sys1...,1) ≈ lpseval(sys2...,1)

A2 = [
    1     0     0     0     0     0     0     0     0
    0     1     0     0     0     0     0     0     0
    0     0     1     0     0     0     0     0     0
    0     0     0     1     0     0     0     0     0
    0     0     0     0     1     0     0     0     0
    0     0     0     0     0     1     0     0     0
    0     0     0     0     0     0     1     0     0
    0     0     0     0     0     0     0     1     0
    0     0     0     0     0     0     0     0     1];
E2 = [
    0     0     0     0     0     0     0     0     0
    1     0     0     0     0     0     0     0     0
    0     1     0     0     0     0     0     0     0
    0     0     0     0     0     0     0     0     0
    0     0     0     1     0     0     0     0     0
    0     0     0     0     1     0     0     0     0
    0     0     0     0     0     0     0     0     0
    0     0     0     0     0     0     1     0     0
    0     0     0     0     0     0     0     1     0];
B2 = [
    -1     0     0
     0     0     0
     0     0     0
     0    -1     0
     0     0     0
     0     0     0
     0     0    -1
     0     0     0
     0     0     0];
C2 = [
    0     1     1     0     3     4     0     0     2
    0     1     0     0     4     0     0     2     0
    0     0     1     0    -1     4     0    -2     2];

D2 = zeros(Int,3,3);

sys = (A2, E2, B2, C2, D2)
@time P = ls2pm(sys...)
@test pmeval(P,1) ≈ lseval(sys...,1)

A2 = [
    -2    -3     0     0     0     0     0     0     0
     1     0     0     0     0     0     0     0     0
     0     0    -2    -3     0     0     0     0     0
     0     0     1     0     0     0     0     0     0
     0     0     0     0     1     0     0     0     0
     0     0     0     0     0     1     0     0     0
     0     0     0     0     0     0     1     0     0
     0     0     0     0     0     0     0     1     0
     0     0     0     0     0     0     0     0     1];
 E2 = [
     1     0     0     0     0     0     0     0     0
     0     1     0     0     0     0     0     0     0
     0     0     1     0     0     0     0     0     0
     0     0     0     1     0     0     0     0     0
     0     0     0     0     0     0     0     0     0
     0     0     0     0     1     0     0     0     0
     0     0     0     0     0     0     0     0     0
     0     0     0     0     0     0     1     0     0
     0     0     0     0     0     0     0     1     0];
 B2 = [
     1     0
     0     0
     0     1
     0     0
    -1     0
     0     0
     0    -1
     0     0
     0     0];
 C2 = [
     1     0     1    -3     0     1     0     2     0
     0     1     1     3     0     1     0     0     1];
D2 = zeros(Int,2,2);

try
   ls2pm(A2,E2,B2,C2,D2,atol1 = 1.e-7,atol2=1.e-7)
   @test false
catch
   @test true
end

# Example 3 - (Varga, Kybernetika, 1990)
A2 = [
1 0 0 0 -1 0 0 0
0 1 0 0 0 -1 0 0
0 0 1 0 0 0 0 0
0 0 0 1 0 0 0 0
0 0 0 0 -1 0 0 0
0 0 0 0 0 -1 0 0
0 0 0 0 3 0 1 0
0 0 0 0 0 2 0 1
]
E2 = [
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
1 0 0 0 0 0 0 0
0 1 0 0 0 0 0 0
0 0 1 0 0 0 0 0
0 0 0 1 0 0 0 0
]
B2 = [
      -1 1
      0 0
      0 0
      0 0
      1 -2
      -2 3
      0 0
      3 -3
]
C2 = [
      0 0 0 0 0 0 -1 0
      0 0 0 0 0 0 0 -1
]
D2 = zeros(Int,2,2);

# case of a non-minimal realization of a polynomial matrix with finite eigenvalues
sys = ([A2 B2; C2 D2], [E2 zeros(8,2); zeros(2,10)], [zeros(8,2);I], [zeros(2,8) -I], zeros(2,2))

@time P = ls2pm(sys...,atol1 = 1.e-7,atol2=1.e-7);
@test pmeval(P,1) ≈ lseval(sys...,1)

end # linearization tools

end # testset
end # module


text:
@testset "rng.jl" begin
    @testset "sample distribution" begin
        rng = AdvancedPS.TracedRNG()
        vns = rand(rng, Distributions.Normal())
        AdvancedPS.save_state!(rng)

        rand(rng, Distributions.Normal())

        AdvancedPS.load_state!(rng)
        new_vns = rand(rng, Distributions.Normal())
        @test new_vns ≈ vns
    end

    @testset "split" begin
        rng = AdvancedPS.TracedRNG()
        key = rng.rng.key
        new_key, = AdvancedPS.split(key, 1)

        @test key ≠ new_key

        Random.seed!(rng, new_key)
        @test rng.rng.key === new_key
    end
end


text:
"""
```
solve(m; nojump = false)
```

func_iter MUST return the values in the same order as they are in `m.functional_variables`.
Should also be clear on the requisites for func_iter.
"""
function solve(m::AbstractNLCTFPModel, init_guess::OrderedDict = OrderedDict{Symbol, Vector{Float64}}();
               nojump::Bool = false, nojump_method::Symbol = :ode, update_method::Symbol = :average,
               init_slm_coefs::OrderedDict = OrderedDict{Symbol, Vector{Float64}}(),
               individual_convergence::OrderedDict{Symbol, Vector{Float64}} = OrderedDict{Symbol, Vector{Float64}}(),
               error_calc::Symbol = :total_error, type_checks::Bool = true, return_sol::Bool = false, verbose::Symbol = :none,
               vars_for_error::Vector{Symbol} = Vector{Symbol}(undef, 0),
               has_verbose::Bool = true, kwargs...)

    if nojump
        return solve_nojump(m, init_guess; method = nojump_method, return_sol = return_sol, verbose = verbose, kwargs...)
    else
        ## Construct functional iteration loop
        stategrid, funcvar, derivs, endo = initialize!(m)
        func_iter = eqcond(m)
        funcvar_dict = get_functional_variables(m)
        θ = parameters_to_named_tuple(get_parameters(m))

        # Update the funcvar guess
        for (k, v) in init_guess
            funcvar[k] = v
        end

        if type_checks
            @assert isa(stategrid, StateGrid)
            @assert isa(funcvar,   OrderedDict)
            @assert isa(derivs,    OrderedDict)
            @assert isa(endo,      OrderedDict)
        end

        # Get settings for the loop
        max_iter     = get_setting(m, :max_iter)        # maximum number of functional loops
        func_tol     = get_setting(m, :tol)             # convergence tolerance
        error_method = get_setting(m, :error_method)    # method for calculating error at the end of each loop
        func_errors = Vector{Float64}(undef, max_iter)  # track error for each iteration

        if update_method == :average
            um = (new, old) -> average_update(new, old, get_setting(m, :learning_rate))
        elseif update_method == :pseudo_transient_continuation
            error("Updating method $update_method not implemented.")
        else
            error("Updating method $update_method not implemented.")
        end

        # We can stop updating individual functional variables prematurely.
        # The OrderedDict individual_convergence maps functional variables
        # to a two-element Vector. The first element must be either 1 or 0,
        # and the second element is the tolerance for convergence.
        converge_individually = !isempty(individual_convergence)

        # Some other set up
        proposal_funcvar = deepcopy(funcvar) # so we can calculate difference between proposal and current values
        total_time = [0.]
        guess_slm_coefs = !isempty(init_slm_coefs)

        if isempty(vars_for_error)
            vars_for_error = collect(keys(funcvar_dict))
        end

        # Start loop!
        if verbose != :none
            println("Beginning functional iteration . . .\n")
        end

        success = falses(1)
        for iter in 1:max_iter
            # Get a new guess
            if verbose != :none
                begin_time = time_ns()
            end

            if has_verbose
                if guess_slm_coefs
                    new_funcvar = converge_individually ?
                        func_iter(stategrid, funcvar, derivs, endo, θ; verbose = verbose,
                                  init_coefs = init_slm_coefs,
                                  individual_convergence = individual_convergence) :
                                      func_iter(stategrid, funcvar, derivs, endo, θ; init_coefs = init_slm_coefs, verbose = verbose)
                else
                    new_funcvar = converge_individually ? func_iter(stategrid, funcvar, derivs, endo, θ; verbose = verbose,
                                                                    individual_convergence = individual_convergence) :
                                                                        func_iter(stategrid, funcvar, derivs, endo, θ; verbose = verbose)
                end
            else
                if guess_slm_coefs
                    new_funcvar = converge_individually ?
                        func_iter(stategrid, funcvar, derivs, endo, θ; init_coefs = init_slm_coefs,
                                  individual_convergence = individual_convergence) :
                                      func_iter(stategrid, funcvar, derivs, endo, θ; init_coefs = init_slm_coefs)
                else
                    new_funcvar = converge_individually ? func_iter(stategrid, funcvar, derivs, endo, θ;
                                                                    individual_convergence = individual_convergence) :
                                                                        func_iter(stategrid, funcvar, derivs, endo, θ)
                end
            end

            # Update the guess
            for (k, v) in proposal_funcvar
                proposal_funcvar[k] = um(new_funcvar[funcvar_dict[k]], v) # um = update_method
            end

            # Calculate errors
            func_errors[iter] = calculate_func_error(proposal_funcvar, funcvar, error_method; vars = vars_for_error)

            if verbose != :none
                spaces1, spaces2, spaces3 = if iter < 10
                    repeat(" ", 16), repeat(" ", 14), repeat(" ", 5)
                elseif iter < 100
                    repeat(" ", 17), repeat(" ", 15), repeat(" ", 6)
                else
                    repeat(" ", 18), repeat(" ", 16), repeat(" ", 7)
                end

                println("Iteration $(iter), current error:            $(func_errors[iter])")
                if verbose == :high
                    for k in vars_for_error
                        indiv_funcvar_err = calculate_func_error(proposal_funcvar[k], funcvar[k], error_method)
                        indiv_space = " " ^ (28 - length(string(k)) + 1)
                        println("Error for $(k):" * indiv_space * string(indiv_funcvar_err))
                    end
                else
                    println("")
                end

                loop_time = (time_ns() - begin_time) / 6e10 # 60 * 1e9 = 6e10
                total_time[1] += loop_time
                expected_time_remaining = (max_iter - iter) * loop_time
                println(verbose, :high, "Duration of loop (min):" * spaces1 * "$(round(loop_time, digits = 4))")
                println("Total elapsed time (min):" * spaces2 * "$(round(total_time[1], digits = 4))")
                println(verbose, :high, "Expected max remaining time (min):" * spaces3 * "$(round(expected_time_remaining, digits = 4))")
                println("\n")
            end

            # Convergence?
            for (k, v) in proposal_funcvar
                funcvar[k] .= v
            end

            if func_errors[iter] < func_tol
                if verbose != :none
                    println("Convergence achieved! Final round error: $(func_errors[iter])")
                end
                success[1] = true
                break
            end
        end

        if success[1]
            if verbose != :none
                println("Calculating remaining variables . . .")
                aug_time = time_ns()
            end

            augment_variables!(m, stategrid, funcvar, derivs, endo)

            if verbose != :none
                total_time[1] += (time_ns() - aug_time) / 6e10
            end
        end

        if verbose != :none
            println("Total elapsed time (min): $(round(total_time[1], digits = 4))\n")
        end

        return stategrid, funcvar, derivs, endo, success[1]
    end
end

"""
```
solve_nojump(m::AbstractNLCTModel)
```

solves the no-jump equilibrium in `m`. The only available method currently is `:ode`, which uses
ODE methods. Planned extensions include `:pseudo_transient_continuation` via EconPDEs.jl
and Chebyshev/Smolyak projection via BasisMatrices.jl/SmolyakApprox.jl
"""
function solve_nojump(m::AbstractNLCTFPModel, init_guess::OrderedDict = OrderedDict{Symbol, Vector{Float64}}();
                      method::Symbol = :ode, return_sol::Bool = false, verbose::Symbol = :none, kwargs...)
    stategrid, functional_variables, derivatives, endogenous_variables = initialize_nojump!(m)

    if (method == :ode || method == :ODE) && ndims(stategrid) == 1 # Univariate no jump model => use ODE methods
        s = collect(keys(get_state_variables(m)))[1] # state variable name
        ode_f, ode_callback = eqcond_nojump(m)

        tspan = (stategrid[s][1], stategrid[s][end])
        θ = parameters_to_named_tuple(haskey(get_settings(m), :nojump_parameters) ?
                                      map(x -> get_parameters(m)[get_keys(m)[x]], get_setting(m, :nojump_parameters)) :
                                      get_parameters(m))
        prob = ODEProblem(ode_f, nojump_ode_init(m), tspan, θ, tstops = stategrid[s][2:end - 1])
        sol = solve(prob, get_setting(m, :ode_integrator),
                    reltol = get_setting(m, :ode_reltol),
                    abstol = get_setting(m, :ode_abstol), callback = ode_callback)

        augment_variables_nojump!(m, stategrid, ode_f, functional_variables,
                                  derivatives, endogenous_variables, sol)

        if return_sol
            return stategrid, functional_variables, derivatives, endogenous_variables, sol
        else
            return stategrid, functional_variables, derivatives, endogenous_variables
        end
    elseif method in [:PTC, :pseudo_transient_continuation]
        stategrid, functional_variables, derivatives, endogenous_variables = initialize_nojump!(m)
        N = length(stategrid)
        θ = parameters_to_named_tuple(m.parameters)
        timestep!, resize_fcn! = eqcond_nojump(m)
        function hjb!(ydot, y)
            resize_fcn!(y, functional_variables) # map y into functional_variables
            ydot .= timestep!(stategrid, functional_variables, derivatives, endogenous_variables, θ)
        end

        # Create sparsity matrix
        y0 = vcat(values(init_guess))
        ydot = similar(y0)
        sparsity_pattern = jacobian_sparsity(hjb!, ydot, y0)
        jac = Float64.(sparse(sparsity_pattern))
        colors = matrix_colors(jac)
        if all(jac .== 0.)
            @warn "Automatic matrix coloring failed."
            jac    = nothing
            colors = 1:length(y0)
        end

        # Run pseudo-transient continuation algorithm
        yfinal, distance = finiteschemesolve(hjb!, y0; J0c = (jac, colors), kwargs... )

        # Populate functional variables
        funcvar_names = keys(functional_variables)
        for (i, k) in enumerate(keys(init_guess))
            functional_variables[k] .= yfinal[1 + (i - 1) * N:(i * N)]
        end

        # Rerun time step to ensure the static step is calculated
        hjb!(ydot, yfinal)

        # Finish remaining calculations
        augment_variables_nojump!(m, stategrid, functional_variables,
                                  derivatives, endogenous_variables)

        return stategrid, functional_variables, derivatives, endogenous_variables
    elseif method in [:pseudo_transient_relaxation, :PTR]

    end
end

# TO DO: write a copy of solve_nojump or a wrapper that takes in spec[<0;22;8Mifically an AbstractNLCTDiffusionModel,
# which is for models that do not have any jumps


text:
"""
    newtonRaphson(f,df,p0)

adalah fungsi untuk menyelesaikan `f(x)=0` menggunakan metode newton. Metode ini
memerlukan turunan dari fungsi `f(x)` yaitu `df` dan nilai tebakan awal `p0`.

# Examples
```jl
julia> f(x) = x^3-3*x+2;

julia> df(x) = 3*x^2-3;

julia> pk,flag,M = newtonRaphson(f,df,-2.4);

julia> pk
-2.0

julia> flag
0

julia> M
6×3 Array{Float64,2}:
 0.0  -2.4      NaN
 1.0  -2.07619    0.32381
 2.0  -2.0036     0.0725945
 3.0  -2.00001    0.00358742
 4.0  -2.0        8.58992e-6
 5.0  -2.0        4.91913e-11
```
return solusi `p0` dengan `flag` bernilai 0 jika metode bisection berhasil menemukan
solusi dan gagal jika tidak 0. Serta, matriks `M` yang berisi catatan proses tiap
iterasi `[k, pk, error]`
"""
function newtonRaphson(f,df,p0)
    # Definisikan nilai toleransi, maksimum iterasi dan tebakan awal
    delta = 10^-7;
    maxi = 100;
    pk = p0
    M = [0 pk NaN];
    # Mulai langkah iterasi
    flag = 1;
    for k = 2:maxi
        # Rumus metode Newton-Raphson
        pk1 = pk
        pk = pk1 - f(pk1)/df(pk1);
        # Hitung nilai galat mutlak dan relatif
        err = abs(pk-pk1);
        rel = 2*err/(abs(pk)+eps());
        M = [M; [k-1 pk err]]
        # Kriteria penghentian iterasi jika galat memenuhi toleransi.
        if err<delta || rel<delta
            flag = 0;
            break
        end
    end
    return pk, flag, M
end


text:
function check_symmetry(tab::Tableau{T}; atol=16eps(T), rtol=16eps(T)) where {T}
    a, b, s = tab.a, tab.b, tab.s
    [isapprox(a[s+1-i, s+1-j] + a[i,j], b[j]; atol=atol, rtol=rtol) for i in axes(a,1), j in axes(a,2)]
end

function issymmetric(tab::Tableau{T}; kwargs...) where {T}
    all(check_symmetry(tab; kwargs...))
end

function check_symmetry(tab::PartitionedTableau{T}; kwargs...) where {T}
    (check_symmetry(tab.q; kwargs...),
     check_symmetry(tab.p; kwargs...))
end

function issymmetric(tab::PartitionedTableau{T}; kwargs...) where {T}
    issymmetric(tab.q) && issymmetric(tab.p)
end


text:
using Base.Test

@testset "getindex compoundinterpolation" begin
	@test CompoundInterpolation(LinearInterpolation, NearestInterpolation)[1.25, 2.2] == InterpolationSupport[(.75,(1,2)), (.25,(2,2))]
end;

@testset "getindex array compoundinterpolation" begin
	@test rand(10,10,10)[CompoundInterpolation(LinearInterpolation, NearestInterpolation, LinearInterpolation), 1.1, 1.25, 2.2] < 1
end;

max_stars_repo_name: JKrehl/Interpol

max_issues_repo_name: JKrehl/Interpol

max_forks_repo_name: JKrehl/Interpol


text:
###############################################################################
#
# create DataFrame
#
###############################################################################
"""
Create a data frame from given parameters:

$(METHODLIST)

"""
function dataframe end

"""
If no parameter is passed to `dataframe`, an empty DataFrame will be returned:

    dataframe()

Return an empty DataFrame

---
Example
```julia
df = dataframe();
```
"""
dataframe() = DataFrame()

"""
When a vector of data and a vector of headers are given, `dataframe` returns a DataFrame with the given data and header:

    dataframe(v_data::Vector, v_name::Vector{String})

Return a DataFrame for given
- `v_data` Vector of data
- `v_name` Vector of data header

Note that the data and header vectors must have the same length.
"""
dataframe(v_data::Vector, v_name::Vector{String}) =
(
    @assert length(v_data) == length(v_name);

    _df = dataframe();

    for _i in eachindex(v_data)
        _df[!,v_name[_i]] = v_data[_i];
    end;

    return _df
)


text:
# MIT License

#Copyright (c) 2021 hondoRandale <jules.rasetaharison@tutanota.com> and contributors

#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:

#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE.

module SimpleFWA

  using InvertedIndices
  using Random
  using Statistics

  export FWA, simpleFWA, metropolisHastingsFWA

  mutable struct FWA
    X::Matrix{Float32};
    fitness_fireworks::Vector{Float32};
    S;
    fitness_sparks;
    x_b;
    y_min::Float32;
    iter::Int;
    err_conv::Float32;
  end

  function dynExplosionAmplitude( y_min::Float32,
                                  y_cf::Float32,
                                  A_cf_old::Float32,
                                  C_a::Float32,
                                  C_r::Float32 )
    A_cf = A_cf_old::Float32
    if y_min - y_cf < 0.0f0
      A_cf *= C_a;
    else
      A_cf *= C_r;
    end
    return A_cf
  end

  function explosionSparks( ;X_i, A_i, s, lower, upper, rng )
    nSparks = size( s, 2 );
    d       = size( s, 1 );
    for j ∈ 1:1:nSparks
      idz = bitrand( rng, d );
      for ii ∈ 1:1:d
        if idz[ii]
          @inbounds ΔX_k    = ( A_i * ( rand( rng, Float32, 1 )[1] * 2.0f0 - 1.0f0 ) )::Float32;
          @inbounds val     = ( X_i[ii] + ΔX_k )::Float32;
          if val >= lower[ii] && val <= upper[ii]
            @inbounds s[ii,j] = val::Float32;
          else
            @inbounds s[ii,j] = ( rand( rng, Float32, 1 )[1] * ( upper[ii] - lower[ii] ) + lower[ii] )::Float32;
          end
        end
      end
    end
  end

  function drawFireworksPositionsBoundary!(;rng,
                                           X::Matrix{Float32},
                                           rand_mat::Matrix{Float32},
                                           upper::Vector{Float32},
                                           lower::Vector{Float32} )
    @assert length( upper ) == length( lower )
    @assert all( size( X )  == size( rand_mat ) )

    nFireworks = size( X, 2 );
    d          = length( lower );
    rand_mat  .= Float32.( rand( rng, d, nFireworks ) );
    for iter ∈ 1:1:nFireworks
      x_sp     = view( X,        :, iter );
      rand_vec = view( rand_mat, :, iter );
      for k ∈ 1:1:d
        lu      = lower[k]::Float32;
        x_sp[k] = ( rand_vec[k] * ( upper[k] - lu ) + lu )::Float32;
      end
    end
  end

  function explosionAmplitudes!(;A::Vector{Float32},
                                 λ_0::Float32,
                                 ϵ_A::Float32,
                                 y_min::Float32,
                                 fitness_fireworks::Vector{Float32} )
    n      = length( fitness_fireworks );
    cacheA = Vector{Float32}( undef, n );
    for ii ∈ 1:1:n
      @inbounds fx         = fitness_fireworks[ii];
      @inbounds cacheA[ii] = ( fx - y_min )::Float32;
    end
    denomA = ( sum( cacheA ) + ϵ_A )::Float32;
    for ii ∈ 1:1:n
      @inbounds A[ii] = ( λ_0 * ( ( cacheA[ii] + ϵ_A ) / denomA ) )::Float32;
    end
  end

  function evaluateCandidates( ;kwargs::Tuple,
                               f::Function,
                               X,
                               fitness )
    @assert size( X, 2 ) == size( fitness, 1 )
    n = size( X, 2 )::Int;
    @assert n > 0
    for ii ∈ 1:1:n
      x            = view( X, :, ii );
      fitness[ii]  = f( x; kwargs=kwargs );
    end
  end

    """
        simpleFWA( nFireworks::Int,
                   nSparks::Int;
                   λ_0::Float32,
                   ϵ_A::Float32,
                   C_a::Float32,
                   C_r::Float32,
                   lower::Vector{Float32},
                   upper::Vector{Float32},
                   objFunction::Function,
                   XPrimary::Vector{ Matrix{Float32} },
                   yPrimary::Vector{ Matrix{Float32} },
                   maxiter::Int,
                   ϵ_conv::Float32=1f-6 )
        minimize objective function objFunction, the solution space is limited by lower and upper bound.
        The optimization algorithm utilized is an simplified version of dynFireWorksAlgorithm. The nFireworks
        parameter governs the number of fireworks being evaluated in parallel in each iteration.  nSparks is
        the number of sparks per firework, in remains constant foreach firework. ϵ_A is the smoothing parameter
        controlling the variance of amplitudes computed foreach fw. C_a ist the upscaling parameter for explosion
        amplitudes. C_r is the downscaling parameter for explosion amplitudes. XPrimary is the feature set of the
        primary algorithm to be tuned. yPrimary is the target set of the primary algorithm. maxiter is the maximum
        number iteraions.  ϵ_conv denotes the convergence parameter.
    """
  @views function simpleFWA( nFireworks::Int,
                             nSparks::Int,
                             kwargs...;
                             λ_0::Float32,
                             ϵ_A::Float32,
                             C_a::Float32,
                             C_r::Float32,
                             lower::Vector{Float32},
                             upper::Vector{Float32},
                             objFunction::Function,
                             maxiter::Int,
                             ϵ_conv::Float32=0.001f0 )

    @assert nFireworks > 0
    @assert nSparks    > 0
    @assert λ_0        > 0.0f0
    @assert ϵ_A        > 0.0f0
    @assert C_a        > 0.0f0
    @assert C_r        > 0.0f0
    @assert length( upper ) == length( lower )
    @assert all( lower .<  upper )
    @assert maxiter   > 0

    d        = length( lower );
    rng      = RandomDevice();

    ## explosion amplitudes of firework roots
    fitness_fireworks = Vector{Float32}( undef, nFireworks );
    fitness_sparks    = Vector{ Vector{Float32} }( undef,  nFireworks );
    A                 = Vector{Float32}( undef, nFireworks );
    A_old             = zeros( Float32, nFireworks );
    ## spark positions for all fireworks
    S                 = Vector{ Matrix{Float32} }( undef, nFireworks );

    for ii in 1:1:nFireworks
      fitness_sparks[ii] = zeros( Float32, nSparks );
      S[ii]              = zeros( Float32, d, nSparks );
    end

    ## firework root explosion positions
    X                 = Matrix{Float32}( undef, d, nFireworks );
    ## random mat needed for inititialization of fw positions
    rand_mat          = Matrix{Float32}( undef, d, nFireworks );
    # generate nFireworks randomly within boundaries
    drawFireworksPositionsBoundary!( ;rng=rng, X=X, rand_mat=rand_mat, upper=upper, lower=lower );
    ## evaluate fireworks
    evaluateCandidates(;kwargs   = kwargs,
                        f        = objFunction,
                        X        = X,
                        fitness  = fitness_fireworks );
    cf        = argmin( fitness_fireworks )::Int;
    y_min     = fitness_fireworks[cf]::Float32;
    y_min_old = fitness_fireworks[Not( cf )][1]::Float32;
    conv_val  = -1.0f0;

    # compute explosion amplitudes
    explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
    A[cf]   = ( minimum( A[ Not( cf ) ] ) / 2.0f0 )::Float32;
    iter    = 1;
    err     = 2.0f0 * ϵ_conv;
    x_b     = X[:,cf];
    x_b_old = zeros( Float32, d );
    while iter <= maxiter
      cf              = argmin( fitness_fireworks )::Int;
      @inbounds y_cf  = fitness_fireworks[cf]::Float32;
      @inbounds A_cf  = A[cf]::Float32;
      # compute explosion amplitudes
      explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
      @inbounds A[cf]  = A_cf::Float32;
      @inbounds A[cf]  = dynExplosionAmplitude( y_min, y_cf, A_old[cf], C_a, C_r )::Float32;
      A_old .= A;
      # compute explosion sparks
      Threads.@threads for jj ∈ 1:1:nFireworks
        explosionSparks(;X_i             = X[:,jj],
                         A_i             = A[jj],
                         s               = S[jj],
                         lower           = lower,
                         upper           = upper,
                         rng             = rng );

        evaluateCandidates(;kwargs   = kwargs,
                            f        = objFunction,
                            X        = S[jj],
                            fitness  = fitness_sparks[jj] );
      end
      for jj ∈ 1:1:nFireworks
        @inbounds val = minimum( fitness_sparks[jj] );
        @inbounds idx = argmin( fitness_sparks[jj] );
        ## update best globally found solution
        if val < y_min
          y_min = val;
          @inbounds x_b   = S[jj][:,idx];
        end
        ## update fw move to best found local solution
        if val < fitness_fireworks[jj]
          @inbounds X[:,jj]              .= S[jj][:,idx];
          @inbounds fitness_fireworks[jj] = val;
        end
      end
      μFW      = mean( fitness_fireworks );
      conv_val = std( fitness_fireworks;mean=mean( fitness_fireworks ) )^0.10;
      if iter > 10 &&  conv_val < ϵ_conv
        break
      end
      y_min_old = y_min
      iter     += 1;
    end
    # check if best global solution is contained in fw position
    if y_min > minimum( fitness_fireworks )
      idx             = argmin( fitness_fireworks );
      @inbounds y_min = fitness_fireworks[ idx ];
      @inbounds x_b   = X[:,idx];
    end
    return FWA( X, fitness_fireworks, S, fitness_sparks, x_b, y_min, iter-1, conv_val )
  end

  @views function metropolisHastingsFWA( nFireworks::Int,
                                         nSparks::Int;
                                         λ_0::Float32,
                                         ϵ_A::Float32,
                                         C_a::Float32,
                                         C_r::Float32,
                                         lower::Vector{Float32},
                                         upper::Vector{Float32},
                                         objFunction::Function,
                                         XPrimary::Vector{ Matrix{Float32} },
                                         yPrimary::Vector{ Matrix{Float32} },
                                         maxiter::Int,
                                         ϵ_conv::Float32=1f-40 )

    @assert nFireworks > 0
    @assert nSparks    > 0
    @assert λ_0        > 0.0f0
    @assert ϵ_A        > 0.0f0
    @assert C_a        > 0.0f0
    @assert C_r        > 0.0f0
    @assert length( upper ) == length( lower )
    @assert all( lower .<  upper )
    @assert size( XPrimary, 2 ) == length( yPrimary )
    @assert maxiter   > 0

    d        = length( lower );
    rng      = RandomDevice();
    ## explosion amplitudes of firework roots
    fitness_fireworks = Vector{Float32}( undef, nFireworks );
    fitness_sparks    = Vector{ Vector{Float32} }( undef,  nFireworks );
    A                 = Vector{Float32}( undef, nFireworks );
    A_old             = zeros( Float32, nFireworks );
    ## spark positions for all fireworks
    S                 = Vector{ Matrix{Float32} }( undef, nFireworks );

    for ii in 1:1:nFireworks
      fitness_sparks[ii] = zeros( Float32, nSparks );
      S[ii]              = zeros( Float32, d, nSparks );
    end

    ## firework root explosion positions
    X                 = Matrix{Float32}( undef, d, nFireworks );
    ## random mat needed for inititialization of fw positions
    rand_mat          = Matrix{Float32}( undef, d, nFireworks );
    # generate nFireworks randomly within boundaries
    drawFireworksPositionsBoundary!( ;rng=rng, X=X, rand_mat=rand_mat, upper=upper, lower=lower );
    ## evaluate fireworks
    evaluateCandidates(;XPrimary = XPrimary,
                        yPrimary = yPrimary,
                        f        = objFunction,
                        X        = X,
                        fitness  = fitness_fireworks );
    cf    = argmin( fitness_fireworks )::Int;
    y_min = fitness_fireworks[cf]::Float32;

    # compute explosion amplitudes
    explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
    A[cf]   = ( minimum( A[ Not( cf ) ] ) / 2.0f0 )::Float32;
    iter    = 0;
    err     = 2.0f0 * ϵ_conv;
    x_b     = X[:,cf];
    x_b_old = zeros( Float32, d );
    while err > ϵ_conv && iter <= maxiter
      cf              = argmin( fitness_fireworks )::Int;
      @inbounds y_cf  = fitness_fireworks[cf]::Float32;
      @inbounds A_cf  = A[cf]::Float32;
      # compute explosion amplitudes
      explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
      @inbounds A[cf]  = A_cf::Float32;
      @inbounds A[cf]  = dynExplosionAmplitude( y_min, y_cf, A_old[cf], C_a, C_r )::Float32;
      A_old .= A;
      # compute explosion sparks
      Threads.@threads for jj ∈ 1:1:nFireworks
        explosionSparks(;X_i             = X[:,jj],
                         A_i             = A[jj],
                         s               = S[jj],
                         lower           = lower,
                         upper           = upper,
                         rng             = rng );

        evaluateCandidates(;XPrimary = XPrimary,
                            yPrimary = yPrimary,
                            f        = objFunction,
                            X        = S[jj],
                            fitness  = fitness_sparks[jj] );
      end
      for jj ∈ 1:1:nFireworks
        @inbounds val = minimum( fitness_sparks[jj] );
        @inbounds idx = argmin( fitness_sparks[jj] );
        ## update best globally found solution
        if val < y_min
          y_min = val;
          @inbounds x_b   = S[jj][:,idx];
        end
        ## update fw by metropolis hastings rejection sampling
        α = val / fitness_fireworks[jj];
        if rand(Float32) < α
          @inbounds X[:,jj]              .= S[jj][:,idx];
          @inbounds fitness_fireworks[jj] = val;
        end
      end
      μFW = mean( fitness_fireworks );
      #err = sqrt( sum( std( X, mean=mean( X, dims=2 )[:,1], dims=2 )[:,1] ) );
      iter += 1;
    end
    # check if best global solution is contained in fw position
    if y_min > minimum( fitness_fireworks )
      idx             = argmin( fitness_fireworks );
      @inbounds y_min = fitness_fireworks[ idx ];
      @inbounds x_b   = X[:,idx];
    end
    return FWA( X, fitness_fireworks, S, fitness_sparks, x_b, y_min, iter, 0.0f0 )
  end

end


text:
"""
    OptSummary

Summary of an `NLopt` optimization

# Fields
* `initial`: a copy of the initial parameter values in the optimization
* `finitial`: the initial value of the objective
* `lowerbd`: lower bounds on the parameter values
* `ftol_rel`: as in NLopt
* `ftol_abs`: as in NLopt
* `xtol_rel`: as in NLopt
* `xtol_abs`: as in NLopt
* `initial_step`: as in NLopt
* `maxfeval`: as in NLopt (`maxeval`)
* `maxtime`: as in NLopt
* `final`: a copy of the final parameter values from the optimization
* `fmin`: the final value of the objective
* `feval`: the number of function evaluations
* `optimizer`: the name of the optimizer used, as a `Symbol`
* `returnvalue`: the return value, as a `Symbol`
* `nAGQ`: number of adaptive Gauss-Hermite quadrature points in deviance evaluation for GLMMs
* `REML`: use the REML criterion for LMM fits
* `sigma`: a priori value for the residual standard deviation for LMM
* `fitlog`: A vector of tuples of parameter and objectives values from steps in the optimization

The latter four fields are MixedModels functionality and not related directly to the `NLopt` package or algorithms.

!!! note
    The internal storage of the parameter values within `fitlog` may change in
    the future to use a different subtype of `AbstractVector` (e.g., `StaticArrays.SVector`)
    for each snapshot without being considered a breaking change.
"""
mutable struct OptSummary{T<:AbstractFloat}
    initial::Vector{T}
    lowerbd::Vector{T}
    finitial::T
    ftol_rel::T
    ftol_abs::T
    xtol_rel::T
    xtol_abs::Vector{T}
    initial_step::Vector{T}
    maxfeval::Int
    maxtime::T
    feval::Int
    final::Vector{T}
    fmin::T
    optimizer::Symbol
    returnvalue::Symbol
    nAGQ::Integer           # don't really belong here but I needed a place to store them
    REML::Bool
    sigma::Union{T,Nothing}
    fitlog::Vector{Tuple{Vector{T},T}} # not SVector because we would need to parameterize on size (which breaks GLMM)
end

function OptSummary(
    initial::Vector{T},
    lowerbd::Vector{T},
    optimizer::Symbol;
    ftol_rel::T=zero(T),
    ftol_abs::T=zero(T),
    xtol_rel::T=zero(T),
    xtol_abs::Vector{T}=zero(initial) .+ 1e-10,
    initial_step::Vector{T}=T[],
    maxfeval=-1,
    maxtime=T(-1),
) where {T<:AbstractFloat}
    fitlog = [(initial, T(Inf))]

    return OptSummary(
        initial,
        lowerbd,
        T(Inf),
        ftol_rel,
        ftol_abs,
        xtol_rel,
        xtol_abs,
        initial_step,
        maxfeval,
        maxtime,
        -1,
        copy(initial),
        T(Inf),
        optimizer,
        :FAILURE,
        1,
        false,
        nothing,
        fitlog,
    )
end

"""
    columntable(s::OptSummary, [stack::Bool=false])

Return `s.fitlog` as a `Tables.columntable`.

When `stack` is false (the default), there will be 3 columns in the result:
- `iter`: the sample number
- `objective`: the value of the objective at that sample
- `θ`: the parameter vector at that sample

(The term `sample` here refers to the fact that when the `thin` argument to the `fit` or
`refit!` call is greater than 1 only a subset of the iterations have results recorded.)

When `stack` is true, there will be 4 columns: `iter`, `objective`, `par`, and `value`
where `value` is the stacked contents of the `θ` vectors (the equivalent of `vcat(θ...)`)
and `par` is a vector of parameter numbers.
"""
function Tables.columntable(s::OptSummary; stack::Bool=false)
    fitlog = s.fitlog
    val = (; iter=axes(fitlog, 1), objective=last.(fitlog), θ=first.(fitlog))
    stack || return val
    θ1 = first(val.θ)
    k = length(θ1)
    return (;
        iter=repeat(val.iter; inner=k),
        objective=repeat(val.objective; inner=k),
        par=repeat(1:k; outer=length(fitlog)),
        value=foldl(vcat, val.θ; init=(eltype(θ1))[]),
    )
end

function Base.show(io::IO, ::MIME"text/plain", s::OptSummary)
    println(io, "Initial parameter vector: ", s.initial)
    println(io, "Initial objective value:  ", s.finitial)
    println(io)
    println(io, "Optimizer (from NLopt):   ", s.optimizer)
    println(io, "Lower bounds:             ", s.lowerbd)
    println(io, "ftol_rel:                 ", s.ftol_rel)
    println(io, "ftol_abs:                 ", s.ftol_abs)
    println(io, "xtol_rel:                 ", s.xtol_rel)
    println(io, "xtol_abs:                 ", s.xtol_abs)
    println(io, "initial_step:             ", s.initial_step)
    println(io, "maxfeval:                 ", s.maxfeval)
    println(io, "maxtime:                  ", s.maxtime)
    println(io)
    println(io, "Function evaluations:     ", s.feval)
    println(io, "Final parameter vector:   ", s.final)
    println(io, "Final objective value:    ", s.fmin)
    return println(io, "Return code:              ", s.returnvalue)
end

Base.show(io::IO, s::OptSummary) = Base.show(io, MIME"text/plain"(), s)

function NLopt.Opt(optsum::OptSummary)
    lb = optsum.lowerbd

    opt = NLopt.Opt(optsum.optimizer, length(lb))
    NLopt.ftol_rel!(opt, optsum.ftol_rel) # relative criterion on objective
    NLopt.ftol_abs!(opt, optsum.ftol_abs) # absolute criterion on objective
    NLopt.xtol_rel!(opt, optsum.xtol_rel) # relative criterion on parameter values
    if length(optsum.xtol_abs) == length(lb)  # not true for fast=false optimization in GLMM
        NLopt.xtol_abs!(opt, optsum.xtol_abs) # absolute criterion on parameter values
    end
    NLopt.lower_bounds!(opt, lb)
    NLopt.maxeval!(opt, optsum.maxfeval)
    NLopt.maxtime!(opt, optsum.maxtime)
    if isempty(optsum.initial_step)
        optsum.initial_step = NLopt.initial_step(opt, optsum.initial, similar(lb))
    else
        NLopt.initial_step!(opt, optsum.initial_step)
    end
    return opt
end

StructTypes.StructType(::Type{<:OptSummary}) = StructTypes.Mutable()
StructTypes.excludes(::Type{<:OptSummary}) = (:lowerbd,)

const _NLOPT_FAILURE_MODES = [
    :FAILURE,
    :INVALID_ARGS,
    :OUT_OF_MEMORY,
    :FORCED_STOP,
    :MAXEVAL_REACHED,
    :MAXTIME_REACHED,
]

function _check_nlopt_return(ret)
    ret == :ROUNDOFF_LIMITED && @warn("NLopt was roundoff limited")
    if ret ∈ _NLOPT_FAILURE_MODES
        @warn("NLopt optimization failure: $ret")
    end
end


text:
#! format: off

using SIIPExamples
using PowerSystems
using JSON3

pkgpath = pkgdir(SIIPExamples)
include(joinpath(pkgpath, "test", "2_PowerSystems_examples", "02_parse_matpower.jl"))

FORECASTS_DIR = joinpath(base_dir, "forecasts", "5bus_ts")
fname = joinpath(FORECASTS_DIR, "timeseries_pointers_da.json")
open(fname, "r") do f
    JSON3.@pretty JSON3.read(f)
end

add_time_series!(sys, fname)
sys

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl

max_stars_repo_name: NREL-SIIP/Examples

max_issues_repo_name: NREL-SIIP/Examples

max_forks_repo_name: NREL-SIIP/Examples


text:
export relocate, relocate!

relocate(df, args...; kwargs...) = relocate!(copy(df), args...; kwargs...)

function relocate!(df, w::Where; kwargs...)
    tmp = [w.fn(col) for col in eachcol(df)]

    cols = names(df)[tmp]

    relocate!(df, cols...; kwargs...)
end

function relocate!(df, f::Function; kwargs...)
    cols = f(df)

    relocate!(df, cols...; kwargs...)
end

function relocate!(df, cols...; before = nothing, after = nothing)
    if !isnothing(before) && !isnothing(after)
        error("before and after cannot both be populated")
    elseif isnothing(before) && isnothing(after)
        select!(df, cols..., :)
    elseif !isnothing(before)
        if before isa Function
            before_idx = columnindex(df, before(df))
        elseif before isa Where
            tmp = [before.fn(col) for col in eachcol(df)]
            where_cols = names(df)[tmp]
            before_idx = columnindex.(Ref(df), where_cols) |> minimum
        elseif before isa Symbol || before isa AbstractString
            before_idx = columnindex(df, before)
        end

        if cols isa Symbol || cols isa AbstractString
            cols_idx = columnindex(df, cols)
        else
            cols_names = reduce(vcat, names.(Ref(df), cols))
            cols_idx = columnindex.(Ref(df), cols_names)
        end
        select!(df, setdiff(1:before_idx-1, cols_idx), cols_idx, before_idx, :)
    elseif !isnothing(after)
        if after isa Function
            after_idx = columnindex(df, after(df))
        elseif after isa Where
            tmp = [after.fn(col) for col in eachcol(df)]
            where_cols = names(df)[tmp]
            after_idx = columnindex.(Ref(df), where_cols) |> maximum
        elseif after isa Symbol || after isa AbstractString
            after_idx = columnindex(df, after)
        end

        if cols isa Symbol || cols isa AbstractString
            cols_idx = columnindex(df, cols)
        else
            cols_names = reduce(vcat, names.(Ref(df), cols))
            cols_idx = columnindex.(Ref(df), cols_names)
        end

        select!(df, setdiff(1:after_idx, cols_idx), cols_idx, :)
    end
end


text:
## Module associated with the MeshArrays package

module MeshArrays

using Pkg, Pkg.Artifacts
thistoml=joinpath(dirname(pathof(MeshArrays)), "..", "Project.toml")
thisversion=Pkg.TOML.parsefile(thistoml)["version"]

p=dirname(pathof(MeshArrays))
artifact_toml = joinpath(p, "../Artifacts.toml")

GRID_LLC90_hash = artifact_hash("GRID_LLC90", artifact_toml)
GRID_LLC90 = joinpath(artifact_path(GRID_LLC90_hash)*"/","GRID_LLC90-1.1/")
GRID_LLC90_download() = artifact"GRID_LLC90"
GRID_LL360_hash = artifact_hash("GRID_LL360", artifact_toml)
GRID_LL360 = joinpath(artifact_path(GRID_LL360_hash)*"/","GRID_LL360-1.0/")
GRID_LL360_download() = artifact"GRID_LL360"
GRID_CS32_hash = artifact_hash("GRID_CS32", artifact_toml)
GRID_CS32 = joinpath(artifact_path(GRID_CS32_hash)*"/","GRID_CS32-1.1/")
GRID_CS32_download() = artifact"GRID_CS32"

include("Types.jl")
include("Grids.jl")
include("Operations.jl")
include("Exchanges.jl")
include("ReadWrite.jl")
include("Solvers.jl")
include("ReIndexing.jl")
include("Interpolation.jl")

export AbstractMeshArray, MeshArray, InnerArray, OuterArray, varmeta
export gcmgrid, exchange, gradient, convergence, smooth, mask
export UnitGrid, simple_periodic_domain, GridSpec, GridLoad, GridOfOnes, GridAddWS!
export Tiles, Interpolate, InterpolationFactors, knn
export ScalarPotential, VectorPotential, ThroughFlow
export StereographicProjection, LatitudeCircles

export location_is_out, NeighborTileIndices_dpdo, NeighborTileIndices_cs, RelocationFunctions_cs
export update_location_cs!, update_location_llc!, update_location_dpdo!

#The following exch_UV differs from normal exchange; incl. exch_UV_N
export exch_UV

end # module


text:
using FileIO: @format_str, load, File
using RData: RString
using DataFrames: DataFrame

using Query: @filter
using Suppressor: @suppress

using OrderedCollections: OrderedDict

import YAML

packdir = "$(@__DIR__)/../data/"

# This way doesn't emit warnings, but requires some change to output format.
#= linkscens = load(File(format"RDataSingle", joinpath(packdir, "processed", "link_scenarios.Rds"))) =#
#= odscens = load(File(format"RDataSingle", joinpath(packdir, "processed", "od_matrices_scenarios.Rds"))) =#

function load_scenarios(packdir=packdir)
    # This way emits warnings, but works fine.
    scenarios = @suppress load(File(format"RDataSingle", joinpath(packdir, "processed", "julia_compat_scenarios.Rds")));

    df = DataFrame(values(scenarios), [:scenario, :year, :type, :data]);
    if !(typeof(df.year) <: Vector{T} where T <: Integer)
        df.year = parse.(Int, df.year)
    end

    links = df |>
        @filter(_.type == "links") |>
        DataFrame |>
        (df -> Dict((name, year) => data for (name, year, _, data) in zip(eachcol(df)...)));

    mats = df |>
        @filter(_.type == "od_matrices") |>
        DataFrame |>
        (df -> Dict((name, year) => data for (name, year, _, data) in zip(eachcol(df)...)))

    metadata = get_metadata(links, mats, packdir)

    return links, mats, metadata
end

function get_metadata(links, mats, packdir)
    # This should probably be reloaded periodically so server doesn't need to be restarted?
    metadata = YAML.load_file(joinpath(packdir, "meta.yaml"); dicttype=OrderedDict{String, Any})

    default_meta = Dict(
        "good" => "smaller",
        "thickness" => "variable",
        "statistics" => "hide",
        "force_bounds" => [],
    )

    for (k,v) in metadata["links"]["columns"]
        metadata["links"]["columns"][k] = merge(default_meta, v)
    end

    for (k,v) in metadata["od_matrices"]["columns"]
        metadata["od_matrices"]["columns"][k] = merge(default_meta, v)
    end

    for (k,v) in metadata["scenarios"]
        metadata["scenarios"][k]["name"] = get(metadata["scenarios"][k],"name",k)
    end

    # Store what years are available for each scenario
    d = Dict(links |> keys .|> x -> (x[1], Int[]))
    for (k, v) in links |> keys
        push!(d[k], v)
    end

    d2 = Dict(mats |> keys .|> x -> (x[1], Int[]))
    for (k, v) in mats |> keys
        push!(d2[k], v)
    end

    scenarioYears = Dict(
        k => sort(intersect(d[k], d2[k]))
        for k in keys(d)
    )

    for (name, years) in scenarioYears
        try
            metadata["scenarios"][name]["at"] = years
        catch x
            if isa(x, KeyError)
                metadata["scenarios"][name] = Dict("use" => false)
            else
                rethrow()
            end
        end
    end

    return metadata
end

#=

using Statistics

# Aggregate to one df, takes a little while and allocates a lot.
# But it's only once, so whatever, could do it.
@time vcat(values(links)...);

# Could just apply the function across an iterable of each thing?
@time Iterators.flatten(map(df -> @view(df[!, :Cost_Pax]), values(links))) |> mean;

# Same kind of thing for the matrices
Iterators.flatten(map(dict -> dict["Pax"], values(mats))) |> mean

Dict((name, year) => matrices for (name, year, matrices) in zip(values(odscens)...))

=#

# Interface

# Get a vector or matrix of a variable for a given year and scenario
# Get a vector of a column for all scenarios and years (maybe not used any more?)

max_stars_repo_name: PolyCement/gcvt

max_issues_repo_name: PolyCement/gcvt

max_forks_repo_name: PolyCement/gcvt


text: #! format: off

max_stars_repo_name: AdelekeBankole/libCEED

max_issues_repo_name: AdelekeBankole/libCEED

max_forks_repo_name: AdelekeBankole/libCEED


text:
import DESPOT:
        lower_bound

mutable struct DESPOTDefaultLowerBound
    #placeholder for now
end

function init_bound{S,A,O}(ub::DESPOTDefaultLowerBound,
                    pomdp::POMDP{S,A,O},
                    config::DESPOTConfig)
    error("Function init_bound for $(typeof(ub)) has not been implemented yet")
end

function lower_bound{S,A,O}(lb::DESPOTDefaultLowerBound,
                     pomdp::POMDP{S,A,O},
                     particles::Vector{DESPOTParticle{S}},
                     config::DESPOTConfig)
    error("Function lower_bound for $(typeof(lb)) has not been implemented yet")
end


text:
function evolve_neutral!( pop::Population, generations::Int64 = 1 )
    # function to evolve the population under neutral selection
    # that is, role models are selected uniformly at random; fitnesses have no impact
    # generations specifies the total number of generations for which to run the simulation
    for i::Int64 in 1:generations
        if pop.verbose println("\ninitiating generation $(pop.generation)") end

        # first, update payoffs
        if pop.verbose println("updating actions only") end
        update_actions_neutral!(pop)  # method that updates actions only

        # then, select a learner and a role model and perform update
        if pop.verbose println("") end
        if pop.verbose println("evolving, generation $(pop.generation)") end
        update_strategies_and_opinions_neutral!(pop)

        if pop.verbose println("") end
        if pop.verbose println("ending generation $(pop.generation)") end
        pop.generation += 1
    end
end

function update_actions_neutral!( pop::Population )
    # update individual actions
    # for neutral selection (β = 0), so payoffs are not needed

    if pop.sets_changed
        # if (previous) learner's sets have changed, re-compute all actions
        if pop.verbose println("updating all actions because set_changed = $(pop.sets_changed)") end
        pop.prev_actions .= zero(Int64)   # reset all prev_actions

        # then, recompute all payoffs
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]   # for each pair within each set
                compute_actions_neutral!( pop, i, j, k )  # update action between i & j when intereacting in set k
            end
        end
    elseif pop.strategies_changed
        # if (previous) learner's strategy has changed, recompute actions involving that individual
        if pop.verbose println("updating actions involving $(pop.prev_learner) b/c strategies_changed = $(pop.strategies_changed)") end
        pop.prev_actions[:,pop.prev_learner,:] .= zero(Int64)
        pop.prev_actions[:,:,pop.prev_learner] .= zero(Int64)

        # then, recompute payoffs involving prev_learner
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]  # for each pair within each set
                if i == pop.prev_learner || j == pop.prev_learner
                    compute_actions_neutral!( pop, i, j, k )  # update action between i & j when intereacting in set k
                end
            end
        end
    else
        if pop.verbose println("not updating actions b/c sets_changed = $(pop.sets_changed) and strategies_changed = $(pop.strategies_changed)") end
        return # no need to update prev_interactions or payoffs
    end
    # sum rows to get total individual payoffs
    pop.prev_interactions = 2 * sum([length(x) for x in pop.sets.set_pairs])

    if pop.verbose println( "\npayoff matrix is now $(pop.payoffs_mat)" ) end
    if pop.verbose println( "payoffs are now $(pop.payoffs)" ) end
    if pop.verbose println( "there were $(pop.prev_interactions) interactions, $(sum(pop.prev_actions)) of which were C" ) end
end

function compute_actions_neutral!( pop::Population, i::Int64, j::Int64, k::Int64 )
    # function to compute payoffs for individuals i and j when they interact in set k
    pop.sim.i_action, pop.sim.j_action = -1, -1  # set to -1 to catch any errors in action assignment

    # actions depend on whether i and j agree on issue k or not
    if pop.sets.h[i,k] == 0 || pop.sets.h[j,k] == 0
        # if either i or j does not belong to set k, something's wrong
        # leave pop.sim.i_action, pop.sim.j_action as -1 so that the compute_payoffs! throws an error when assigning payoffs
        @warn("$i and $j should not be interacting in set $k"); return
    elseif pop.sets.h[i,k] == pop.sets.h[j,k]
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,1], pop.strategies[j,1] # if they agree, play s_agree
    else
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,2], pop.strategies[j,2] # if they disagree, play s_disagree
    end

    if pop.verbose
        println("\n\tupdating actions of $i and $j in set $k; agree on issue $k? $(pop.sets.h[i,k] == pop.sets.h[j,k])")
        println("\t\t$i has strats $(pop.strategies[i,:]); $i plays strat $(pop.sim.i_action)")
        println("\t\t$j has strats $(pop.strategies[j,:]); $j plays strat $(pop.sim.j_action)")
    end
    # track actions; no need to track payoffs under neutral selection
    pop.prev_actions[k,i,j] = pop.sim.i_action
    pop.prev_actions[k,j,i] = pop.sim.j_action
end

function update_strategies_and_opinions_neutral!( pop::Population )
    # function to perform death-birth update
    if pop.verbose println("Strategies before updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships before updating: $(pop.sets.h)") end

    # choose a learner uniformly at random
    pop.sim.learner = sample( 1:pop.sets.N )
    # choose a role model also uniformly at random (for neutral selection)
    pop.sim.role_model = sample( 1:pop.sets.N )

    # learner updates
    if pop.verbose println("\n\trandomly chosen learner is $(pop.sim.learner), role_model is $(pop.sim.role_model)") end
    if pop.verbose println("\t\tall payoffs are $(pop.payoffs)") end
    if pop.verbose println("\t\tfitnesses are $(pop.sim.fitnesses)") end

    # compute probability of imitation: 1 if same party, 1-p if different parties
    pop.sets.affils[pop.sim.learner] == pop.sets.affils[pop.sim.role_model] ? imit_prob = 1.0 : imit_prob = 1-pop.game.ps[pop.sets.affils[pop.sim.learner]]

    if pop.verbose
        println("\t\tlearner and role model in the same party? $(pop.sets.affils[ pop.sim.learner ] == pop.sets.affils[ pop.sim.role_model ]), imitate with prob $(imit_prob)")
    end

    # update strategies and sets: imitation / mutation occurs with probability imit_prob
    if rand() < imit_prob
        if pop.verbose println("\twill consider imitating") end
        update_strategies!( pop ) # includes mutation
        update_opinions!( pop )   # includes mutation
        pop.num_imitating += 1
    else
        if pop.verbose println("\tnah won't imitate") end
        no_strategy_or_opinion_update!( pop )
        pop.num_not_imitating += 1
    end

    # track updates
    pop.prev_learner = pop.sim.learner
    if pop.sets_changed == true
        # h is updated above; need to update h_bar, h_base10, h_bar_base10, set_members, set_pairs
        pop.sets.set_members, pop.sets.set_pairs = set_members_and_pairs(pop.sets.h)  # update set_members and set_pairs
        pop.sets.h_bar        = [abs(x) for x in pop.sets.h]
        pop.sets.h_base10     = [vectobase3(pop.sets.h[i,:]) for i in 1:pop.sets.N]
        pop.sets.h_bar_base10 = [vectobase3(pop.sets.h_bar[i,:]) for i in 1:pop.sets.N]
        pop.num_sets_changed += 1
    end
    if pop.strategies_changed == true
        # strategies are updated above; need to update strategies_base10
        pop.strategies_base10 = [2 * pop.strategies[i,1] + pop.strategies[i,2] for i in 1:pop.sets.N]
        pop.num_strategies_changed += 1
    end
    if pop.strategies_changed && pop.sets_changed
        pop.num_both_changed += 1
    end

    if pop.verbose println("\nStrategies after updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships after updating:\t $(pop.sets.h)") end
    if pop.verbose println("tracking: prev_learner = $(pop.prev_learner), strategies_changed = $(pop.strategies_changed), sets_changed = $(pop.sets_changed)") end
end