text:
##################
## Broadcasting ##
##################

using Base.Broadcast: BroadcastStyle, ArrayStyle, Broadcasted, broadcasted
using ForwardDiff: ForwardDiff, Dual
import Base.Broadcast: materialize
const RDBroadcasted{F, T} = Broadcasted{<:Any, <:Any, F, T}

"""
    NotTracked(f::Function)

A struct that can be used to wrap around closures, structs and arrays of structs declaring that they do not contain tracked variables. This enables a more efficient broadcasting of such functions and structs when doing automatic differentiation with `ReverseDiff` producing a `TrackedArray` instead of an `Array{<:TrackedReal}`.
"""
struct NotTracked{F} <: Function
    f::F
end
(f::NotTracked{<:Union{Function, Type}})(args...; kwargs...) = f.f(args...; kwargs...)

istypeorclosure(::F) where {F} = _istypeorclosure(F)
istypeorclosure(::AbstractArray{F}) where {F} = _istypeorclosure(F)
istypeorclosure(::Base.RefValue{F}) where {F} = _istypeorclosure(F)
istypeorclosure(::AbstractArray{<:Real}) = false
istypeorclosure(::TrackedArray) = false
istypeorclosure(::AbstractArray{<:TrackedReal}) = true
istypeorclosure(::Real) = false
@generated _istypeorclosure(::Type{F}) where {F} = :($(fieldcount(F) > 0))

mayhavetracked(b) = istypeorclosure(b)
mayhavetracked(b::Type) = false
mayhavetracked(b::NotTracked) = false
mayhavetracked(b::Base.RefValue{<:NotTracked}) = false
mayhavetracked(b::Broadcasted) = mayhavetracked(b.f) || any(mayhavetracked, b.args)

struct TrackedStyle <: BroadcastStyle end

Broadcast.BroadcastStyle(::Type{<:Union{TrackedArray, TrackedReal}}) = TrackedStyle()
Broadcast.BroadcastStyle(::TrackedStyle, b::BroadcastStyle) = TrackedStyle()

# We have to re-build the original broadcast struct to get the appropriate array
# style. We need this primarily to support CuArrays' broadcasting fixes.
broadcast_rebuild(xs) = recur_value(xs)
recur_value(xs) = xs
recur_value(xs::Union{TrackedReal, TrackedArray}) = recur_value(value(xs))

function broadcast_rebuild(bc::Broadcasted)
    broadcasted(bc.f, broadcast_rebuild.(bc.args)...)
end

getstyle(::Broadcasted{Style}) where {Style} = Style
remove_not_tracked(f) = f
remove_not_tracked(f::NotTracked) = f.f
remove_not_tracked(f::Base.RefValue{<:NotTracked}) = Ref(remove_not_tracked(f[]))
remove_not_tracked(f::Base.RefValue{<:NotTracked{<:AbstractArray}}) = remove_not_tracked(f[])
function remove_not_tracked(b::Broadcasted{style}) where {style}
    return Broadcasted{style}(remove_not_tracked(b.f), remove_not_tracked.(b.args), b.axes)
end

onlyrealarrays(args::Tuple) = onlyrealarray(first(args)) && onlyrealarrays(Base.tail(args))
onlyrealarrays(::Tuple{}) = true
onlyrealarray(::AbstractArray{<:Real}) = true
onlyrealarray(::AbstractArray) = false
onlyrealarray(::Any) = true

anyreals(args::Tuple) = first(args) isa Real || anyreals(Base.tail(args))
anyreals(args::Tuple{}) = false

function get_implementation(bc, f, T, args)
    outputisreal = (T <: AbstractArray{<:Real}) && (T !== Union{})
    # Each arg is either a real number, an array of untraked reals, a tracked array of reals or an array of untracked non-reals,
    # Output is real, and
    # No tracked closure or arguments, except TrackedReal and TrackedArray.
    if !mayhavetracked(bc) && outputisreal && (anyreals(args) || !onlyrealarrays(args))
        return Val(:tracker)
    # No arg is a real number and array args must be arrays of untracked reals or tracked arrays of reals,
    # Output is real, and
    # No tracked closure or arguments, except TrackedReal and TrackedArray.
    elseif !mayhavetracked(bc) && outputisreal
        return Val(:reversediff)
    # Function or any arg is possibly a tracked non-real or an array of tracked reals/non-reals,
    # Or output is not an array of reals
    else
        return Val(:fallback)
    end
end
function Base.copy(_bc::Broadcasted{TrackedStyle})
    bc = remove_not_tracked(_bc)
    flattened_bc = Broadcast.flatten(bc)
    untracked_bc = broadcast_rebuild(bc)
    flattened_untracked_bc = Broadcast.flatten(untracked_bc)
    T = Core.Compiler.return_type(copy, Tuple{typeof(untracked_bc)})
    f, args = flattened_untracked_bc.f, flattened_bc.args
    implementation = get_implementation(_bc, f, T, args)
    if implementation isa Val{:reversediff}
        return ∇broadcast(f, args...)
    elseif implementation isa Val{:tracker}
        return tracker_∇broadcast(f, args...)
    else
        style, axes = getstyle(flattened_untracked_bc), flattened_bc.axes
        return copy(Broadcasted{style, typeof(axes), typeof(f), typeof(args)}(f, args, axes))
    end
end

getouttype(::TrackedReal{<:Any, D}) where {D} = D
getouttype(::TrackedArray{<:Any, D}) where {D} = D
getouttype(::Any) = Union{}

deref(x) = x
deref(x::Base.RefValue) = x[]

@generated function splatcall(f, x::SVector{N}, utargs::T, ::Val{tinds}) where {N, T <: Tuple, tinds}
    args = []
    ti = 1
    uti = 1
    for i in 1:(N + length(T.types))
        if i in tinds
            push!(args, :(deref(x[$ti])))
            ti += 1
        else
            push!(args, :(deref(utargs[$uti])))
            uti += 1
        end
    end
    return quote
        $(Expr(:meta, :inline))
        $(Expr(:call, :f, args...))
    end
end

@generated function splitargs(args::T) where {T <: Tuple}
    N = length(T.types)
    RealOrArray = Union{Real, AbstractArray}
    inds = [i for i in 1:N if T.types[i] <: RealOrArray]
    indsval = :(Val{$(Expr(:tuple, [:($i) for i in inds]...))}())
    maybetracked = Expr(:tuple, [:(args[$i]) for i in inds]...)
    untracked = Expr(:tuple, [:(args[$i]) for i in 1:N if !(i in inds)]...)
    return :($indsval, $maybetracked, $untracked)
end

## A generalization of the broadcasting approach in ReverseDiff for general functions

@inline function ∇broadcast(f::F, args::Vararg{<:Any}) where {F}
    inds, targs, untracked = splitargs(args)
    N = length(targs)
    D = promote_type(getouttype.(targs)...)
    result = DiffResults.GradientResult(zero(SVector{N, D}))
    function df(x...)
        return ForwardDiff.gradient!(
            result,
            s -> splatcall(f, s, untracked, inds),
            SVector(x),
        )
    end
    results = broadcast(df, value.(targs)...)
    tp = tape(targs...)
    out_value = DiffResults.value.(results)
    eltype(out_value) == Bool && return out_value
    out = track(out_value, D, tp)
	cache = (results, df, index_bound.(targs, (out,)))
	record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache)
    return out
end
@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(∇broadcast)})
    input = instruction.input
    output = instruction.output
    output_deriv = deriv(output)
    results, _, bounds = instruction.cache
    N = length(input)
    if N == 1 || all(isequal(size(input[1])), size.(Base.tail(input)))
        _br_add_to_deriv!(input, output_deriv, results)
    else
        _br_add_to_deriv!(input, output_deriv, results, bounds)
    end
    unseed!(output)
    return nothing
end

@generated function _br_add_to_deriv!(xs::T, o, r) where {T <: Tuple}
    N = length(T.types)
    return Expr(:block, [:(_br_add_to_deriv!(xs[$i], o, r, Val($i))) for i in 1:N]...)
end
_br_add_to_deriv!(_, _, _, _) = nothing
function _br_add_to_deriv!(x::Union{TrackedReal, TrackedArray}, out_deriv, results, ::Val{i}) where {i}
    return istracked(x) && diffresult_increment_deriv!(x, out_deriv, results, i)
end

@generated function _br_add_to_deriv!(xs::T, o, r, bounds) where {T <: Tuple}
    N = length(T.types)
    return Expr(:block, [:(_br_add_to_deriv!(xs[$i], o, r, Val($i), bounds[$i])) for i in 1:N]...)
end
_br_add_to_deriv!(_, _, _, _, _) = nothing
function _br_add_to_deriv!(x::Union{TrackedReal,TrackedArray}, out_deriv, results, ::Val{i}, bound) where {i}
    return istracked(x) && diffresult_increment_deriv!(x, out_deriv, results, i, bound)
end

@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)})
    input, output = instruction.input, instruction.output
    results, df, _ = instruction.cache
    pull_value!.(input)
    broadcast!(df, results, value.(input)...)
    output_value = value(output)
    output_value .= DiffResults.value.(results)
    return nothing
end

## Tracker style broadcasting
## Good for broadcasting real numbers or arrays of non-tracked structs

trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))

unbroadcast(x::AbstractArray, Δ) =
  size(x) == size(Δ) ? Δ :
  length(x) == length(Δ) ? trim(x, Δ) :
    trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ)))))

unbroadcast(x::Number, Δ) = sum(Δ)
unbroadcast(x::Base.RefValue, _) = nothing

dual(x, p) = x
dual(x::Real, p) = Dual(x, p)

function _deriv(f, G, ::Val{i}, args::Vararg{Any, N}) where {N, i}
    dargs = ntuple(j -> dual(args[j], i==j), Val(N))
    return f(dargs...).partials[1] * G
end
@generated function _derivs(f, G, args::Vararg{Any, N}) where {N}
    return Expr(:tuple, [:(_deriv.(f, G, Val($i), args...)) for i in 1:N]...)
end
@inline function tracker_∇broadcast(f, args::Vararg{Any, N}) where {N}
    args_values = map(value, args)
    out_value = broadcast(f, args_values...)
    tp = tape(args...)
    eltype(out_value) == Bool && return out_value
	out = track(out_value, tp)
    cache = (f,)
	record!(tp, SpecialInstruction, tracker_∇broadcast, args, out, cache)
    return out
end

@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)})
    input, output = instruction.input, instruction.output
    f = instruction.cache[1]
    output_value = value(output)
    pull_value!.(input)
    broadcast!(f, output_value, value.(input)...)
    return nothing
end

@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)})
    input = instruction.input
    output = instruction.output
    f = instruction.cache[1]
    output_deriv = deriv(output)
    N = length(input)
    Δargs = _derivs(f, output_deriv, value.(input)...)
    dxs = map(unbroadcast, input, Δargs)
    map(_add_to_deriv!, input, dxs)
    unseed!(output)
    return nothing
end

## Limited ReverseDiff broadcasting
## Efficient broadcasting for specific functions, e.g. +, -

@inline _materialize(f, args) = broadcast(f, args...)

for (M, f, arity) in DiffRules.diffrules()
    isdefined(ReverseDiff, M) || continue
    if arity == 1
        @eval @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray}}) = _materialize(bc.f, bc.args)
    elseif arity == 2
        @eval begin
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedArray}}) = _materialize(bc.f, bc.args)
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedReal}}) = _materialize(bc.f, bc.args)
            @noinline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,TrackedArray}}) = _materialize(bc.f, bc.args)
        end
        for A in ARRAY_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A{<:Number},TrackedArray}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray, $A{<:Number}}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A{<:Number}, TrackedReal}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,$A{<:Number}}}) = _materialize(bc.f, bc.args)
            end
        end
        for R in REAL_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$R,TrackedArray}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,$R}}) = _materialize(bc.f, bc.args)
            end
        end
    end
end


text:
#################################################################################################
#
# demo 7:  Image Histogram in OpenCV.jl
#
# https://github.com/bsdnoobz/opencv-code/blob/master/display-histogram.cpp
#################################################################################################

# Julia filename with full path
inputfile = joinpath(Pkg.dir("OpenCV"), "./test/images/mandrill.jpg")
outfile = joinpath(Pkg.dir("OpenCV"), "./test/images/mandrillHist.jpg")

cxx"""
void showHistogram(cv::Mat& img)
{
	int bins = 256;             // number of bins
	int nc = img.channels();    // number of channels

	std::vector<cv::Mat> hist(nc);      // histogram arrays

	// Initalize histogram arrays
	for (int i = 0; i < hist.size(); i++)
		hist[i] = cv::Mat::zeros(1, bins, CV_32SC1);

	// Calculate the histogram of the image
	for (int i = 0; i < img.rows; i++)
	{
		for (int j = 0; j < img.cols; j++)
		{
			for (int k = 0; k < nc; k++)
			{
				uchar val = nc == 1 ? img.at<uchar>(i,j) : img.at<cv::Vec3b>(i,j)[k];
				hist[k].at<int>(val) += 1;
			}
		}
	}

	// For each histogram arrays, obtain the maximum (peak) value
	// Needed to normalize the display later
	int hmax[3] = {0,0,0};
	for (int i = 0; i < nc; i++)
	{
		for (int j = 0; j < bins-1; j++)
			hmax[i] = hist[i].at<int>(j) > hmax[i] ? hist[i].at<int>(j) : hmax[i];
	}

	const char* wname[3] = {"blue", "green", "red"};
	cv::Scalar colors[3] = {cv::Scalar(255,0,0), cv::Scalar(0,255,0), cv::Scalar(0,0,255)};

	std::vector<cv::Mat> canvas(nc);

	// Display each histogram in a canvas
	for (int i = 0; i < nc; i++)
	{
		canvas[i] = cv::Mat::ones(125, bins, CV_8UC3);

		for (int j = 0, rows = canvas[i].rows; j < bins-1; j++)
		{
			cv::line(
				canvas[i],
				cv::Point(j, rows),
				cv::Point(j, rows - (hist[i].at<int>(j) * rows/hmax[i])),
				nc == 1 ? cv::Scalar(200,200,200) : colors[i],
				1, 8, 0
			);
		}

		cv::imshow(nc == 1 ? "value" : wname[i], canvas[i]);
	}
}

// Test the `showHistogram` function above
void imHist(const char *inputfile, const char *outfile)
{
	cv::Mat src = cv::imread(inputfile);
	if (src.empty()) {
     std::cout << "Can not open the file!" << std::endl;
	   exit(0);
  }
	showHistogram(src);
	cv::imshow("src", src);
	cv::waitKey(0);
  cv::destroyAllWindows();
}
"""

@cxx imHist(pointer(inputfile), pointer(outfile))

ext: jl


text:
const eightschools ="
data {
  int<lower=0> J; // number of schools
  real y[J]; // estimated treatment effects
  real<lower=0> sigma[J]; // s.e. of effect estimates
}
parameters {
  real mu;
  real<lower=0> tau;
  real eta[J];
}
transformed parameters {
  real theta[J];
  for (j in 1:J)
    theta[j] <- mu + tau * eta[j];
}
model {
  eta ~ normal(0, 1);
  y ~ normal(theta, sigma);
}
"


text:
# This file was generated, do not modify it. # hide
unique(iris.Species)


text:
"""
    @alert [duration=2.0] [message="Done!"] [onerror=true] [body...]

Calls [`alert`](@ref) if `body` takes longer than `duration` (default to 2.0)
seconds to complete; Posts the alert even if `body` throws an exception so long
as `onerror=true`.

Settings (e.g. `duration`) are specified as keyword arguments.

> NOTE: when `onerror=true` the body is run in a `try/catch` block, meaning
any variables you define inside `@alert` are local.
"""
macro alert(args...)
    if length(args) == 0
        error("Missing body for @alert macro.")
    end
    options, body = parse_alert_args(args)
    options = NamedTuple(options)

    if get(options, :onerror, true)
        return quote
            start_time = Dates.now()
            result = nothing
            try
                result = begin
                    $(esc(body))
                end
            finally
                delay, msg = Alert.__at_alert_options__(;$options...)
                if !isinf(delay) && (Dates.now() - start_time) >= Dates.Millisecond(round(Int,1000delay))
                    alert(msg)
                end
            end

            result
        end
    else
        return quote
            start_time = Dates.now()
            result = nothing
            result = begin
                $(esc(body))
            end

            delay, msg = Alert.__at_alert_options__(;$options...)
            if !isinf(delay) && (Dates.now() - start_time) >= Dates.Millisecond(round(Int,1000delay))
                alert(msg)
            end

            result
        end
    end
end

function parse_alert_args(args)
    opts = []

    i = 0
    for outer i in 1:length(args)-1
        if @capture(args[i], var_ = body_)
            push!(opts, var => body)
        else
            break
        end
    end

    opts, args[i+1]
end

__at_alert_options__(;duration=2.0,message="Done!",onerror=true) = duration, message, onerror


text:
module Gridworld

using Requires

const GW = Gridworld
export GW

include("objects.jl")
include("grid_world_base.jl")
include("abstract_grid_world.jl")
include("envs/envs.jl")
include("render_in_terminal.jl")

# function __init__()
#     @require Makie="ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" include("render_with_Makie.jl")
# end
include("render_with_Makie.jl")

end


text:
################################################################################
#
#   module LatPhysBandstructuresPlottingPyPlot
#   -> PyPlot
#   -> LatPhysBase
#   -> LatPhysReciprocal
#   -> LatPhysBandstructures
#   -> LinearAlgebra
#
#   --> PLOTTING OF BANDSTRUCTURES AND FERMI SURFACES
#           - 2D lattices
#           - 3D lattices
#
################################################################################

# start of module
module LatPhysBandstructuresPlottingPyPlot

# used libraries
using LatPhysBase
using LatPhysReciprocal
using LatPhysBandstructures
using LinearAlgebra
using LatPhysPlottingPyPlot
using LatPhysReciprocalPlottingPyPlot
using LaTeXStrings
using PyPlot

# explicitly import PyPlot.plot to overwrite
import PyPlot.plot

# include plotting of bandstructures
include("plotting_bandstructures/bandstructure_plotting.jl")

# include plotting of energy manifoldss
include("plotting_energy_manifolds/energy_manifold_plotting.jl")

end # module


text:
module Certificate

import MutableArithmetics
const MA = MutableArithmetics
import MultivariatePolynomials
const MP = MultivariatePolynomials
import MultivariateBases
const MB = MultivariateBases
using SemialgebraicSets

using SumOfSquares

include("newton_polytope.jl")

# Certificate              denominator * p = ...                                      domain
# Polya       (1 + x1 + ... + xn)^(2d) * p = ∑ λ_α x.^α                  λ_α ∈ R_+    R_+^n
# Putinar                                p = σ_0 + ∑ σ_i g_i             σ_i SOS      g_i ≥ 0
# Krivine                                p = ∑ λ_{αβ} g.^α (1 .- g).^β   λ_{αβ} ≥ 0   0 ≤ g_i ≤ 1
# Schmüdgen                              p = ∑ σ_α g^α                   σ_α SOS      g_i ≥ 0
# Hilbert                            d * p = σ                           d, σ SOS     R^n

struct PreorderIndex
    value::Int
end
struct IdealIndex
    value::Int
end

const Index = Union{PreorderIndex, IdealIndex}

abstract type Attribute end

# For get
struct Cone <: Attribute end
struct GramBasis <: Attribute end
# FIXME currently, this returns `MB.MonomialBasis` instead of `MB.MonomialBasis{MT, MVT}`
struct GramBasisType <: Attribute end
struct ReducedPolynomial <: Attribute end
struct IdealCertificate <: Attribute end
struct PreprocessedDomain <: Attribute end

# Only for PreorderIndex
struct PreorderIndices <: Attribute end
struct Generator <: Attribute end
struct MultiplierBasis <: Attribute end
struct MultiplierBasisType <: Attribute end

# For set
#struct Monomials <: Attribute end

# PreorderCertificate
# IdealCertificate
# FullSpaceCertificate

abstract type AbstractCertificate end

abstract type AbstractPreorderCertificate <: AbstractCertificate end
abstract type AbstractIdealCertificate <: AbstractCertificate end

"""
    struct Putinar{IC <: AbstractIdealCertificate, CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: AbstractPreorderCertificate
        ideal_certificate::IC
        cone::CT
        basis::Type{BT}
        maxdegree::Int
    end

The `Putinar` certificate ensures the nonnegativity of `p(x)` for all `x` such that
`g_i(x) >= 0` and `h_i(x) = 0` by exhibiting Sum-of-Squares polynomials `σ_i(x)`
such that `p(x) - ∑ σ_i(x) g_i(x)` is guaranteed to be nonnegativity for all `x`
such that `h_i(x) = 0`.
The polynomials `σ_i(x)` are search over `cone` with a basis of type `basis` such that
the degree of `σ_i(x) g_i(x)` does not exceed `maxdegree`.
"""
struct Putinar{IC <: AbstractIdealCertificate, CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: AbstractPreorderCertificate
    ideal_certificate::IC
    cone::CT
    basis::Type{BT}
    maxdegree::Int
end

get(certificate::Putinar, ::Cone) = certificate.cone

struct DomainWithVariables{S, V}
    domain::S
    variables::V
end
function get(::Putinar, ::PreprocessedDomain, domain::BasicSemialgebraicSet, p)
    return DomainWithVariables(domain, MP.variables(p))
end

function get(::Putinar, ::PreorderIndices, domain::DomainWithVariables)
    return map(PreorderIndex, eachindex(domain.domain.p))
end

function maxdegree_gram_basis(B::Type, variables, maxdegree::Int)
    return MB.maxdegree_basis(B, variables, div(maxdegree, 2))
end
function multiplier_maxdegree(maxdegree, q)
    maxdegree_s2 = maxdegree - MP.maxdegree(q)
    # If maxdegree_s2 is odd, `div(maxdegree_s2, 2)` would make s^2 have degree up to maxdegree_s2-1
    # for this reason, we take `div(maxdegree_s2 + 1, 2)` so that s^2 have degree up to maxdegree_s2+1
    return maxdegree_s2 + 1
end
function get(certificate::Putinar, ::MultiplierBasis, index::PreorderIndex, domain::DomainWithVariables)
    q = domain.domain.p[index.value]
    vars = sort!([domain.variables..., MP.variables(q)...], rev = true)
    unique!(vars)
    return maxdegree_gram_basis(certificate.basis, vars, multiplier_maxdegree(certificate.maxdegree, q))
end
function get(::Type{Putinar{IC, CT, BT}}, ::MultiplierBasisType) where {IC, CT, BT}
    return BT
end

function get(::Putinar, ::Generator, index::PreorderIndex, domain::DomainWithVariables)
    return domain.domain.p[index.value]
end

get(certificate::Putinar, ::IdealCertificate) = certificate.ideal_certificate
get(::Type{<:Putinar{IC}}, ::IdealCertificate) where {IC} = IC

SumOfSquares.matrix_cone_type(::Type{<:Putinar{IC, CT}}) where {IC, CT} = SumOfSquares.matrix_cone_type(CT)

######################
# Ideal certificates #
######################

abstract type SimpleIdealCertificate{CT, BT} <: AbstractIdealCertificate end
get(::SimpleIdealCertificate, ::ReducedPolynomial, poly, domain) = poly

get(certificate::SimpleIdealCertificate, ::Cone) = certificate.cone
SumOfSquares.matrix_cone_type(::Type{<:SimpleIdealCertificate{CT}}) where {CT} = SumOfSquares.matrix_cone_type(CT)

# TODO return something else when `PolyJuMP` support other bases.
zero_basis(certificate::SimpleIdealCertificate) = MB.MonomialBasis
zero_basis_type(::Type{<:SimpleIdealCertificate{CT, BT}}) where {CT, BT} = MB.MonomialBasis

"""
    struct MaxDegree{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: SimpleIdealCertificate{CT, BT}
        cone::CT
        basis::Type{BT}
        maxdegree::Int
    end

The `MaxDegree` certificate ensures the nonnegativity of `p(x)` for all `x` such that
`h_i(x) = 0` by exhibiting a Sum-of-Squares polynomials `σ(x)`
such that `p(x) - σ(x)` is guaranteed to be zero for all `x`
such that `h_i(x) = 0`.
The polynomial `σ(x)` is search over `cone` with a basis of type `basis` such that
the degree of `σ(x)` does not exceed `maxdegree`.
"""
struct MaxDegree{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: SimpleIdealCertificate{CT, BT}
    cone::CT
    basis::Type{BT}
    maxdegree::Int
end
function get(certificate::MaxDegree, ::GramBasis, poly) where CT
    return maxdegree_gram_basis(certificate.basis, MP.variables(poly), certificate.maxdegree)
end
function get(::Type{MaxDegree{CT, BT}}, ::GramBasisType) where {CT, BT}
    return BT
end

"""
    struct FixedBasis{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: SimpleIdealCertificate{CT, BT}
        cone::CT
        basis::BT
    end

The `FixedBasis` certificate ensures the nonnegativity of `p(x)` for all `x` such that
`h_i(x) = 0` by exhibiting a Sum-of-Squares polynomials `σ(x)`
such that `p(x) - σ(x)` is guaranteed to be zero for all `x`
such that `h_i(x) = 0`.
The polynomial `σ(x)` is search over `cone` with basis `basis`.
"""
struct FixedBasis{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis} <: SimpleIdealCertificate{CT, BT}
    cone::CT
    basis::BT
end
function get(certificate::FixedBasis, ::GramBasis, poly) where CT
    return certificate.basis
end
function get(::Type{FixedBasis{CT, BT}}, ::GramBasisType) where {CT, BT}
    return BT
end

"""
    struct Newton{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis, NPT <: Tuple} <: SimpleIdealCertificate{CT, BT}
        cone::CT
        basis::Type{BT}
        variable_groups::NPT
    end

The `Newton` certificate ensures the nonnegativity of `p(x)` for all `x` such that
`h_i(x) = 0` by exhibiting a Sum-of-Squares polynomials `σ(x)`
such that `p(x) - σ(x)` is guaranteed to be zero for all `x`
such that `h_i(x) = 0`.
The polynomial `σ(x)` is search over `cone` with a basis of type `basis`
chosen using the multipartite Newton polytope with parts `variable_groups`.
If `variable_groups = tuple()` then it falls back to the classical Newton polytope
with all variables in the same part.
"""
struct Newton{CT <: SumOfSquares.SOSLikeCone, BT <: MB.AbstractPolynomialBasis, NPT <: Tuple} <: SimpleIdealCertificate{CT, BT}
    cone::CT
    basis::Type{BT}
    variable_groups::NPT
end
function get(certificate::Newton{CT, B}, ::GramBasis, poly) where {CT, B}
    return MB.basis_covering_monomials(B, monomials_half_newton_polytope(MP.monomials(poly), certificate.variable_groups))
end
function get(::Type{<:Newton{CT, BT}}, ::GramBasisType) where {CT, BT}
    return BT
end

"""
    struct Remainder{GCT<:AbstractIdealCertificate} <: AbstractIdealCertificate
        gram_certificate::GCT
    end

The `Remainder` certificate ensures the nonnegativity of `p(x)` for all `x` such that
`h_i(x) = 0` by guaranteeing the remainder of `p(x)` modulo the ideal generated by
`⟨h_i⟩` to be nonnegative for all `x` such that `h_i(x) = 0` using the certificate
`gram_certificate`.
For instance, if `gram_certificate` is [`SumOfSquares.Certificate.Newton`](@ref),
then the certificate `Remainder(gram_certificate)` will take the remainder before
computing the Newton polytope hence might generate a much smaller Newton polytope
hence a smaller basis and smaller semidefinite program.
However, this then corresponds to a lower degree of the hierarchy which might
be insufficient to find a certificate.
"""
struct Remainder{GCT<:AbstractIdealCertificate} <: AbstractIdealCertificate
    gram_certificate::GCT
end

function get(::Remainder, ::ReducedPolynomial, poly, domain)
    return convert(typeof(poly), rem(poly, ideal(domain)))
end

function get(certificate::Remainder, attr::GramBasis, poly)
    return get(certificate.gram_certificate, attr, poly)
end
function get(::Type{Remainder{GCT}}, attr::GramBasisType) where GCT
    return get(GCT, attr)
end

get(certificate::Remainder, attr::Cone) = get(certificate.gram_certificate, attr)
SumOfSquares.matrix_cone_type(::Type{Remainder{GCT}}) where {GCT} = SumOfSquares.matrix_cone_type(GCT)
zero_basis(certificate::Remainder) = zero_basis(certificate.gram_certificate)
zero_basis_type(::Type{Remainder{GCT}}) where {GCT} = zero_basis_type(GCT)

include("sparse/ChordalExtensionGraph.jl")
using .ChordalExtensionGraph: ChordalCompletion, ClusterCompletion
export ChordalCompletion, ClusterCompletion

include("sparse/sparse_putinar.jl")

include("Symmetry/Symmetry.jl")
export Symmetry

end


text:
using CompScienceMeshes
using StaticArrays
using Test

p1 = point(1.0, 0.0, 0.0)
p2 = point(10.0, 0.0, 1.0)
p3 = point(0.0, 2.0, 5.0)

ch = simplex(p1,p2,p3)

dom = domain(ch)
q, u = quadpoints(dom, 7)
@which neighborhood(ch, q[1])

pw = CompScienceMeshes.quadpoints(ch, 7);

#t1 = sum(w)
t1 = sum(q[2] for q in pw)
t2 = 0.5*norm(cross(p1-p3, p2-p3))

@test 2 * (t1-t2) / (t1+t2) + 1  ≈ 1

p1 = point(1.0, 0.0, 0.0)
p2 = point(2.0, 0.0, 0.0)
seg = simplex(p1,p2)

wpvs = quadpoints(x->x[1]^2, [seg], (20,))
r = 0.0
for wpv in wpvs[1,1]
    w = wpv.weight
    v = wpv.value
    global r += w * v
end

@test r ≈ 1/3*(2^3-1^3)


text:
using Test
using POMDPs
using POMDPTesting
using POMDPModelTools

import POMDPs: transition, observation, initialstate_distribution, updater

struct TestPOMDP <: POMDP{Bool, Bool, Bool} end
updater(problem::TestPOMDP) = DiscreteUpdater(problem)
initialstate_distribution(::TestPOMDP) = BoolDistribution(0.0)
transition(p::TestPOMDP, s::Bool, a::Bool) = BoolDistribution(0.5)
observation(p::TestPOMDP, a::Bool, sp::Bool) = BoolDistribution(0.5)
@testset "model" begin
    POMDPTesting.probability_check(TestPOMDP())
end


text:
cd(@__DIR__); runapprox = false; include("setups/sunflower.jl")
gr(dpi = 200)

## (a) sunflower graph
gplot(W, X; width = 1)
scatter_gplot!(X; c = :red, ms = LinRange(1, 9, N))
plt = plot!(frame = :none)
savefig(plt, "../figs/SunFlower.png")

## (b) voronoi tessellation
xx, yy = getplotxy(voronoiedges(tess))
plt = plot(xx, yy, xlim = [1, 2], ylim = [1, 2], linestyle = :auto, linewidth = 1,
    linecolor = :blue, grid = false, label = "", ratio = 1, frame = :box, ticks = false)
savefig(plt, "../figs/Sunflower_Barbara_Voronoi_cells.png")


text:
### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Plots, LaTeXStrings
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.6"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	xGrid = 0:0.01:10
	uGrid = 0:0.01:1
	busy = 0.8
end

# ╔═╡ 80359808-20f8-11eb-01d8-a96d178126b4
begin
	F(t)= t<=0 ? 0 : 1 - busy*exp(-(1-busy)t)
	infimum(B) = isempty(B) ? Inf : minimum(B)
	invF(u) = infimum(filter((x) -> (F(x) >= u),xGrid))
end

# ╔═╡ 805b0906-20f8-11eb-1142-e99c34f90ba2
begin
	p1 = plot(xGrid, F.(xGrid), c=:blue, xlims=(-0.1,10), ylims=(0,1),
		xlabel=L"x", ylabel=L"F(x)")

	p2 = plot(uGrid, invF.(uGrid), c=:blue, xlims=(0,0.95), ylims=(0,maximum(xGrid)),
		xlabel=L"u", ylabel=L"F^{-1}(u)")

	plot(p1, p2, legend=:none, size=(800, 400))
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.6"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═80359808-20f8-11eb-01d8-a96d178126b4
# ╠═805b0906-20f8-11eb-1142-e99c34f90ba2
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


text:
VERSION >= v"0.4.0-dev+6521" && __precompile__()

module MDCT
export mdct, imdct

import Base.FFTW.fftwNumber, Base.FFTW.r2r, Base.FFTW.r2r!, Base.FFTW.REDFT11

fftwsimilar{T<:fftwNumber}(X::AbstractArray{T}, sz) = Array(T, sz...)
fftwsimilar{T<:Real}(X::AbstractArray{T}, sz) = Array(Float64, sz...)
fftwsimilar{T<:Complex}(X::AbstractArray{T}, sz) = Array(Complex128, sz...)

# The following two routines compute the MDCT and IMDCT via
# a type-IV DCT (FFTW's REDFT11 r2r transform).  For a review
# of this relationship, see the notes I posted on Wikipedia:
#     http://en.wikipedia.org/wiki/Modified_discrete_cosine_transform

function mdct{T<:Number}(X::AbstractVector{T})
    sz = length(X)
    if isodd(sz)
        throw(ArgumentError("mdct requires an even-length vector"))
    end
    N = div(sz, 2);
    Y = fftwsimilar(X, N)
    if isodd(N)
        throw(ArgumentError("mdct requires a multiple-of-4 vector length"))
        # FIXME: handle odd case via DCT-III?
    else
        N2 = div(N,2)
        for i = 1:N2
            Y[i] = -0.5 * (X[(3*N2+1)-i] + X[3*N2+i])
            Y[N2+i] = -0.5 * (X[(2*N2+1)-i] - X[i])
        end
        r2r!(Y, REDFT11)
    end
    return Y
end

function imdct{T<:Number}(X::StridedVector{T})
    N = length(X)
    Z = fftwsimilar(X, 2*N)
    if isodd(N)
        throw(ArgumentError("imdct requires an even vector length"))
        # FIXME: handle odd case via DCT-II?
    else
        Y = r2r(X, REDFT11)
        N2 = div(N,2)
        s = 0.5 / N
        for i = 1:N2
            Z[N+1-i] = -(Z[i] = Y[N2+i]*s)
            Z[N+N2+1-i] = (Z[N+N2+i] = -Y[i]*s)
        end
    end
    return Z
end

imdct{T<:Number}(X::AbstractVector{T}) = imdct(copy!(fftwsimilar(X, size(X)),
                                                     X))

end # MDCT


text:
using SafeTestsets
@safetestset "My_calc Tests" begin include("My_calc_test.jl") end


text:
export CountingMeasure

struct CountingMeasure{X} <: PrimitiveMeasure end

@primitive CountingMeasure

function Base.show(io::IO, ::MIME"text/plain", μ::CountingMeasure{X}) where {X}
    io = IOContext(io, :compact => true)
    print(io, "CountingMeasure(", X, ")")
end

CountingMeasure(X) = CountingMeasure{X}()

# sampletype(::CountingMeasure{ℝ}) = Float64
# sampletype(::CountingMeasure{ℝ₊}) = Float64
# sampletype(::CountingMeasure{𝕀}) = Float64

sampletype(::CountingMeasure) = Int

testvalue(μ::CountingMeasure{X}) where {X} = testvalue(X)

logdensity(::CountingMeasure, x) = zero(float(x))

# (::CountingMeaure)(s) = length(Set(s))


text:
using StatsPlots, Dates, CSV, DataFrames, JLD2, NamedArrays

## Get vaccine scenarios

t_0 = (Date(2021, 9, 27) - Date(2020, 2, 20)).value #Vaccines start Sept 2021
t_end = (Date(2021, 9, 27) + Month(18) - Date(2020, 2, 20)).value #Option where vaccines finish in 18 months
t_end_rapid = (Date(2021, 9, 27) + Month(6) - Date(2020, 2, 20)).value #Option where vaccines finish in 6 months
ts = [0; t_0; t_end; t_end + 100] .+ 14 # Time points for 18 month deployment lagged 14 days to account for delay between jab and maximum effectiveness
ts_rapid = [0; t_0; t_end_rapid; t_end_rapid + 100] .+ 14 # Time points for 6 month deployment

# Number of people DOUBLE jabbed by time points under three coverages

fully_vaccinated_30perc = [0; 0; 15.9e6 / 2; 15.9e6 / 2]
fully_vaccinated_50perc = [0; 0; 25e6 / 2; 25e6 / 2]
fully_vaccinated_70perc = [0; 0; 35e6 / 2; 35e6 / 2]

#Daily rate of vaccination
scen_str = ["no vac", "30 perc", "50 perc", "70 perc", "30 perc rapid", "50 perc rapid", "70 perc rapid"]

scen_vac_rates = [0.0, (fully_vaccinated_30perc[3] - fully_vaccinated_30perc[2]) / (ts[3] - ts[2]),
    (fully_vaccinated_50perc[3] - fully_vaccinated_50perc[2]) / (ts[3] - ts[2]),
    (fully_vaccinated_70perc[3] - fully_vaccinated_70perc[2]) / (ts[3] - ts[2]),
    (fully_vaccinated_30perc[3] - fully_vaccinated_30perc[2]) / (ts_rapid[3] - ts_rapid[2]),
    (fully_vaccinated_50perc[3] - fully_vaccinated_50perc[2]) / (ts_rapid[3] - ts_rapid[2]),
    (fully_vaccinated_70perc[3] - fully_vaccinated_70perc[2]) / (ts_rapid[3] - ts_rapid[2])]

N_ovr_50 = sum(N_kenya[11:end, :])
N_18_49 = sum(N_kenya[5:10, :]) + 0.4 * sum(N_kenya[4, :])
max_level_ovr_50s = N_ovr_50 .* 0.8
max_level_18_49 = [0, fully_vaccinated_30perc[end], fully_vaccinated_50perc[end], fully_vaccinated_70perc[end],
    fully_vaccinated_30perc[end], fully_vaccinated_50perc[end], fully_vaccinated_70perc[end]] .- max_level_ovr_50s

max_level_18_49_perc = 100 .* max_level_18_49 ./ N_18_49

t_hit_max_lvl_ovr_50 = map(rate -> max_level_ovr_50s / rate, scen_vac_rates)
t_hit_max_lvl_18_49 = map((rate, max_lvl) -> max_lvl / rate, scen_vac_rates, max_level_18_49) .+ t_hit_max_lvl_ovr_50

## Get population sizes

@load("data/N_kenya.jld2")

## X ticks and dates

xtickdates = [Date(2020, 2, 1) + Month(k) for k = 19:29]
xtickpos = [(d - Date(2020, 2, 19)).value for d in xtickdates]
xticklab = [monthname(d)[1:3] * "-" * string(year(d) - 2000) for d in xtickdates]
t_sept_model = (Date(2021, 9, 1) - Date(2021, 1, 1)).value
t_sept_data = (Date(2021, 9, 1) - Date(2020, 2, 19)).value
t_end_plot = (Date(2022, 7, 1) - Date(2020, 2, 19)).value

t_start_model = (Date(2021, 1, 1) - Date(2020, 2, 19)).value
t_end_model = (Date(2022, 7, 1) - Date(2020, 2, 19)).value
t_end_model_fit = (Date(2021, 9, 1) - Date(2021, 1, 1)).value
t_end_fit = (Date(2021, 9, 1) - Date(2020, 2, 19)).value

xs_model_proj = (t_end_fit+1):t_end_model
xs_model_fitted_to_data = t_start_model:t_end_fit

n_model_proj = length(xs_model_proj)
n_model_fit = length(xs_model_fitted_to_data)

## Plot cumulative vaccines over 50s

plt_vac_ovr_50 = plot(xticks = (xtickpos, xticklab),
    size = (800, 400),
    left_margin = 5mm, right_margin = 2mm,
    xlims = (t_end_fit, t_end_plot + 5),
    ylims = (0, 100),
    legend = :bottomright,
    ylabel = "Fully vaccinated (%)",
    title = "Scenarios: proportion fully vaccinated (over 50 yo)",
    tickfont = 11, guidefont = 13, titlefont = 17)

plot!(plt_vac_ovr_50, [ts[1:2]; [10000, 100000]], [0, 0, 0, 0],
    lw = 2, color = :grey, lab = "No vaccination")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[2] + ts[2], t_hit_max_lvl_ovr_50[2] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, color = 1, lab = "30% coverage")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[3] + ts[2], t_hit_max_lvl_ovr_50[3] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, color = 2, lab = "50% coverage")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[4] + ts[2], t_hit_max_lvl_ovr_50[4] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, color = 3, lab = "70% coverage")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[5] + ts[2], t_hit_max_lvl_ovr_50[5] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, ls = :dash, color = 1, lab = "30% coverage, rapid rollout")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[6] + ts[2], t_hit_max_lvl_ovr_50[6] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, ls = :dash, color = 2, lab = "50% coverage, rapid rollout")
plot!(plt_vac_ovr_50, [ts[1:2]; [t_hit_max_lvl_ovr_50[7] + ts[2], t_hit_max_lvl_ovr_50[7] + ts[2] + 1000]], [0, 0, 80, 80],
    lw = 2, ls = :dash, color = 3, lab = "70% coverage, rapid rollout")

# savefig(plt_vac_ovr_50,"generation_of_data_for_econ_analysis/main_plots/vaccine_rate_ovr_50.png")

## Plot cumulative vaccines 18-49

plt_vac_18_49 = plot(xticks = (xtickpos, xticklab),
    size = (800, 400),
    left_margin = 5mm, right_margin = 2mm,
    xlims = (t_end_fit, t_end_plot + 5),
    ylims = (0, 100),
    legend = :topleft,
    ylabel = "Fully vaccinated (%)",
    title = "Scenarios: proportion fully vaccinated (18-49 yo)",
    tickfont = 11, guidefont = 13, titlefont = 17)

plot!(plt_vac_18_49, [ts[1:2]; [10000, 100000]], [0, 0, 0, 0],
    lw = 2, color = :grey, lab = "No vaccination")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[2] + ts[2], t_hit_max_lvl_18_49[2] + ts[2], t_hit_max_lvl_18_49[2] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[2], max_level_18_49_perc[2]],
    lw = 2, color = 1, lab = "30% coverage")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[3] + ts[2], t_hit_max_lvl_18_49[3] + ts[2], t_hit_max_lvl_18_49[3] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[3], max_level_18_49_perc[3]],
    lw = 2, color = 2, lab = "50% coverage")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[4] + ts[2], t_hit_max_lvl_18_49[4] + ts[2], t_hit_max_lvl_18_49[4] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[4], max_level_18_49_perc[4]],
    lw = 2, color = 3, lab = "70% coverage")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[5] + ts[2], t_hit_max_lvl_18_49[5] + ts[2], t_hit_max_lvl_18_49[5] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[5], max_level_18_49_perc[5]],
    lw = 2, ls = :dash, color = 1, lab = "30% coverage, rapid rollout")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[6] + ts[2], t_hit_max_lvl_18_49[6] + ts[2], t_hit_max_lvl_18_49[6] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[6], max_level_18_49_perc[6]],
    lw = 2, ls = :dash, color = 2, lab = "50% coverage, rapid rollout")
plot!(plt_vac_18_49, [ts[1:2]; [t_hit_max_lvl_ovr_50[7] + ts[2], t_hit_max_lvl_18_49[7] + ts[2], t_hit_max_lvl_18_49[7] + ts[2] + 1000]], [0, 0, 0, max_level_18_49_perc[7], max_level_18_49_perc[7]],
    lw = 2, ls = :dash, color = 3, lab = "70% coverage, rapid rollout")

# savefig(plt_vac_18_49, "generation_of_data_for_econ_analysis/main_plots/vaccine_rate_18_49.png")


text:
UA_Client_getContext(client::UA_Client) = UA_Client_getConfig(client).clientContext

function UA_Client_connectUsername(client::Ptr{UA_Client}, endpointUrl::AbstractString, username::AbstractString, password::AbstractString)
    identityToken = UA_UserNameIdentityToken_new()
    identityToken == CNULL && return UA_STATUSCODE_BADOUTOFMEMORY

    identityToken.userName = UA_STRING_ALLOC(username)
    identityToken.password = UA_STRING_ALLOC(password)
    cc = UA_Client_getConfig(client)
    UA_ExtensionObject_clear(cc.userIdentityToken)
    cc.userIdentityToken.encoding = UA_EXTENSIONOBJECT_DECODED
    cc.userIdentityToken.content.decoded.type = UA_TYPES_PTRS[UA_TYPES_USERNAMEIDENTITYTOKEN]
    cc.userIdentityToken.content.decoded.data = identityToken
    return UA_Client_connect(client, endpointUrl)
end

## UA_Client_Service functions
for att in attributes_UA_Client_Service
    fun_name = Symbol(att[1])
    req_type = Symbol("UA_", uppercasefirst(att[2]), "Request")
    resp_type = Symbol("UA_", uppercasefirst(att[2]), "Response")
    req_type_ptr = Symbol("UA_TYPES_", uppercase(String(att[2])), "REQUEST")
    resp_type_ptr = Symbol("UA_TYPES_", uppercase(String(att[2])), "RESPONSE")

    @eval begin
        if @isdefined $(req_type) # Skip functions that use undefined types, e.g. deactivated historizing types
            function $(fun_name)(client::Ptr{UA_Client}, request::Ptr{$(req_type)})
                response = Ref{$(resp_type)}()
                retval = __UA_Client_Service(
                    client,
                    request,
                    UA_TYPES_PTRS[$(req_type_ptr)],
                    response,
                    UA_TYPES_PTRS[$(resp_type_ptr)]
                )
                if isnothing(retval) || retval == UA_STATUSCODE_GOOD
                    return response[]
                else
                    error("Service request of type ´$(req_type)´ from UA_Client failed with statuscode \"$(UA_StatusCode_name_print(retval))\".")
                end
            end
        end
    end
end

## Read attribute functions
for att in attributes_UA_Client_read
    fun_name = Symbol(att[1])
    attr_name = Symbol(att[2])
    ret_type = Symbol(att[3])
    ret_type_ptr = Symbol("UA_TYPES_", uppercase(String(ret_type)[4:end]))
    ua_attr_name = Symbol("UA_ATTRIBUTEID_", uppercase(att[2]))

    @eval begin
        function $(fun_name)(client::Ptr{UA_Client}, nodeId::Ref{UA_NodeId})
            data_type_ptr = UA_TYPES_PTRS[$(ret_type_ptr)]
            out = Ref{$(ret_type)}()
            retval = __UA_Client_readAttribute(client, nodeId, $(ua_attr_name), out, data_type_ptr)
            if retval == UA_STATUSCODE_GOOD
                return out[]
            else
                error("Reading attribute ´$(attr_name)´ from UA_Client failed with statuscode \"$(UA_StatusCode_name_print(retval))\".")
            end
        end

        function $(fun_name)(client::Ptr{UA_Client}, nodeId::UA_NodeId)
            return $(fun_name)(client, Ref(nodeId))
        end
    end
end

## Write attribute functions
for att in attributes_UA_Client_write
    fun_name = Symbol(att[1])
    attr_name = Symbol(att[2])
    attr_type = Symbol(att[3])
    attr_type_ptr = Symbol("UA_TYPES_", uppercase(String(attr_type)[4:end]))
    ua_attr_name = Symbol("UA_ATTRIBUTEID_", uppercase(att[2]))

    @eval begin
        function $(fun_name)(client::Ptr{UA_Client}, nodeId::Ref{UA_NodeId}, new_attr::$(attr_type))
            data_type_ptr = UA_TYPES_PTRS[$(attr_type_ptr)]
            retval = __UA_Client_writeAttribute(client, nodeId, $(ua_attr_name), new_attr, data_type_ptr)
            if retval == UA_STATUSCODE_GOOD
                return retval
            else
                error("Writing attribute ´$(attr_name)´ on UA_Client failed with statuscode \"$(UA_StatusCode_name_print(retval))\".")
            end
        end

        function $(fun_name)(client::Ptr{UA_Client}, nodeId::UA_NodeId, new_attr::$(attr_type))
            return $(fun_name)(client, Ref(nodeId), new_attr)
        end
    end
end

## Read attribute async functions
for att in attributes_UA_Client_read_async
    fun_name = Symbol(att[1])
    attr_name = Symbol(att[2])
    ret_type = Symbol(att[3])
    ret_type_ptr = Symbol("UA_TYPES_", uppercase(String(ret_type)[4:end]))
    ua_attr_name = Symbol("UA_ATTRIBUTEID_", uppercase(att[2]))

    @eval begin
        function $(fun_name)(client::Ptr{UA_Client}, nodeId::Ref{UA_NodeId}, callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
            data_type_ptr = UA_TYPES_PTRS[$(ret_type_ptr)]
            retval = __UA_Client_readAttribute_async(
                client,
                nodeId,
                $(ua_attr_name),
                data_type_ptr,
                reinterpret(UA_ClientAsyncServiceCallback, callback),
                userdata,
                reqId
            )
            if retval == UA_STATUSCODE_GOOD
                return retval
            else
                error("Reading attribute ´$(attr_name)´ asynchronously from UA_Client failed with statuscode \"$(UA_StatusCode_name_print(retval))\".")
            end
        end

        function $(fun_name)(client::Ptr{UA_Client}, nodeId::UA_NodeId, callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
            return $(fun_name)(client, Ref(nodeId), callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
        end
    end
end

# ## Write attribute async functions
for att in attributes_UA_Client_write_async
    fun_name = Symbol(att[1])
    attr_name = Symbol(att[2])
    attr_type = Symbol(att[3])
    attr_type_ptr = Symbol("UA_TYPES_", uppercase(String(attr_type)[4:end]))
    ua_attr_name = Symbol("UA_ATTRIBUTEID_", uppercase(att[2]))

    @eval begin
        function $(fun_name)(client::Ptr{UA_Client}, nodeId::Ref{UA_NodeId}, out::$(attr_type), callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
            data_type_ptr = UA_TYPES_PTRS[$(attr_type_ptr)]
            retval = __UA_Client_writeAttribute_async(client, nodeId, $(ua_attr_name), out, data_type_ptr, callback, userdata, reqId)
            if retval == UA_STATUSCODE_GOOD
                return retval
            else
                error("Writing attribute ´$(attr_name)´ asynchronously on UA_Client failed with statuscode \"$(UA_StatusCode_name_print(retval))\".")
            end
        end

        function $(fun_name)(client::Ptr{UA_Client}, nodeId::UA_NodeId, out::$(attr_type), callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
            return $(fun_name)(client, Ref(nodeId), out, callback::Ptr{Nothing}, userdata::Ptr{Nothing}, reqId::Integer)
        end
    end
end


text:
using SparseArrays
using MinimalRLCore

include("iht.jl")

# using MimimalRLCore
"""
    TileCoder(num_tilings, num_tiles, num_features, num_ints; wrap=false, wrapwidths=0.0)
Tile coder for coding all features together.
"""
mutable struct TileCoder{F}
    # Main Arguments
    tilings::Int
    tiles::Int
    dims::Int
    ints::Int

    # Optional Arguments
    wrap::Bool
    wrapwidths::Float64

    iht::TileCoding.IHT
    TileCoder{F}(num_tilings, num_tiles, num_features, num_ints=1; wrap=false, wrapwidths=0.0) where {F} =
        new{F}(num_tilings,
               num_tiles,
               num_features,
               num_ints,
               wrap,
               0.0,
               TileCoding.IHT(num_tilings*(num_tiles+1)^num_features * num_ints))
end

TileCoder(args...;kwargs...) = TileCoder{Vector{Int}}(args...; kwargs...)
SparseTileCoder(args...;kwargs...) = TileCoder{SparseVector{Int}}(args...; kwargs...)

function get_tc_indicies(fc::TileCoder, s; ints=[], readonly=false)
    if fc.wrap
        return 1 .+ TileCoding.tileswrap!(fc.iht, fc.tilings, s.*fc.tiles, fc.wrapwidths, ints, readonly)
    else
        return 1 .+ TileCoding.tiles!(fc.iht, fc.tilings, s.*fc.tiles, ints, readonly)
    end
end

MinimalRLCore.create_features(fc::TileCoder{Vector{Int}}, s; ints=[], readonly=false) =
    get_tc_indicies(fc, s; ints=[], readonly=false)

function MinimalRLCore.create_features(fc::TileCoder{SparseVector{N}}, s; ints=[], readonly=false) where {N}
    idx = get_tc_indicies(fc, s; ints=[], readonly=false)
    s = spzeros(N, size(fc))
    s[idx] .= N(1)
    return s
end

# feature_size(fc::TileCoder) = fc.tilings*(fc.tiles+1)^fc.dims * fc.ints
Base.size(fc::TileCoder) = fc.tilings*(fc.tiles+1)^fc.dims * fc.ints

(fc::TileCoder)(s; ints=[], readonly=false) =
    create_features(fc, s; ints=ints, readonly=readonly)


text:
using JSON
using Profile
using LightGraphs

include("graphindex.jl")
include("plan.jl")
include("flip.jl")
include("cut.jl")
include("bounds.jl")

function load_graph(path::String)::Dict
    f = open(path)
    parsed = JSON.parse(f)
    n_nodes = length(parsed["nodes"])
    assignment = zeros(Int64, n_nodes)
    population = zeros(Int64, n_nodes)
    graph = SimpleGraph(n_nodes)
    for (index, node) in enumerate(parsed["nodes"])
        population[index] = node["population"]
        assignment[index] = node["district"]
    end
    for (index, edges) in enumerate(parsed["adjacency"])
        for edge in edges
            if edge["id"] + 1 > index
                # Julia arrays are 1-indexed, but NetworkX (and Python in general)
                # use 0-indexing. We increment each value in the ID field of each
                # edge by 1 to account for this.
                add_edge!(graph, index, edge["id"] + 1)
            end
        end
    end
    return Dict("population" => population,
                "assignment" => assignment,
                "graph" => graph)
end

function update!(plan::Plan{Int64}, flip::Flip{Int64}, cut::CutDelta)
    @inbounds plan.assignment[flip.node] = flip.new_assignment
    setdiff!(plan.cut_edges, cut.cut_edges_before)
    union!(plan.cut_edges, cut.cut_edges_after)
    @inbounds plan.district_populations[flip.old_assignment] -= flip.population
    @inbounds plan.district_populations[flip.new_assignment] += flip.population
    plan.n_cut_edges += cut.Δ
end

function next!(plan::Plan{Int64}, graph::GraphIndex{Int64},
               pop_bounds::Bounds{Int64}, cut_edges_bound::Int64)::Bool
        flip = random_flip(graph, plan)
        if !population_bounds(flip, plan, pop_bounds)
            return false
        end
        cut = CutDelta(flip, plan, graph)
        if plan.n_cut_edges + cut.Δ  > cut_edges_bound
            return false
        end
        if !contiguous(flip, plan, graph, copy(cut.neighbors))
            return false
        end
        update!(plan, flip, cut)
        return true
end

function main(graph_path)
    # Constants (TODO: config file)
    n_steps = 10000000 #12000000 * 60 * 8
    window_size = 1000000
    min_pop = Int64(0.99 * 1000)  # constraint (hardcoded): minimum district population
    max_pop = Int64(1.01 * 1000)  # constraint (hardcoded): maximum district population
    pop_bounds = Bounds{Int64}(min_pop, max_pop)

    graph_data = load_graph(graph_path)
    graph = GraphIndex(graph_data["graph"], graph_data["population"])
    plan = Plan(graph, graph_data["assignment"])

    max_cut_edges = 2 * plan.n_cut_edges  # constraint: max. cut edges
    accepted_in_window = 0
    accepted = 0
    for step in 1:n_steps
        if step % window_size == 0
            println("window: ", accepted_in_window, "/", window_size)
            accepted_in_window = 0
        end
        if next!(plan, graph, pop_bounds, max_cut_edges)
            accepted += 1
            accepted_in_window += 1
        end
    end
    println()
    println(plan.assignment)
    println("accepted ", accepted, "/", n_steps)
end

@time main("tests/fixtures/horizontal.json")

max_stars_repo_name: pjrule/gridwalk

max_issues_repo_name: pjrule/gridwalk

max_forks_repo_name: pjrule/gridwalk


text:
struct TuneCBTSSim
    mdp
    seed::Int
    n_iters::Int
    max_steps::Int
    A_max::Int
    n_proposes::Int
    log_length_scale::Float64
    log_signal_sigma::Float64
    log_obs_noise::Float64
    ec::Float64
    k_state::Float64
    alpha_state::Float64
    metadata::Dict
end
function TuneCBTSSim()
end
function POMDPs.simulate(sim::TuneCBTSSim)
    rng = MersenneTwister(sim.seed)
    mdp = sim.mdp
    b = CBTSBandit(mdp,
                     A_max=sim.A_max,
                     n_proposes=sim.n_proposes,
                     exploration_constant=sim.ec,
                     log_length_scale=sim.log_length_scale,
                     log_signal_sigma=sim.log_signal_sigma,
                     log_obs_noise=sim.log_obs_noise,
                    )
    solver = ModularSolver(mdp, b, sim.seed; n_iterations=sim.n_iters,
                     k_state=sim.k_state,
                     alpha_state=sim.alpha_state)
    policy = solve(solver, mdp)
    s0 = initial_state(mdp, rng)
    hr = HistoryRecorder(; max_steps=sim.max_steps, rng=rng)
    h = simulate(hr, mdp, policy, s0)
    r_total = sum(h.reward_hist)
    r_total
end

@with_kw struct TuneCBTS
    mdp=Walk1D()
    n_seeds=5
    n_iters=[10,25,50]
    max_steps=20
    A_max=[10.20]
    n_proposes=[10,20]
    log_length_scale=[-1.0,0.0]
    log_signal_sigma=[-1.0,0.0]
    log_obs_noise=[-2.0,-1.0]
    ecs=[0.0,1.0]
    k_states=[0.0,1.0]
    alpha_states=[0.0,0.5]
end
struct TuneCBTSResult
    title::String
    data::DataFrame
end
TuneCBTS(::Type{Walk1D}) = TuneCBTS(Walk1D(), 200,
                                       [300],  #n_iters
                                       20, #max_steps
                                       [10,20,30,50], #A_max
                                       [20,50,100], #n_proposes
                                       [-2.0,-1.0,0.0,0.5], #log_length_scale
                                       [0.0], #log_signal_sigma
                                       [-2.0,-1.0,0.0,0.5], #log_obs_noise
                                       [0.25,0.5,0.75,1.0], #ec
                                       [0.1,0.5,1.0,2.0,5.0,10.0,20.0], #k_states
                                       [0.0,0.1,0.25,0.5,0.75,0.9]) #alpha_states

function generate_sim_q(study::TuneCBTS)
    q = []
    for seed in 1:study.n_seeds, n_iters in study.n_iters, A_max in study.A_max, n_proposes in study.n_proposes, log_length_scale in study.log_length_scale, log_signal_sigma in study.log_signal_sigma, log_obs_noise in study.log_obs_noise, ec in study.ecs, k_state in study.k_states, alpha_state in study.alpha_states
        push!(q, TuneCBTSSim(study.mdp, seed, n_iters, study.max_steps, A_max, n_proposes, log_length_scale, log_signal_sigma, log_obs_noise, ec, k_state, alpha_state, Dict()))
    end
    q
end
function run_study(study::TuneCBTS)
    df = DataFrame([String, Int, Int, Int, Int, Int, Float64, Float64, Float64, Float64, Float64, Float64, Float64], [:alg, :seed, :n_iters, :max_steps, :A_max, :n_proposes, :log_length_scale, :log_signal_sigma, :log_obs_noise, :ec, :k_state, :alpha_state, :r_total], 0)
    q = generate_sim_q(study)
    result = pmap(POMDPs.simulate, q)
    for (r, sim) in zip(result, q)
        push!(df, ["CBTSBandit", sim.seed, sim.n_iters, sim.max_steps, sim.A_max, sim.n_proposes, sim.log_length_scale, sim.log_signal_sigma, sim.log_obs_noise, sim.ec, sim.k_state, sim.alpha_state, r])
    end
    TuneCBTSResult(string(study.mdp), df)
end
function report(result::TuneCBTSResult)
    df = aggregate(result.data, [:alg, :max_steps, :A_max, :n_proposes, :log_length_scale, :log_signal_sigma, :log_obs_noise, :ec, :k_state, :alpha_state], [mean, std])[[:alg, :A_max, :n_proposes, :log_length_scale, :log_signal_sigma, :log_obs_noise, :ec, :k_state, :alpha_state, :r_total_mean, :r_total_std]]
    dd = sort(df[:], :r_total_mean, rev=true)
    dd
end


text:
using FractionalTransforms
using Test

@testset "Test FRFT" begin
    @test isapprox(real(frft([1,2,3], 0.5)), [0.100464613588213, 3.07460010423315, 1.31566432887283], atol=1e-5)
    @test isapprox(imag(frft([1,2,3], 0.5)), [-0.259409480977277, 0.199349381540632, -0.93645356561191], atol=1e-5)
end

@testset "Test FRFT auxillary functions" begin
    @testset "Test Sinc Interpolation" begin
        @test isapprox(real(sinc_interp([1, 2, 3], 3)), [1.0, 1.1577906803857638, 1.4472383504822042, 2.0, 2.6877283651812367, 3.1425747039042147, 3.0]; atol=1e-5)
        @test isapprox(imag(sinc_interp([1, 2, 3], 3)), [-4.0696311822644287e-17, -2.4671622769447922e-17, -2.522336560207922e-16, -3.331855648569948e-17, 0.0, 1.9624582529744518e-17, 3.331855648569948e-17]; atol=1e-5)
    end

    @testset "Test Frequency Shear" begin
        @test isapprox(real(freq_shear([1, 2, 3], 3)), [0.0707372016677029, 2.0, 0.2122116050031087]; atol=1e-5)
        @test isapprox(imag(freq_shear([1, 2, 3], 3)), [0.9974949866040544, 0.0, 2.9924849598121632]; atol=1e-5)
    end
end


text:
### A Pluto.jl notebook ###
# v0.12.17

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3c92f3e0-4561-11eb-37f1-8592dae8553f
begin
	using Plots
	using PlutoUI
	using Unitful
	using UnitfulRecipes
end

# ╔═╡ 089cf360-4561-11eb-2df0-8f3849683471
include("GasLaws.jl")

# ╔═╡ fae31c72-46f2-11eb-08b9-0b24d9a602cc
md"""
Enter $3$ of $P$, $V$, $n$, and $T$ and just the units for the remaining quantity. The remaining quantity will be calculated with the chosen units using the Ideal Gas law.

A number of units are available for each quantity and most SI prefixes can be used. To specifiy Celcius/Fahrenheit temperature units, you must include the degree symbol (e.g. °C/°F)
"""

# ╔═╡ ea7c3510-459e-11eb-19e1-0b4793eb57fb
begin
	md"""
	P = $(@bind P NumberField(0.0:100.0))  $(@bind P_unit TextField(;default="bar"))

	V= $(@bind V NumberField(0.0:100.0))  $(@bind V_unit TextField(;default="L"))

	n= $(x=@bind n NumberField(0.0:100.0)) $(@bind n_unit TextField(;default="mol"))

	T= $(@bind T NumberField(0.0:1000.0))  $(@bind T_unit TextField(;default="K"))
	"""
end

# ╔═╡ b26d3f00-45a0-11eb-358c-e10f44d8a1cc
begin
	if isnan(P)
		P_full=uparse(P_unit)
		quantity="Pressure"
	else
	    P_full=P*uparse(P_unit)
	end
	if isnan(V)
		V_full=uparse(V_unit)
		quantity="Volume"
	else
	    V_full=V*uparse(V_unit)
	end
	if isnan(n)
		n_full=uparse(n_unit)
		quantity="amount"
    else
	    n_full=n*uparse(n_unit)
	end
	if isnan(T)
		T_full=T*uparse(T_unit)
		quantity="Temperature"
	else
		T_full=T*uparse(T_unit)
	end
	try
	    result=GasLaws.ideal(P_full,V_full,n_full,T_full)
	    md"""The resulting $quantity is $result"""
	catch err
		md"""Enter $3$ values to calculate the blank value"""
	end
end

# ╔═╡ 59abc110-4731-11eb-3392-c5d3db261f1a

# ╔═╡ 99c5e560-46f4-11eb-2e31-254b08bf6a2b
md"""
Can calculate pressure for a real gas using the Van der Waals or Redlich-Kwong equations. The Slider below allows you to change the temperature for a plots of $P$ vs $V_m$ (molar volume) with each of these equations. The $a$ and $b$ constants for $\ce{CO2}$ are used for each function.
"""

# ╔═╡ 57be7b8e-4727-11eb-0c8d-aff4ab6b8835
@bind newT Slider(0.0:1000.0,show_value=true)

# ╔═╡ 86087550-4727-11eb-3b46-fd0ef0c05d93
begin
	realT=newT*1u"K"
	Vm_range=(1.0:0.25:5.0)u"L/mol"
	VdW(v)=GasLaws.Van_der_Waals(v,realT)
	RK(v)=GasLaws.Redlich_Kwong(v,realT)
	plot(VdW,Vm_range,xlabel="Molar Volume",ylabel="Pressure",label="Van der Waals")
	plot!(RK,Vm_range,xlabel="Molar Volume",ylabel="Pressure",label="Redlich-Kwong")
end

# ╔═╡ Cell order:
# ╟─089cf360-4561-11eb-2df0-8f3849683471
# ╟─3c92f3e0-4561-11eb-37f1-8592dae8553f
# ╟─fae31c72-46f2-11eb-08b9-0b24d9a602cc
# ╟─ea7c3510-459e-11eb-19e1-0b4793eb57fb
# ╟─b26d3f00-45a0-11eb-358c-e10f44d8a1cc
# ╟─59abc110-4731-11eb-3392-c5d3db261f1a
# ╟─99c5e560-46f4-11eb-2e31-254b08bf6a2b
# ╟─57be7b8e-4727-11eb-0c8d-aff4ab6b8835
# ╟─86087550-4727-11eb-3b46-fd0ef0c05d93

max_stars_repo_name: TyBalduf/IdealGas

max_issues_repo_name: TyBalduf/IdealGas

max_forks_repo_name: TyBalduf/IdealGas


text:
"""
This struct contains the capacity market parameters
for price forcasts using CEM and endogeneous Economic Dispatch.
The Capacity market is implemented using a piece-wise Capacity Demand Curve,
with break-points and price-points parameterized within this struct.
"""
struct CapacityMarket # Assumes a piece-wise linear demand curve

    break_points::Vector{Float64}   # demand curve break points
    price_points::Vector{Float64}    # $/MW-year

    function CapacityMarket(b, p)
        @assert all(b .>= 0)
        @assert all(p .>= 0)
        @assert length(b) > 1
        @assert length(p) == length(b)
        new(b, p)
    end

end


text:
########## Salsa Trace Free List

# This section defines a thread-local storage free-list for salsa derived function
# traces. This is needed because each derived function needs its own vector to track
# dependencies, and vectors have to be heap allocated, but allocations are expensive. We can
# avoid those allocations by pre-allocating a pool of trace objects, and reusing them with a
# freelist. We have one instance of pool+freelist per OS Thread, so that no locking is
# required when taking/putting trace objects into/out of the freelist.
#
# The implementation works as follows:
# On each thread we keep a pool and a freelist:
#   - The pool stores all the TraceOfDependencyKeys objects that are allocated, and
#   - The freelist stores _indexes_ into that pool vector, representing traces that are
#     free to use.
#
# When a derived function starts, it takes an index out of the freelist, which is stored in
# the instance of the _TracingRuntime. Then, while executing user code, when any subsequent
# derived functions are called (bringing us back into the Salsa codebase), we access the
# trace via the provided index, and register the newly called derived function dependency.
#
# We use an Int index rather than simply a reference to the trace itself in order to ensure
# that `_TracingRuntime` remains isbits, since we create one instance per derived function
# call, and if it is not isbits, it will incur allocations.
#
# When the derived function call finishes, we copy the dependencies out of the trace, clear
# it, and put its index back into the freelist. If we ever attempt to take out a new index,
# but the pool is empty, the pool will be doubled with new traces. This should be fairly
# rare, since it will only happen the first time we exceed some number of derived functions
# in flight.
#
# NOTE: In the future, we intend to rewrite Salsa to use a Stack, as depth-first-search,
# rather than using function calls as we do now. Based on our early sketches, we think such
# a design will allow us to remove all of this silliness since each Salsa Task will have its
# own local variables that can be freely reused, much more simply. For now, this suffices.

const g_threadlocal_trace_pools = Vector{TraceOfDependencyKeys}[]
const g_threadlocal_trace_freelists = Vector{Int}[]
const g_threadlocal_pool_locks = Base.ReentrantLock[]

# Start off with some large number of traces in the pool. This is arbitrary, but we want it
# to be big enough where it doesn't have to get doubled very often. The traces will double
# whenever we have more active derived functions than available traces, which should only
# happen for very-long chains of derived functions or very-wide task parallelism.
const N_INIT_TRACES = 1024

# This function is called in Salsa.__init__() because we don't know the
# number of threads until runtime. (__init__() is defined at the end of this file.)
function _init_thread_local_pools_and_freelists()
    append!(
        g_threadlocal_trace_pools,
        TraceOfDependencyKeys[TraceOfDependencyKeys() for _ = 1:N_INIT_TRACES] for _ = 1:Threads.nthreads()
    )
    append!(
        g_threadlocal_trace_freelists,
        Int[i for i = 1:N_INIT_TRACES] for _ = 1:Threads.nthreads()
    )
    for _ = 1:Threads.nthreads()
        push!(g_threadlocal_pool_locks, Base.ReentrantLock())
    end
    return nothing
end

# Thread-safe accessors: provides one pool+freelist per OS thread.
tls_trace_pool() = g_threadlocal_trace_pools[Threads.threadid()]
tls_trace_freelist() = g_threadlocal_trace_freelists[Threads.threadid()]
tls_trace_pool_lock() = g_threadlocal_pool_locks[Threads.threadid()]

# This is sort of terrible, but since a _TracingRuntime can be _referenced_ from a different
# thread than the one that created it, we need both a thread id and a trace id to retrieve
# its trace from the new thread. Note that this is safe because the parent thread will never
# release its trace until all children have returned, so there can be no dangling references
# to the trace.
struct TraceId
    thread_idx::Int
    trace_idx::Int
end

# We store the call_stack in the Trace instead of in the runtime, because the trace is a
# linked-list (meaning not isbits), and we want to keep the Runtime isbits to avoid
# allocations. Since the Traces are pre-allocated it's okay for them to contain heap ptrs.
@inline function get_trace_with_call_stack(call_stack::Union{SalsaStackFrame,Nothing})
    trace_id = get_free_trace_id()
    get_trace(trace_id).call_stack = call_stack
    return trace_id
end

function get_trace(id::TraceId)
    return g_threadlocal_trace_pools[id.thread_idx][id.trace_idx]
end

function get_free_trace_id()::TraceId
    trace_freelist = tls_trace_freelist()
    trace_pool = tls_trace_pool()
    if isempty(trace_freelist)
        # Even with a thread-local pool we could potentially run into race conditions here
        # if we didn't lock around this section, when multiple tasks on the same thread
        # attempt to grow the pool but block inside this function on e.g. the `@info` call.
        @lock tls_trace_pool_lock() begin

            # Because we only acquire the lock once we've noticed the empty freelist, we can
            # have multiple tasks ending up in the critical section here, one after
            # another. Of course only the first one should actually perform the doubling. We
            # therefore have to check again whether the freelist is still empty. The
            # alternative would be to lock _before_ checking if the freelist is empty, but
            # that seems like unnecessary overhead since the empty freelist should be a very
            # rare occurrence.
            if !isempty(trace_freelist)
                return TraceId(Threads.threadid(), pop!(trace_freelist))
            end

            @debug_mode @info "DOUBLING SALSA TRACE POOL ON THREAD $(Threads.threadid())"
            old_size = length(trace_pool)
            num_new_traces = old_size # Double the pool
            @debug_mode @assert num_new_traces > 0
            # Grow the pool:
            append!(
                trace_pool,
                TraceOfDependencyKeys[TraceOfDependencyKeys() for _ = 1:num_new_traces],
            )
            for i = (old_size + 1) : (old_size + num_new_traces)
                push!(trace_freelist, i)
            end
        end
    end
    return TraceId(Threads.threadid(), pop!(trace_freelist))
end

function release_trace_id(id::TraceId)
    @assert id.thread_idx === Threads.threadid()
    # NOTE: no locking is needed here because this is all happening per-thread.
    empty_trace!(tls_trace_pool()[id.trace_idx])
    push!(tls_trace_freelist(), id.trace_idx)
    return nothing
end

function empty_trace!(trace::TraceOfDependencyKeys)
    empty!(trace.ordered_deps)
    empty!(trace.seen_deps)
end


text:
function build_transcription_reaction_table(gene_name::String, sequence::BioSequences.LongSequence;
    polymeraseSymbol::Symbol=:RNAP, ecnumber::String="2.7.7.6", logger::Union{Nothing,SimpleLogger}=nothing)::Some

    try

        # initialize -
        id_array = String[]
        forward_reaction_string = String[]
        reverse_reaction_string = String[]
        reversibility_array = Bool[]
        ec_number_array = Union{Missing,String}[]
        lower_bound_array = Union{Missing,Float64}[]
        upper_bound_array = Union{Missing,Float64}[]
        polymeraseSymbol = string(polymeraseSymbol)

        # get a count on G,C,A,U -
        number_of_A = count(sequence, DNA_A)
        number_of_U = count(sequence, DNA_T)
        number_of_C = count(sequence, DNA_C)
        number_of_G = count(sequence, DNA_G)
        total_nucleotides = number_of_A + number_of_U + number_of_C + number_of_G

        # setup reactions -
        # gene + RNAP <=> gene_RNAP_closed
        push!(id_array, "$(gene_name)_binding")
        push!(forward_reaction_string, "$(gene_name)+$(polymeraseSymbol)")
        push!(reverse_reaction_string, "$(gene_name)_$(polymeraseSymbol)_closed")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # gene_RNAP_closed => gene_RNAP_open
        push!(id_array, "$(gene_name)_open")
        push!(forward_reaction_string, "$(gene_name)_$(polymeraseSymbol)_closed")
        push!(reverse_reaction_string, "$(gene_name)_$(polymeraseSymbol)_open")
        push!(reversibility_array, false)
        push!(ec_number_array, missing)
        push!(lower_bound_array, 0.0)
        push!(upper_bound_array, Inf)

        # transcription -
        push!(id_array, "$(gene_name)_transcription")
        push!(forward_reaction_string, "$(gene_name)_$(polymeraseSymbol)_open+$(number_of_A)*M_atp_c+$(number_of_U)*M_utp_c+$(number_of_C)*M_ctp_c+$(number_of_G)*M_gtp_c+$(total_nucleotides)*M_h2o_c")
        push!(reverse_reaction_string, "mRNA_$(gene_name)+$(gene_name)+$(polymeraseSymbol)+$(total_nucleotides)*M_ppi_c")
        push!(reversibility_array, false)
        push!(ec_number_array, ecnumber)
        push!(lower_bound_array, 0.0)
        push!(upper_bound_array, Inf)

        # mRNA degradation -
        push!(id_array, "mRNA_$(gene_name)_degradation")
        push!(forward_reaction_string, "mRNA_$(gene_name)")
        push!(reverse_reaction_string, "$(number_of_A)*M_amp_c+$(number_of_U)*M_ump_c+$(number_of_C)*M_cmp_c+$(number_of_G)*M_gmp_c")
        push!(reversibility_array, false)
        push!(ec_number_array, missing)
        push!(lower_bound_array, 0.0)
        push!(upper_bound_array, Inf)

        # package into DataFrame -
        reaction_dataframe = DataFrame(id=id_array, forward_reaction=forward_reaction_string, reverse_reaction=reverse_reaction_string, reversibility=reversibility_array,
            ec_number=ec_number_array, lower_bound=lower_bound_array, upper_bound=upper_bound_array)

        # return -
        return Some(reaction_dataframe)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function build_translation_reaction_table(protein_name::String, sequence::BioSequences.LongAminoAcidSeq;
    ribosomeSymbol::Symbol=:RIBOSOME, ecnumber::String="3.1.27.10", logger::Union{Nothing,SimpleLogger}=nothing)::Some

    try

        # initailize -
        id_array = String[]
        forward_reaction_string = String[]
        reverse_reaction_string = String[]
        reversibility_array = Bool[]
        ec_number_array = Union{Missing,String}[]
        lower_bound_array = Union{Missing,Float64}[]
        upper_bound_array = Union{Missing,Float64}[]
        ribosomeSymbol = string(ribosomeSymbol)
        total_residue_count = 0
        protein_aa_map = Dict{BioSymbol,Int64}()
        aa_biosymbol_array = [
            AA_A, AA_R, AA_N, AA_D, AA_C, AA_Q, AA_E, AA_G, AA_H, AA_I, AA_L, AA_K, AA_M, AA_F, AA_P, AA_S, AA_T, AA_W, AA_Y, AA_V
        ];

        # load AA map -
        aa_metabolite_map = TOML.parsefile(joinpath(_PATH_TO_CONFIG, "AAMap.toml"))

        # build the protein_aa_map -
        for residue in aa_biosymbol_array
            protein_aa_map[residue] = count(sequence, residue)
        end

        # total residue count -
        for residue in aa_biosymbol_array
            number_per_AA = protein_aa_map[residue]
            total_residue_count = total_residue_count + number_per_AA
        end

        # setup reactions -
        # mRNA + RIBOSOME <=> mRNA_RIBOSOME_closed
        push!(id_array, "$(protein_name)_binding")
        push!(forward_reaction_string, "mRNA_$(protein_name)+$(ribosomeSymbol)")
        push!(reverse_reaction_string, "mRNA_$(protein_name)_$(ribosomeSymbol)_closed")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # mRNA_RIBOSOME_closed => mRNA_RIBOSOME_start
        push!(id_array, "mRNA_$(protein_name)_open")
        push!(forward_reaction_string, "mRNA_$(protein_name)_$(ribosomeSymbol)_closed")
        push!(reverse_reaction_string, "mRNA_$(protein_name)_$(ribosomeSymbol)_start")
        push!(reversibility_array, false)
        push!(ec_number_array, missing)
        push!(lower_bound_array, 0.0)
        push!(upper_bound_array, Inf)

        # mRNA_RIBOSOME_translation -
        push!(id_array, "mRNA_$(protein_name)_translation")
        push!(forward_reaction_string, "mRNA_$(protein_name)_$(ribosomeSymbol)_start+$(2 * total_residue_count)*M_gtp_c+$(2 * total_residue_count)*M_h2o_c")
        push!(reverse_reaction_string, "mRNA_$(protein_name)+$(ribosomeSymbol)+P_$(protein_name)+$(2 * total_residue_count)*M_gdp_c+$(2 * total_residue_count)*M_pi_c+$(total_residue_count)*tRNA_c")
        push!(reversibility_array, false)
        push!(ec_number_array, missing)
        push!(lower_bound_array, 0.0)
        push!(upper_bound_array, Inf)

        # charge the tRNA -
        for residue in aa_biosymbol_array

            # get the key symbol -
            key_value = "AA_" * (string(residue))

            # get the number and M_* of this residue -
            number_of_AA_residue = protein_aa_map[residue]
            metabolite_symbol = aa_metabolite_map[key_value]

            # build the reaction record -
            push!(id_array, "tRNA_charging_$(metabolite_symbol)_$(protein_name)")
            push!(forward_reaction_string, "$(number_of_AA_residue)*$(metabolite_symbol)+$(number_of_AA_residue)*M_atp_c+$(number_of_AA_residue)*tRNA_c+$(number_of_AA_residue)*M_h2o_c")
            push!(reverse_reaction_string, "$(number_of_AA_residue)*$(metabolite_symbol)_tRNA_c+$(number_of_AA_residue)*M_amp_c+$(number_of_AA_residue)*M_ppi_c")
            push!(reversibility_array, false)
            push!(ec_number_array, missing)
            push!(lower_bound_array, 0.0)
            push!(upper_bound_array, Inf)
        end

        # package into DataFrame -
        reaction_dataframe = DataFrame(id=id_array, forward_reaction=forward_reaction_string, reverse_reaction=reverse_reaction_string, reversibility=reversibility_array,
            ec_number=ec_number_array, lower_bound=lower_bound_array, upper_bound=upper_bound_array)

        # return -
        return Some(reaction_dataframe)

    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function build_txtl_transport_reaction_table()::Some

    try

        # initailize -
        id_array = String[]
        forward_reaction_string = String[]
        reverse_reaction_string = String[]
        reversibility_array = Bool[]
        ec_number_array = Union{Missing,String}[]
        lower_bound_array = Union{Missing,Float64}[]
        upper_bound_array = Union{Missing,Float64}[]
        aa_biosymbol_array = [
            AA_A, AA_R, AA_N, AA_D, AA_C, AA_Q, AA_E, AA_G, AA_H, AA_I, AA_L, AA_K, AA_M, AA_F, AA_P, AA_S, AA_T, AA_W, AA_Y, AA_V
        ];

        # load AA map -
        aa_metabolite_map = TOML.parsefile(joinpath(_PATH_TO_CONFIG, "AAMap.toml"))

        # add exchange water reactions -
        push!(id_array, "M_h2o_c_exchange")
        push!(forward_reaction_string, "M_h2o_e")
        push!(reverse_reaction_string, "M_h2o_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_ppi_e exchange -
        push!(id_array, "M_ppi_c_exchange")
        push!(forward_reaction_string, "M_ppi_e")
        push!(reverse_reaction_string, "M_ppi_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_amp_c exchange -
        push!(id_array, "M_amp_c_exchange")
        push!(forward_reaction_string, "M_amp_e")
        push!(reverse_reaction_string, "M_amp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_gmp_c exchange -
        push!(id_array, "M_gmp_c_exchange")
        push!(forward_reaction_string, "M_gmp_e")
        push!(reverse_reaction_string, "M_gmp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_cmp_c exchange -
        push!(id_array, "M_cmp_c_exchange")
        push!(forward_reaction_string, "M_cmp_e")
        push!(reverse_reaction_string, "M_cmp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_ump_c exchange -
        push!(id_array, "M_ump_c_exchange")
        push!(forward_reaction_string, "M_ump_e")
        push!(reverse_reaction_string, "M_ump_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_atp_c exchange -
        push!(id_array, "M_atp_c_exchange")
        push!(forward_reaction_string, "M_atp_e")
        push!(reverse_reaction_string, "M_atp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_gtp_c exchange -
        push!(id_array, "M_gtp_c_exchange")
        push!(forward_reaction_string, "M_gtp_e")
        push!(reverse_reaction_string, "M_gtp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_ctp_c exchange -
        push!(id_array, "M_ctp_c_exchange")
        push!(forward_reaction_string, "M_ctp_e")
        push!(reverse_reaction_string, "M_ctp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # add M_utp_c exchange -
        push!(id_array, "M_utp_c_exchange")
        push!(forward_reaction_string, "M_utp_e")
        push!(reverse_reaction_string, "M_utp_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        push!(id_array, "tRNA_c_exchange")
        push!(forward_reaction_string, "tRNA_e")
        push!(reverse_reaction_string, "tRNA_c")
        push!(reversibility_array, true)
        push!(ec_number_array, missing)
        push!(lower_bound_array, -Inf)
        push!(upper_bound_array, Inf)

        # transfer AAs -
        for residue in aa_biosymbol_array

            # get the key symbol -
            key_value = "AA_" * (string(residue))

            # metabilite -
            metabolite_symbol_c = aa_metabolite_map[key_value]
            metabolite_symbol_e = replace(metabolite_symbol_c, "_c" => "_e")

            # write the record -
            push!(id_array, "$(metabolite_symbol_c)_exchange")
            push!(forward_reaction_string, metabolite_symbol_e)
            push!(reverse_reaction_string, metabolite_symbol_c)
            push!(reversibility_array, true)
            push!(ec_number_array, missing)
            push!(lower_bound_array, -Inf)
            push!(upper_bound_array, Inf)
        end

        # package into DataFrame -
        reaction_dataframe = DataFrame(id=id_array, forward_reaction=forward_reaction_string, reverse_reaction=reverse_reaction_string, reversibility=reversibility_array,
            ec_number=ec_number_array, lower_bound=lower_bound_array, upper_bound=upper_bound_array)

        # return -
        return Some(reaction_dataframe)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function transcribe(sequence::BioSequences.LongSequence; complementOperation::Function=!,
    logger::Union{Nothing,SimpleLogger}=nothing)::Some

    try

        # initialize -
        new_sequence_array = Array{BioSymbol,1}()

        # ok: let's iterate the sequence, call the complementOperation function for each nt -
        for nt in sequence

            # call the complementOperation function to get the complementary nucleotide -
            complementary_nt = complementOperation(nt)

            # push -
            push!(new_sequence_array, complementary_nt)
        end

        # create new sequence -
        new_seq_from_array = LongSequence{RNAAlphabet{4}}(new_sequence_array)

        # return -
        return Some(new_seq_from_array)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function transcribe(table::DataFrame, complementOperation::Function=!;
    logger::Union{Nothing,SimpleLogger}=nothing)

    try

        # initailize -
        transcription_dictionary = Dict{String,Any}()

        # get the length of the table -
        (number_of_rows, _) = size(table)
        for row_index = 1:number_of_rows

            # get id -
            sequence_id = table[row_index, :id]
            sequence = table[row_index, :sequence]

            # transcribe -
            result = transcribe_sequence(sequence, complementOperation; logger=logger) |> check

            # package -
            transcription_dictionary[sequence_id] = result
        end

        # return -
        return Some(transcription_dictionary)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function translate(sequence::BioSequences.LongAminoAcidSeq, complementOperation::Function;
    logger::Union{Nothing,SimpleLogger}=nothing)::Some

    try
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function translate(table::DataFrame, complementOperation::Function;
    logger::Union{Nothing,SimpleLogger}=nothing)::Some

    try

        # get the size of the table -
        (number_of_proteins, _) = size(table)
        for protein_index in number_of_proteins

            # get the sequence -
            protein_seq = table[protein_index, :protein_sequence]
        end

    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end


text:
export DegreeIndependentSet

struct DegreeIndependentSet end

"""
    independent_set(g, DegreeIndependentSet())

Obtain an [independent set](https://en.wikipedia.org/wiki/Independent_set_(graph_theory)) of
`g` using a greedy heuristic based on the degree of the vertices.

### Implementation Notes
A vertex is said to be valid if it is not in the independent set or adjacent to any vertex
in the independent set.
Initilalise the independent set to an empty set and iteratively choose the vertex that is
adjacent to the fewest valid vertices in the independent set until all vertices are invalid.

### Performance
Runtime: O((|V|+|E|)*log(|V|))
Memory: O(|V|)
"""
function independent_set(
    g::AbstractGraph{T},
    alg::DegreeIndependentSet
    ) where T <: Integer

    nvg = nv(g)
    ind_set = Vector{T}()
    sizehint!(ind_set, nvg)
    deleted = falses(nvg)
    degree_queue = PriorityQueue(enumerate(degree(g)))

    while !isempty(degree_queue)
        v = dequeue!(degree_queue)
        deleted[v] && continue
        deleted[v] = true
        push!(ind_set, v)

        for u in neighbors(g, v)
            deleted[u] && continue
            deleted[u] = true
            @inbounds @simd for w in neighbors(g, u)
                if !deleted[w]
                    degree_queue[w] -= 1
                end
            end
        end
    end
    return ind_set
end


text:
@testset "EqualTo Constraint" begin
    @testset "digits form a number (Issue 200)" begin
        use_diff = false
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "traverse_strategy" => :DBFS,
            "branch_split" => :InHalf,
            "logging" => [],
        ))

        @variable(model, 0 <= x[1:10] <= 9, Int)
        @variable(model, 10000 <= v[1:2] <= 99999, Int) # ABCDE and FGHIJ
        if use_diff
            @variable(model, 0 <= diff <= 999, Int) # adding this is much slower than using v[1:2] only
        end

        a, b, c, d, e, f, g, h, i, j = x
        @constraint(model, x in CS.AllDifferent())
        @constraint(model, v[1] == 10000 * a + 1000 * b + 100 * c + 10 * d + e) # ABCDE
        @constraint(model, v[2] == 10000 * f + 1000 * g + 100 * h + 10 * i + j) # FGHIJ

        # Using diff is slower
        if use_diff
            @constraint(model, diff == v[1] - v[2])  # much slower
            @constraint(model, diff >= 1)
        else
            @constraint(model, v[1] - v[2] >= 1)
        end

        if use_diff
            @objective(model, Min, diff) # slower
        else
            @objective(model, Min, v[1] - v[2])
        end

        # Solve the problem
        optimize!(model)

        status = JuMP.termination_status(model)
        @test status == MOI.OPTIMAL
        x_vals = convert.(Int, JuMP.value.(x))
        @test x_vals == [5, 0, 1, 2, 3, 4, 9, 8, 7, 6]
    end
end


text:
module MASW

using FFTW, Plots

export masw

"""
    masw(U::AbstractArray, x::AbstractArray, t::AbstractArray, ct::AbstractArray;
               plotfreqlimits::Tuple=(),
               plotwaterlevel::Real=0.0,
               figdir::String="",
               figname::String="")

Return dispersion curve inferred from Multichannel Analysis of Surface Waves (MASW)

# Arguments
- `U::AbstractArray`: 2D Array for seismci traces: size(U)=(T:length of timeseries, N:number of traces)
- `x::AbstractArray`: Vector for offsets (distance between source and receiver)
- `t::AbstractArray`: Vector for timeseries (ex. 0:1/fs:Tmax)
- `ct::AbstractArray`: Vector for scanning velocity of interest.
- `plotfreqlimits::Tuple=()`: Tuple for minfreq and maxfreq of plotting (ex. (0.0, 10.0))
- `plotwaterlevel::Real=0.0`: Dispersion Image waterlevel expressed as a percent for display purposes.
- `figdir::String=""`: Figure directory name if you want to plot the result figure.
- `figname::String=""`: figname if you want to plot the result figure (please include extention.).

# Usage

`(f, ct, Udisp, DispersionCurve) = masw(U, x, t, ct, plotfreqlimits, plotwaterlevel, figdir, figname)`

# Output

- `f`: Vector for frequency associated with dispersion image.
- `ct`: Vector for scanning velocity associated with dispersion image.
- `Udisp`: Normalized 2D array for dispersion image.
- `DispersionCurve`: DispersionCurve including frequency and automatically picked phase velocity.

# Reference
- Park, C.B., Miller, R.D. and Xia, J. (2001) Offset and resolution of dispersion curve in multichannel analysis of surface waves (MASW). doi:10.4133/1.2922953

Julia version written by kura-okubo (https://github.com/kura-okubo)

This is the translation in Julia based on matlab script written by John Schuh.

Matlab script: Copyright (c) 2018, John Schuh
All rights reserved.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

- John Schuh (2020). MASW Dispersion Curve (https://www.mathworks.com/matlabcentral/fileexchange/65138-masw-dispersion-curve), MATLAB Central File Exchange. Retrieved January 14, 2020.
"""
function masw(U::AbstractArray, x::AbstractArray, t::AbstractArray, ct::AbstractArray;
                plotfreqlimits::Tuple=(),
                plotwaterlevel::Real=0.0,
                figdir::String="",
                figname::String="")

    # precheck of array size
    if size(U) != (length(t), length(x))
        error("size of U is not correct. Please check that size(U) == (length(t), length(x))")
    end

    datanorm = U./maximum(U)

    L=length(t); # number os samples
    T=maximum(t); # max time
    if mod(L,2) == 0
        f = (1/T)*(0:L/2-1); # n is even
    else
        f = (1/T)*(0:(L-1)/2); # n is odd
    end

    Ufft=FFTW.rfft(U, 1) # fft is performed on each column (the time vector)
    w=f*2*pi; # convert the freq vector to rads/sec

    ## 2. Normalize the spectrum
    Rnorm=Ufft./abs.(Ufft); # Normalize the spectrum

    ## 3. Perfrom the summation over N in the frequency domain
    AsSum = zeros(length(w),length(ct))
    @simd for ii=1:length(w) # send in frequencies one at a time
        As=zeros(Complex, length(x),length(ct)); #initialize As or overwite the previous As
        for n=1:length(x) # send in a positon (x) one at a time
            # Iterate As over position
            As[n,:]=exp.(1im*w[ii].*(x[n]./ct)).*Rnorm[ii,n]
        end
         AsSum[ii,:]=abs.(sum(As, dims=1));
    end
    AsSum=AsSum'; # transpose the matrix so velocity is on vertical and freq is on horizontal
    normed=AsSum./maximum(AsSum, dims=1);# normalize the dispersion image so max column values=1

    ############ curve autopicking ###################
    # This section will auto pick the dispersion curve based on the peak
    # value in each column. This will need updating when
    # higher modes are introduced
    remaindat=deepcopy(normed)
    remaindat[remaindat .< 1] .= 0
    ind=findall(x -> x==1, remaindat)
    row = []; col = []
    for i = 1:length(ind)
        push!(row, ind[i][1]); push!(col, ind[i][2])
    end

    autovel=ct[row]; # autopicked velocity indices that contain the peak value
    autofreq=f[col]; # autopicked frequency indices that contain the peak value
    # get rid of all the extra picks at zero frequency
    nz=length(autofreq)-count(x -> x != 0, autofreq); # determine the number of repeated picks at zero
    autofreq=autofreq[nz:end]; # remove the repeated zero picks
    autovel=autovel[nz:end]; # remove the repeated zero picks
    DispersionVelocity=autovel; # need to get the rlowess working again, I no longer have access to the toolbox that allows it
    normed[normed .< plotwaterlevel/100] .= NaN; # throw out data less than plotwaterlevel
    ################ end autopicking ######################

    DispersionCurve=[f, DispersionVelocity]; # this only has meaning if the curve was picked using 'manual' or 'auto'

    if !isempty(figdir)

        if Plots.backend() != Plots.PyPlotBackend()
            @warn ("MASW.jl plotting is available only with pyplot backend. Please type
                    `using Plots;Plots.pyplot()` before call this function. Return result
                     without plotting figure.")
            return (f, ct, normed, DispersionCurve)
        end

        # plot input traces
        l = Plots.@layout [a b]
        p1 = Plots.contourf(x, t, U, levels=200,  yflip = true, xlabel="Distance", ylabel="Time",
                            linewidth=0.0, c=:balance,  title="Traces")
        # Plot the dispersion results
        p2 = Plots.contourf(f,ct,normed, levels=200, xlabel="Freq (Hz)", ylabel="Phase Velocity",
                            clims=(0.0,1.0), linewidth=0.0, c=:thermal, title="Dispersion Image") # plot the normalized dispersion image
        Plots.plot!(DispersionCurve[1],DispersionCurve[2], line=(:black, 2.0), label="dispersion curve") # Plot the picked dispersion curve. This is nothing if 'none' picking was set
        if !isempty(plotfreqlimits)
            Plots.xlims!(plotfreqlimits)
        end
        Plots.scalefontsizes(3)
        p = Plots.plot(p1, p2, size=(1600, 800), layout=l)
        if !ispath(figdir)
            mkpath(figdir)
        end
        Plots.savefig(p, figdir*"/"*figname)
    end

    return (f, ct, normed, DispersionCurve)
end

end # module


text:
########################################################################################
# Plotting
########################################################################################
@recipe function f(
    sol::AbstractProbODESolution;
    vars=nothing,
    denseplot=sol.dense,
    plotdensity=1000,
    tspan=nothing,
    ribbon_width=1.96,
)
    tstart, tend = isnothing(tspan) ? (sol.t[1], sol.t[end]) : tspan
    times = denseplot ? range(tstart, tend, length=plotdensity) : sol.t
    sol_rvs = denseplot ? sol(times).u : sol.pu
    if !isnothing(tspan)
        sol_rvs = sol_rvs[tstart.<=times.<=tend]
        times = times[tstart.<=times.<=tend]
    end
    values = stack(mean(sol_rvs))
    stds = stack(std(sol_rvs))

    if isnothing(vars)
        ribbon --> ribbon_width * stds
        xguide --> "t"
        yguide --> "u(t)"
        label --> hcat(["u$(i)(t)" for i in 1:length(sol.u[1])]...)
        return times, values
    else
        int_vars = interpret_vars(vars, sol, getsyms(sol))
        varsizes = unique(length.(int_vars))
        @assert length(varsizes) == 1 "`vars` argument has an errors"
        ndims = varsizes[1] - 1  # First argument is not about dimensions
        if ndims == 2
            _x = []
            _y = []
            _labels = []
            for (_, i, j) in int_vars
                push!(_x, i == 0 ? times : values[:, i])
                push!(_y, j == 0 ? times : values[:, j])
            end
            return _x, _y
        elseif ndims == 3
            _x = []
            _y = []
            _z = []
            for (_, i, j, k) in int_vars
                push!(_x, i == 0 ? times : values[:, i])
                push!(_y, j == 0 ? times : values[:, j])
                push!(_z, k == 0 ? times : values[:, k])
            end
            return _x, _y, _z
        else
            error("Error with `vars` argument")
        end
    end
end