text:
using LinearAlgebra
"""
ExKIObj{FT<:AbstractFloat, IT<:Int}
Extended Kalman Inversion struct (ExKI)
For solving the inverse problem
    y = G(θ) + η

"""

mutable struct ExKIObj{FT<:AbstractFloat, IT<:Int}
     method::String
    "a vector of arrays of size  N_parameters containing the mean of the parameters (in each exki iteration a new array of mean is added)"
     θ_mean::Vector{Array{FT, 1}}
     "a vector of arrays of size (N_parameters x N_parameters) containing the covariance of the parameters (in each exki iteration a new array of cov is added)"
     θθ_cov::Vector{Array{FT, 2}}
     "a vector of arrays of size N_y containing the predicted observation (in each exki iteration a new array of predicted observation is added)"
     y_pred::Vector{Array{FT, 1}}
     "vector of observations (length: N_y)"
     y::Array{FT, 1}
     "covariance of the observational noise"
     Σ_η::Array{FT, 2}
     "size of θ"
     N_θ::IT
     "size of y"
     N_y::IT
     "prior mean"
     r_0::Array{FT, 1}
     "prior covariance"
     Σ_0::Array{FT, 2}
     "hyperparameter"
     γ::FT
     "current iteration number"
     iter::IT
     N_iter::IT
end

"""
ExKIObj Constructor
N_iter: number of iterations
y::Array{FT,1} : observation
Σ_η::Array{FT, 2} : observation error covariance
γ::FT : hyperparameter
r_0::Array{FT, 1} : prior mean
Σ_0::Array{FT, 2} : prior covariance
θ0_mean::Array{FT, 1} : initial mean
θθ0_cov::Array{FT, 2} : initial covariance
"""
function ExKIObj(method::String,
                 N_iter::IT,
                 y::Array{FT,1},
                 Σ_η::Array{FT, 2},
                 γ::FT,
                 r_0::Array{FT, 1},
                 Σ_0::Array{FT, 2},
                 θ0_mean::Array{FT,1} = r_0,
                 θθ0_cov::Array{FT, 2} = Σ_0) where {FT<:AbstractFloat, IT<:Int}

    N_θ = size(θ0_mean,1)
    N_y = size(y, 1)

    θ_mean = Array{FT,1}[]  # array of Array{FT, 2}'s
    push!(θ_mean, θ0_mean) # insert parameters at end of array (in this case just 1st entry)
    θθ_cov = Array{FT,2}[] # array of Array{FT, 2}'s
    push!(θθ_cov, θθ0_cov) # insert parameters at end of array (in this case just 1st entry)

    y_pred = Array{FT, 1}[]  # array of Array{FT, 2}'s

    iter = IT(0)

    ExKIObj{FT,IT}(method,
                  θ_mean, θθ_cov, y_pred,
                  y,   Σ_η,
                  N_θ, N_y,
                  r_0, Σ_0,
                  γ, iter, N_iter)

end

"""
update exki struct
forward(θ) = G(θ) dG(θ)
"""
function update_ensemble!(exki::ExKIObj{FT, IT}, forward::Function) where {FT<:AbstractFloat, IT<:Int}

    exki.iter += 1

    θ_mean  = exki.θ_mean[end]
    θθ_cov = exki.θθ_cov[end]
    y = exki.y
    γ = exki.γ
    r_0, Σ_0 = exki.r_0, exki.Σ_0
    N_θ, N_y = exki.N_θ, exki.N_y
    ############# Prediction step:

    if exki.method == "ExKI"
        θ_p_mean  = θ_mean
        θθ_p_cov =  (γ + 1)*θθ_cov

        g_mean, dg = forward(θ_p_mean)
        θg_cov = θθ_p_cov * dg'

        # extended system
        ff_cov = [ dg*θg_cov+(1 + 1/γ)*Σ_η  θg_cov';
                    θg_cov     θθ_p_cov+(1 + 1/γ)*Σ_0]

        θf_cov = [θg_cov   θθ_p_cov]

        tmp = θf_cov/ff_cov

        θ_mean =  θ_p_mean + tmp*([y ; r_0] - [g_mean; θ_p_mean])

        θθ_cov =  θθ_p_cov - tmp*θf_cov'
    elseif exki.method == "ExKF"
        Δt = 1.0/exki.N_iter

        θ_p_mean  = θ_mean
        θθ_p_cov =   θθ_cov

        g_mean, dg = forward(θ_p_mean)
        θg_cov = θθ_p_cov * dg'

        # extended system

        gg_cov = dg*θg_cov + Σ_η/Δt

        tmp = θg_cov/gg_cov

        θ_mean =  θ_p_mean + tmp*(y - g_mean)

        θθ_cov =  θθ_p_cov - tmp*θg_cov'
    end

    ########### Save resutls
    push!(exki.y_pred, g_mean) # N_ens x N_data
    push!(exki.θ_mean, θ_mean) # N_ens x N_params
    push!(exki.θθ_cov, θθ_cov) # N_ens x N_data
end

function ExKI_Run(s_param, forward::Function,
    method,
    y, Σ_η,
    r_0,
    Σ_0,
    θ0_mean, θθ0_cov,
    N_iter,
    γ = 1.0)

    exkiobj = ExKIObj(
        method,
        N_iter,
        y,
        Σ_η,
        γ,
        r_0,
        Σ_0,
        θ0_mean,
        θθ0_cov)

    func(θ) = forward(s_param, θ)

    for i in 1:N_iter

        update_ensemble!(exkiobj, func)

        @info "optimization error at iter $(i) = ", 0.5*(exkiobj.y_pred[i] - exkiobj.y)'*(exkiobj.Σ_η\(exkiobj.y_pred[i] - exkiobj.y))
        @info "Frobenius norm of the covariance at iter $(i) = ", norm(exkiobj.θθ_cov[i])
    end

    return exkiobj

end

function plot_param_iter(exkiobj::ExKIObj{FT, IT}, θ_ref::Array{FT,1}, θ_ref_names::Array{String}) where {FT<:AbstractFloat, IT<:Int}

    θ_mean = exkiobj.θ_mean
    θθ_cov = exkiobj.θθ_cov

    N_iter = length(θ_mean) - 1
    ites = Array(LinRange(1, N_iter+1, N_iter+1))

    θ_mean_arr = abs.(hcat(θ_mean...))

    N_θ = length(θ_ref)
    θθ_std_arr = zeros(Float64, (N_θ, N_iter+1))
    for i = 1:N_iter+1
        for j = 1:N_θ
            θθ_std_arr[j, i] = sqrt(θθ_cov[i][j,j])
        end
    end

    for i = 1:N_θ
        errorbar(ites, θ_mean_arr[i,:], yerr=3.0*θθ_std_arr[i,:], fmt="--o",fillstyle="none", label=θ_ref_names[i])
        plot(ites, fill(θ_ref[i], N_iter+1), "--", color="gray")
    end

    xlabel("Iterations")
    legend()
    tight_layout()
end

function plot_opt_errors(exkiobj::ExKIObj{FT, IT},
    θ_ref::Union{Array{FT,1}, Nothing} = nothing,
    transform_func::Union{Function, Nothing} = nothing) where {FT<:AbstractFloat, IT<:Int}

    θ_mean = exkiobj.θ_mean
    θθ_cov = exkiobj.θθ_cov
    y_pred = exkiobj.y_pred
    Σ_η = exkiobj.Σ_η
    y = exkiobj.y

    N_iter = length(θ_mean) - 1

    ites = Array(LinRange(1, N_iter, N_iter))
    N_subfigs = (θ_ref === nothing ? 2 : 3)

    errors = zeros(Float64, N_subfigs, N_iter)
    fig, ax = PyPlot.subplots(ncols=N_subfigs, figsize=(N_subfigs*6,6))
    for i = 1:N_iter
        errors[N_subfigs - 1, i] = 0.5*(y - y_pred[i])'*(Σ_η\(y - y_pred[i]))
        errors[N_subfigs, i]     = norm(θθ_cov[i])

        if N_subfigs == 3
            errors[1, i] = norm(θ_ref - (transform_func === nothing ? θ_mean[i] : transform_func(θ_mean[i])))/norm(θ_ref)
        end

    end

    markevery = max(div(N_iter, 10), 1)
    ax[N_subfigs - 1].plot(ites, errors[N_subfigs - 1, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
    ax[N_subfigs - 1].set_xlabel("Iterations")
    ax[N_subfigs - 1].set_ylabel("Optimization error")
    ax[N_subfigs - 1].grid()

    ax[N_subfigs].plot(ites, errors[N_subfigs, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
    ax[N_subfigs].set_xlabel("Iterations")
    ax[N_subfigs].set_ylabel("Frobenius norm of the covariance")
    ax[N_subfigs].grid()
    if N_subfigs == 3
        ax[1].set_xlabel("Iterations")
        ax[1].plot(ites, errors[1, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
        ax[1].set_ylabel("L₂ norm error")
        ax[1].grid()
    end

end

##################

# function Two_Param_Linear_Test(problem_type::String, θ0_bar, θθ0_cov)

#     N_θ = length(θ0_bar)

#     if problem_type == "under-determined"
#         # under-determined case
#         θ_ref = [0.6, 1.2]
#         G = [1.0 2.0;]

#         y = [3.0;]
#         Σ_η = Array(Diagonal(fill(0.1^2, size(y))))

#     elseif problem_type == "over-determined"
#         # over-determined case
#         θ_ref = [1/3, 8.5/6]
#         G = [1.0 2.0; 3.0 4.0; 5.0 6.0]

#         y = [3.0;7.0;10.0]
#         Σ_η = Array(Diagonal(fill(0.1^2, size(y))))

#     elseif problem_type == "Hilbert"

#         G = zeros(N_θ, N_θ)
#         for i = 1:N_θ
#             for j = 1:N_θ
#                 G[i,j] = 1/(i + j - 1)
#             end
#         end

#         θ_ref = fill(1.0, N_θ)
#         y   = G*θ_ref
#         Σ_η = Array(Diagonal(fill(0.5^2, N_θ)))

#     else
#         error("Problem type : ", problem_type, " has not implemented!")
#     end

#     Σ_post = inv(G'*(Σ_η\G) + inv(θθ0_cov))
#     θ_post = θ0_bar + Σ_post*(G'*(Σ_η\(y - G*θ0_bar)))

#     return θ_post, Σ_post, G, y, Σ_η
# end

# struct Setup_Param{MAT, IT<:Int}
#     G::MAT
#     N_θ::IT
#     N_y::IT
# end

# function forward(s_param::Setup_Param, θ::Array{FT, 1}) where {FT<:AbstractFloat}
#     G = s_param.G
#     return G * θ, G
# end

# N_iter = 30
# N_θ = 2

# r_0, Σ_0 = zeros(Float64, N_θ), Array(Diagonal(fill(1.0^2, N_θ)))
# θ0_mean, θθ0_cov = r_0, Σ_0

# problem_type = "under-determined"

# θ_post, Σ_post, G, y, Σ_η = Two_Param_Linear_Test(problem_type, r_0, Σ_0)

# N_y = length(y)

# s_param = Setup_Param(G, N_θ, N_y)

# method = "ExKI"
# # UKI-1
# exki_obj = ExKI_Run(s_param, forward,
# method,
# y, Σ_η,
# r_0,
# Σ_0,
# θ0_mean, θθ0_cov,
# N_iter,)

# @info "ERRORs are :" , norm(exki_obj.θ_mean[end] - θ_post), norm(exki_obj.θθ_cov[end] - Σ_post)


text:
module ScPoExample
export hello, domath

hello(who::String) = "Hello, $who"
domath(x::Number) = (x + 5)
end

ext: jl


text:
"""
    calculateUtility(weights::Vector{Float64}, objs::Vector{Obj})

    returns the utility "profit/weight" for both the objectives, for the variables
"""
function calculateUtility(weights::Vector{Float64}, objs::Vector{Obj})
    utility1 = zeros(length(weights))
    utility2 = zeros(length(weights))

    for iter = 1:length(weights)
        utility1[iter] = objs[1].profits[iter]/weights[iter]
        utility2[iter] = objs[2].profits[iter]/weights[iter]
    end

    return(utility1, utility2)
end

"""
    permOrder(prob::Problem, order::String = "random")

    returns the permutation vector to order variables according to the given order. A reverse permutation vector is returned too.
"""
function permOrder(prob::Problem, order::String = "random"; verbose = false)

    @assert prob.nbObj == 2 "This problem is 20KP"
    @assert order in ("averageUtility", "bestMin", "rankSum", "random") "Unknown criteria for order the variables"

    u1, u2 = calculateUtility(prob.constraint.weights, prob.objs)

    if order == "bestMin"

        worst = zeros(prob.nbVar)

        for iter in 1:prob.nbVar
            worst[iter] = min(u1[iter], u2[iter])
        end
        verbose && println(worst)

        perm = sortperm(worst, rev=true)
        verbose && println(perm)

    elseif order == "averageUtility"

        avg = zeros(prob.nbVar)
        for iter = 1:prob.nbVar
            avg[iter] = (u1[iter] + u2[iter])/2
        end

        perm = sortperm(avg, rev=true)

    elseif order == "rankSum"

        rank1 = sortperm(u1)
        rank2 = sortperm(u2)

        meanRank = zeros(prob.nbVar)

        for iter in 1:prob.nbVar
            meanRank[iter] = (rank1[iter] + rank2[iter])
        end

        perm = sortperm(meanRank)

    else #random
        perm = Random.shuffle(1:prob.nbVar)
    end

    revPerm = Vector{Int}(undef, prob.nbVar)
    for iter = 1:prob.nbVar
        revPerm[perm[iter]] = iter
    end

    return(perm, revPerm)
end


text:
# Optimization of bin counts for histograms, heatmaps, hexbin plots, etc.
#
# I'm using the penalized maximum-likelihood method proposed in
#   Birge, L, and Rozenholc, Y. (2006) How many bins should be put in a regular
#   histogram?
#
# There has been quite a bit written on this problem, but there are a number of
# methods that all seem to give good results with little difference. Birge's
# method is simple (it's just AIC with an extra logarithmic term), has a decent
# theoretical justification, and is general enough to apply to multidimensional
# and non-regular bin selecetion problems. Though, the penalty they use was
# optimized for regular histograms, so may need to be tweaked.
#
# The Birge penalty is
#    penalty(D) = D - 1 + log(D)^2.5
# where D is the number of bins. The 2.5 constant was arrived at emperically by
# optimizing over samples from example density functions.
#

# Penalized log-likelihood function for a histogram with d regular bins.
#
# Args:
#   d: Number of bins in the histogram.
#   n: Number of sample (which should equal sum(bincounts[1:d])).
#   bincounts: An array giving the number occurrences in each bin.
#   binwidth: Width of each bin in the histogram.
#
# Returns:
#   Log-likelihood with Birge's penalty applied.
#
function bincount_pll(d::Int, n::Int, bincounts::Vector{Int}, binwidth::Float64)
    ll = 0
    for i in 1:d
        if bincounts[i] > 0
            ll += bincounts[i] * log(bincounts[i] / (n * binwidth))
        end
    end
    ll - (d - 1 + log(d)^2.5)
end

# Optimize the number of bins for a regular one dimensional histogram.
#
# Args:
#   xs: A sample.
#
# Returns:
#   A tuple of the form (d, bincounts), where d gives the optimal number of
#   bins, and bincounts is an array giving the number of occurances in each bin.
#
function choose_bin_count_1d(xs::Vector)
    n = length(xs)
    if n <= 1
        return 1
    end

    x_min, x_max = min(xs), max(xs)
    span = x_max - x_min

    d_min = 3
    d_max = min(150, int(ceil(n / log(n))))
    bincounts = zeros(Int, d_max)

    d_best = d_min
    pll_best = -Inf

    # Brute force optimization: since the number of bins has to be reasonably
    # small to plot, this is pretty quick and very simple.
    for d in d_min:d_max
        binwidth = span / d
        bincounts[1:d] = 0

        for x in xs
            if !isconcrete(x)
                continue
            end
            bincounts[max(1, min(d, int(ceil((x - x_min) / binwidth))))] += 1
        end

        pll = bincount_pll(d, n, bincounts, binwidth)
        if pll > pll_best
            d_best = d
            pll_best = pll
        end
    end

    bincounts[1:d_best] = 0
    binwidth = span / d_best
    for x in xs
        if !isconcrete(x)
            continue
        end
        bincounts[max(1, min(d_best, int(ceil((x - x_min) / binwidth))))] += 1
    end

    (d_best, bincounts)
end

# Optimize the number of bins for regular two dimensional histograms.
#
# Args:
#   xs: Dimension one data.
#   ys: Dimension two data.
#
# Returns:
#   A tuple of the form (dx, dy, bincounts), where dx, dy gives the number of
#   bins in each respective dimension and bincounts is a dx by dy matrix giving
#   the count in each bin.
#
function choose_bin_count_2d(xs::Vector, ys::Vector)
    # For two demensions, I'm just going to optimize the marginal bin counts.
    # This might not be optimal, but its simple and fast.
    dx, _ = choose_bin_count_1d(xs)
    dy, _ = choose_bin_count_1d(ys)

    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)

    # bin widths
    wx = (x_max - x_min) / dx
    wy = (y_max - y_min) / dy

    bincounts = zeros(Int, (dy, dx))
    for (x, y) in zip(xs, ys)
        i = max(1, min(dx, int(ceil((x - x_min) / wx))))
        j = max(1, min(dy, int(ceil((y - y_min) / wy))))
        bincounts[j, i] += 1
    end

    (dx, dy, bincounts)
end


text:
"""
    TableDatasetRecipe(tablefile; catcols, contcols, kwargs...])

Recipe for loading a `TableDataset`. `tablefile` is the path of a file that can
be read as a table. `catcols` and `contcols` indicate the categorical and
continuous columns of the table. If they are not given, they are detected
automatically.
"""
Base.@kwdef struct TableDatasetRecipe <: DatasetRecipe
    file = ""
    catcols = nothing
    contcols = nothing
    loadfn = loadfile
end

recipeblocks(::Type{TableDatasetRecipe}) = TableRow

function loadrecipe(recipe::TableDatasetRecipe, path)
    tablepath = joinpath(path, recipe.file)
    table = recipe.loadfn(tablepath)
    Tables.istable(table) || error("Expected `recipe.loadfn($(tablepath))` to return a table, instead got type $(typeof(table))")
    data = TableDataset(table)

    catcols, contcols = if isnothing(recipe.catcols) || isnothing(recipe.contcols)
        cat, cont = FastAI.getcoltypes(data)
        cat = isnothing(recipe.catcols) ? cat : recipe.catcols
        cont = isnothing(recipe.contcols) ? cont : recipe.contcols
        cat, cont
    else
        recipe.catcols, recipe.contcols
    end

    block = TableRow(
        catcols,
        contcols,
        FastAI.gettransformdict(data, DataAugmentation.Categorify, catcols),
    )
    return data, block
end

struct TableClassificationRecipe <: DatasetRecipe
    recipe::TableDatasetRecipe
    targetcol
end

recipeblocks(::Type{TableClassificationRecipe}) = (TableRow, Label)

function loadrecipe(recipe::TableClassificationRecipe, args...; kwargs...)
    data, block::TableRow = loadrecipe(recipe.recipe, args...; kwargs...)
    recipe.targetcol in block.catcols || error("Expected categorical column $(recipe.targetcol) to exist.")

    data = rows, labels = (data, Tables.getcolumn(data.table, recipe.targetcol))
    blocks = (
        removecol(block, recipe.targetcol),
        Label(unique(eachobs(labels))),
    )

    return data, blocks
end

struct TableRegressionRecipe <: DatasetRecipe
    recipe::TableDatasetRecipe
    targetcol
end

recipeblocks(::Type{TableRegressionRecipe}) = (TableRow, Continuous)

function loadrecipe(recipe::TableRegressionRecipe, args...; kwargs...)
    data, block::TableRow = loadrecipe(recipe.recipe, args...; kwargs...)
    recipe.targetcol in block.contcols || error("Expected continuous column $(recipe.targetcol) to exist.")

    data = rows, labels = (data, Tables.getcolumn(data.table, recipe.targetcol))
    blocks = (
        removecol(block, recipe.targetcol),
        Continuous(1),
    )
    return data, blocks
end

## Utils

removecol(block::TableRow, col) = TableRow(
        Tuple(cat for cat in block.catcols if cat != col),
        Tuple(cont for cont in block.contcols if cat != col),
        Dict(catcol => cats for (catcol, cats) in block.categorydict if catcol != col)
    )


text:
using DAI

bnd = MCBN.BayesNetDAI(3)
MCBN.add_edge!(bnd, 2, 1)
MCBN.add_edge!(bnd, 3, 1)
MCBN.add_edge!(bnd, 3, 2)
DAI.clearBackups!(bnd)

MCBN.set_factor!(bnd, 1, [0.9717959704448,
                  0.02820402955519996,
                  0.4884633419082751,
                  0.5115366580917249,
                  0.4671602284144756,
                  0.5328397715855244,
                  0.09147912958666593,
                  0.9085208704133341])
MCBN.set_factor!(bnd, 2, [0.6038715782985324,
                          0.39612842170146756,
                          0.16877438390015786,
                          0.8312256160998421])
MCBN.set_factor!(bnd, 3, [0.3137903917753795, 0.6862096082246205])

MCBN.check_bnd(bnd)
@show true_cod1 = MCBN.cod(bnd, 1, [2,3])
@show true_cod2 = MCBN.cod(bnd, 1, [2])
@show true_cod3 = MCBN.cod(bnd, 2, [1])

MCBN.setupIA!(bnd)
@show marginal(bnd.ia, VarSet(bnd.verts[1], bnd.verts[2]))
@show marginal(bnd.ia, VarSet(bnd.verts[1]))


text:
using FIGlet
using Test

@testset "FIGlet.jl" begin
    iob = IOBuffer(b"flf2a", read=true);
    @test FIGlet.readmagic(iob) == UInt8['f', 'l', 'f', '2', 'a']

    @test FIGlet.FIGletHeader('$', 6, 5, 16, 15, 11) == FIGlet.FIGletHeader('$', 6, 5, 16, 15, 11, 0, 2, 0)

    @test FIGlet.availablefonts() |> length == 680

    for font in FIGlet.availablefonts()
        @test FIGlet.readfont(font) isa FIGlet.FIGletFont
    end

    ff = FIGlet.readfont("jiskan16")
    iob = IOBuffer()
    print(iob, ff)
    @test String(take!(iob)) == "FIGletFont(n=7098)"

    iob = IOBuffer()
    print(iob, ff.font_characters['㙤'])
    @test String(take!(iob)) == "FIGletChar(ord='㙤')"

    @test_throws FIGlet.FontNotFoundError FIGlet.readfont("wat")
    @test_throws FIGlet.FontError FIGlet.readfont(joinpath(@__DIR__, "..", "README.md"))

    @test FIGlet.smushem(FIGlet.readfont("standard"), ' ', ' ') == ' '
    @test FIGlet.smushem(FIGlet.readfont("standard"), '_', ' ') == '_'
    @test FIGlet.smushem(FIGlet.readfont("standard"), ' ', '_') == '_'

end


text:
#################################
#
#           WENO3
#
##################################

##-------------------------
## Positive

@doc raw"""
    posWENO3!(flux, U)

Updates the vector "flux" so that flux[i] represents the positive WENO3
flux at x_{i+1/2}, i.e. a WENO3 approximation of the flux going from the
left cell to the right cell. Necessarily, `flux` is supposed to be an
OffsetArray with indices from 0 to N.

The leftmost flux is not computed, as it depends on the boundary
condition. For the stencil to be valid for the second-to-leftmost and the
rightmost fluxes, an Upwind approximation is used instead.
"""
function posWENO3!(flux, U)
    N = length(U)
    flux[1] = posUpwind(U, 1)
    flux[N] = posUpwind(U, N)

    local ε = 10^-8
    # Squared smoothness indicators ("non-zero"-ified)
    local β = diff(U).^2
    @. β = 1.0 / ( ε + β )^2
    # Non-normalized weights
    local ω0 = 0.333333333333333333 .* β[1:N-2]
    local ω1 = 0.666666666666666667 .* β[2:N-1]
    local Σα = ω0 + ω1
    # Normalize weights
    @. ω0 /= Σα
    @. ω1 /= Σα
    # Pointwise values
    local U0 = -0.5*U[1:N-2] + 1.5*U[2:N-1]
    local U1 =  0.5*U[2:N-1] + 0.5*U[3:N]

    @. flux[2:N-1] = ω0*U0 + ω1*U1

    nothing
end

@doc raw"""
    posWENO3(U,i::Int)

Computes the positive WENO5 flux of U at x_{i+1/2}.

Ref: https://www.cs.usask.ca/faculty/spiteri/63786-gg.pdf
"""
function posWENO3(U,i::Int)
    local ε = 10^-8

    ## Smoothness indicators
    local β0 = (U[i] - U[i-1])^2
    local β1 = (U[i+1] - U[i])^2

    ## Non-normalized weights
    local α0 = 0.333333333333333333/(ε + β0)^2
    local α1 = 0.666666666666666667/(ε + β1)^2
    local Σα = α0 + α1

    ## Weights (normalized)
    local ω0, ω1 = α0/Σα, α1/Σα

    ## Reconstruction of pointwise values
    local U0 = -0.5*U[i-1] + 1.5*U[i]
    local U1 =  0.5*U[i]   + 0.5*U[i+1]

    ## Final computation
    ω0*U0 + ω1*U1
end

##---------------------------
## Negative

@doc raw"""
    negWENO3!(flux, U)

Updates the vector "flux" so that flux[i] represents the negative WENO3
flux at x_{i+1/2}, i.e. a WENO3 approximation of the flux going from the
right cell to the left cell. Necessarily, `flux` is supposed to be an
OffsetArray with indices from 0 to N.

The rightmost flux is not computed, as it depends on the boundary
condition. For the stencil to be valid for the second-to-rightmost and the
leftmost fluxes, an upwind approximation is used instead.
"""
function negWENO3!(flux, U)
    N = length(U)
    flux[0] = negUpwind(U, 0)
    flux[N-1] = negUpwind(U, N-1)

    local ε = 10^-8
    # Inverse of squared smoothness indicators ("non-zero"-ified)
    local β = diff(U) .^2
    @. β = 1.0 / (ε + β)^2
    # Non-normalized weights
    local ω0 = 0.333333333333333333 .* β[2:N-1]
    local ω1 = 0.666666666666666667 .* β[1:N-2]
    local Σα = ω0 + ω1
    # Normalize weights
    @. ω0 /= Σα
    @. ω1 /= Σα
    # Pointwise values
    local U0 = -0.5*U[3:N]   + 1.5*U[2:N-1]
    local U1 =  0.5*U[2:N-1] + 0.5*U[1:N-2]

    @. flux[1:N-2] = ω0*U0 + ω1*U1

    # @inbounds for i in 2 : N-2
    #     flux[i] = negWENO3(U, i)
    # end
    nothing
end

@doc raw"""
    negWENO3(U,i::Int)

Computes the negative WENO5 flux of U at x_{i+1/2}. This could also be
obtained by symmetry from posWENO3.

Ref: https://www.cs.usask.ca/faculty/spiteri/63786-gg.pdf
"""
function negWENO3(U,i::Int)
    local ε = 10^-8

    ## Smoothness indicators
    local β0 = (U[i+2] - U[i+1])^2
    local β1 = (U[i+1] - U[i])^2

    ## Non-normalized weights
    local α0 = 0.333333333333333333/(ε + β0)^2
    local α1 = 0.666666666666666667/(ε + β1)^2
    local Σα = α0 + α1

    ## Weights (normalized)
    local ω0, ω1 = α0/Σα, α1/Σα

    ## Reconstruction of pointwise values
    local U0 = -0.5*U[i+2] + 1.5*U[i+1]
    local U1 =  0.5*U[i+1] + 0.5*U[i]

    ## Final computation
    ω0*U0 + ω1*U1
end

max_stars_repo_name: tremelow/StiffKinetic

max_issues_repo_name: tremelow/StiffKinetic

max_forks_repo_name: tremelow/StiffKinetic


text:
function mdtext(s::String; escape_underscores = false, kwargs...)
    escape_underscores && (s = replace(s, "_"=>"\\_"))
    m = Markdown.parse(s)
    return m
end


text:
# A direction generator generates the search directions to use at each step of
# a GSS search.
abstract DirectionGenerator

type ConstantDirectionGen <: DirectionGenerator
  directions::Array{Float64, 2}

  ConstantDirectionGen(directions) = begin
    new(directions)
  end
end

function directions_for_k(cg::ConstantDirectionGen, k)
  cg.directions # Always the same regardless of k...
end

# We can easily do a compass search with GSS by generating directions
# individually (+ and -) for each coordinate.
compass_search_directions(n) = ConstantDirectionGen([eye(n,n) -eye(n, n)])

using BlackBoxOptim

# Generating set search as described on page 21 (405) in Kolda2003 but extended
# with random ordering or frequency adapted ordering of directions.
function generating_set_search(problem;
  max_seconds = 2*numdims(problem), max_evals_per_dim = 1e5,
  ftol = 1e-7,
  delta_tol = 1e-13,
  step_size = false, x = false, random_order = false,
  freq_adapt_order = false,
  direction_gen = compass_search_directions(numdims(problem)))

  n = numdims(problem)
  ss = search_space(problem)

  step_size = 0.50 * minimum(diameters(ss))

  @assert delta_tol > 0
  @assert step_size > delta_tol

  max_evals = max_evals_per_dim * n

  directions = [eye(n,n) -eye(n, n)]

  x = rand_individual(ss)
  fbest = eval1(x, problem)
  num_fevals = 1

  archive = BlackBoxOptim.TopListArchive(n, 10)
  add_candidate!(archive, fbest, x[:], num_fevals)

  fc = fbest
  candidate = zeros(n, 1)
  k = 0
  termination_reason = :unknown
  order = collect(1:size(directions, 2))

  if freq_adapt_order
    fa_directions = BlackBoxOptim.FrequencyAdapter(size(directions, 2))
  end

  start_time = time()

  while(true)

    if step_size < delta_tol
      termination_reason = "Step size smaller than delta_tol"
      break
    end

    if (time() - start_time) > max_seconds
      termination_reason = "Exceeded time budget"
      break
    end

    if num_fevals > max_evals
      termination_reason = "Exceeded function eval budget"
      break
    end

    # Get the directions for this iteration
    k += 1
    directions = directions_for_k(direction_gen, k)

    # Check all directions to find a better point
    found_better = false

    # Freq adaptation takes precedence over random ordering so check first.
    if freq_adapt_order
      order = BlackBoxOptim.create_new_block!(fa_directions)
    elseif random_order == true
      shuffle!(order)
    end

    for(direction in order)
      candidate = x + step_size .* directions[:, direction]

      fc = eval1(candidate, problem)
      num_fevals += 1

      # Update freq based on progress
      if freq_adapt_order
        improvement = (fbest - fc) / fbest
        #BlackBoxOptim.update!(fa_directions, direction, max(0, improvement))
        BlackBoxOptim.update!(fa_directions, direction, improvement)
      end

      if fc < fbest
        found_better = true
        break
      end

    end

    if found_better
      fbest = fc
      x = candidate
      add_candidate!(archive, fbest, x[:], num_fevals)
      println(num_fevals, ": fbest = ", fbest)

      if fitness_is_within_ftol(problem, ftol, fbest)
        termination_reason = "Within ftol"
        break
      end
    else
      step_size = step_size / 2
    end
  end

  return x, fbest, num_fevals, termination_reason, archive
end

function min_distance(point, points)
  dist = norm(point - points[:,1])
  for(i in 2:size(points, 2))
    dist = minimum([dist, norm(point - points[:, i])])
  end
  dist
end

function maximize_min_distance(n, points, num_samples = 100,
  sampler = () -> randn(n,1))
  xbest = xnew = sampler()
  dist = min_distance(xnew, points)
  for(i in 1:num_samples)
    xnew = sampler()
    newdist = min_distance(xnew, points)
    if newdist > dist
      xbest = xnew
      dist = newdist
    end
  end
  xbest
end

function restart_gss(f, n; max_seconds = 2*n, delta_tol = 1e-7,
  step_size = 1.0,
  x = false, diverse_starting_points = false, diameter = 1.0,
  random_order = false, freq_adapt_order = false,
  known_fmin = :unknown, ftol = 1e-8,
  direction_gen = compass_search_directions(n))

  now = start_time = time()
  total_fbest = Inf
  xb = total_fevals = termination_reason = 0
  num_runs = 0

  sampler_func = () -> diameter * randn(n, 1)

  if x == false
    # Generate the first starting point.
    starting_points = sampler_func()
  else
    starting_points = x
  end

  while(now < start_time + max_seconds)
    time_left = max_seconds - (now - start_time)

    num_runs += 1

    x, fbest, num_fevals, termination_reason = generating_set_search(f, n;
      max_seconds = time_left, delta_tol = delta_tol, random_order = random_order,
      freq_adapt_order = freq_adapt_order,
      step_size = step_size, x = starting_points[:,num_runs],
      known_fmin = known_fmin, ftol = ftol,
      direction_gen = direction_gen)

    step_size = step_size * 2
    #delta_tol /= 10

    total_fevals += num_fevals

    if fbest < total_fbest
      total_fbest = fbest
      xb = x
      println(num_runs, ": New global best = $(total_fbest)")
      if termination_reason == "Within ftol"
        break
      end
    end

    if diverse_starting_points
      xnew = maximize_min_distance(n, starting_points, 1000, sampler_func)
    else
      xnew = sampler_func()
    end
    starting_points = hcat(starting_points, xnew)

    now = time()
  end

  if time() > start_time + max_seconds
    termination_reason = "Exceeded time budget (elasped = $(time()-start_time))"
  end

  return xb, total_fbest, total_fevals, termination_reason
end

# Run parameter study with given parameters.
function run_with_params(of, n)
  start_time = time()
  xb, fb, fes, tr = restart_gss(of, n; known_fmin = 0.0, random_order = false)
  end_time = time()
  elapsed = end_time - start_time

  prefix_header = "StartTime,ElapsedTime,Function,N,FitnessBest,TotalFuncEvals,TerminationReason"
  println(join([prefix_header, map((i) -> "x$(i)", 1:length(xb))], ","))
  println(join([strftime("%Y%m%d-%H%M%S", start_time), elapsed, of, n, fb, fes, tr, xb], ","))
end

#of = rastrigin
#n = 5
#@time fbs, fes = repeated_runs(
#  ((n) -> restart_gss(of, n; known_fmin = 0.0, random_order = false)),
#  n);

#@time generating_set_search(sphere, 100; known_fmin = 0.0, random_order = true)
#@time generating_set_search(rosenbrock, 16; known_fmin = 0.0)
#@time restart_gss(rosenbrock, 32; known_fmin = 0.0, random_order = false)
#
#@time restart_gss(rosenbrock, 16; known_fmin = 0.0)
#@time restart_gss(xtransform(16, rosenbrock), 16; known_fmin = 0.0)
#
#@time restart_gss(cigar, 32; known_fmin = 0.0)

#of = rosenbrock
#@time ts, fbs, fes = compare_params([
#  (of, 2, false, false, false, 1.0),
#  (of, 2, false, false, false, 2.0),
#  (of, 2, false, false, true, 1.0),
#  (of, 2, false, false, true, 2.0)
#  ],
#  ((f, n, fao, ro, div, diam) ->
#    restart_gss(xtransform(n, f), n;
#      known_fmin = 0.0, freq_adapt_order = fao, random_order = ro,
#      diverse_starting_points = div, diameter = diam)),
#  5
#);
#
#
#of = griewank
#n = 2
#@time fbs, fes = repeated_runs(
#  ((n) -> restart_gss(of, n; known_fmin = 0.0, random_order = false)),
#  n);

#of = rosenbrock
#@time ts, fbs, fes = compare_params([
#  (of, 10, true, false),
#  #(of, 10, false, true),
#  #(of, 10, false, false),
#  #(of, 30, true, false),
#  #(of, 30, false, true),
#  (of, 30, false, false)
#  ],
#  ((f, n, fao, ro) ->
#    restart_gss(xtransform(n, f), n;
#      known_fmin = 0.0, freq_adapt_order = fao, random_order = ro)),
#  2
#)
#
#f = schwefel2_21
#n = 16
#
#of = deceptive_cuccu2011(30, 2)
#of = rastrigin
#n = 5
#@time fbs, fes = repeated_runs(
#  ((n) -> restart_gss(of, n; known_fmin = 0.0, random_order = false)),
#  n);
#
#
## Comparing to Cuccu2011
#of = deceptive_cuccu2011(15, 2)
#of = rastrigin
#@time ts, fbs, fes = compare_params([
#  (of, 2, true, false),
#  #(of, 2, false, true),
#  (of, 2, false, false),
#  (of, 5, true, false),
#  #(of, 5, false, true),
#  (of, 5, false, false),
#  (of, 10, true, false),
#  #(of, 10, false, true),
#  (of, 10, false, false),
#  (of, 20, true, false),
#  #(of, 20, false, true),
#  (of, 20, false, false)
#  ],
#  ((f, n, fao, ro) ->
#    restart_gss(xtransform(n, f), n;
#      known_fmin = 0.0, freq_adapt_order = fao, random_order = ro)),
#  25
#);
#
## Cuccu tells the median number of generations but not clear how many func evals
## this corresponds to. Lets assume he uses the standard 4+3*log(d) rule to set
## population size => for d=20 we should multiply by 12.
## Table IV on page 5 of Cuccu2011 should thus read:
##   d=2,  2112
##   d=5,  1872
##   d=10, 633
##   d=20, 16488
##
## These results for d=20 are impressive.


text:
# acc reachability script
# ENV["JULIA_DEBUG"] = Main # for debugging

using OVERTVerify
using LazySets
using Dates
using JLD2

@testset "reachability" begin
    controller = "acc_controller"
    controller_filepath = "../nnet_files/jmlr/"*controller*".nnet"
    # println("Controller is: ", controller)
    query = OvertQuery(
        ACC,  # problem
        controller_filepath,    # network file
        Id(),      	# last layer activation layer Id()=linear, or ReLU()=relu
        "MIP",     	# query solver, "MIP" or "ReluPlex"
        4,        	# ntime
        0.1,       	# dt
        -1,        	# N_overt
        )

    # x1,x2,x3 are lead vehicle variables
    # x4,x5,x6 are ego vehicle variables
    x_lead = [90,110]
    v_lead = [32, 32.2]
    gamma_lead = [0,0]
    x_ego = [10,11]
    v_ego = [30, 30.2]
    gamma_ego = [0, 0]
    var_list = [x_lead, v_lead, gamma_lead, x_ego, v_ego, gamma_ego]
    input_set = Hyperrectangle(
        low=[v_range[1] for v_range in var_list],
        high=[v_range[2] for v_range in var_list]
        )

    concretization_intervals = [2, 2]

    if DEFAULT_MODEL == "glpk" || DEFAULT_MODEL == "GLPK"
        @warn "Skipping reachability test because GLPK is slow; try testing with Gurobi installed"
    else
        t1 = Dates.time()
        concrete_state_sets, symbolic_state_sets, concrete_meas_sets, symbolic_meas_sets = symbolic_reachability_with_concretization(query, input_set, concretization_intervals)
        t2 = Dates.time()
        dt = (t2-t1)
        print("elapsed time= $(dt) seconds")

        # Intersect all sets with output constraint and see if
        # reachable set is fully within safe set OR check to see if it ever intersects unsafe set
        # they are equivalent
        # We want the measurement to be greater than 10, always. So the unsafe set if <= 10
        avoid_sets = [HalfSpace([1.], 10.)] # 1*y <= 10

        reachable_meas_sets = clean_up_meas_sets(concrete_meas_sets, symbolic_meas_sets, concretization_intervals)

        t1 = time()
        safe, violations = check_avoid_set_intersection(reachable_meas_sets, input_set, avoid_sets)
        dt_check = time() - t1
        @test safe == true
    end
end


text:
using StaticArrays

struct Member{Tval, Tid}
  val::Tval
  id::Tid
end

import Base.isless
function isless(a::Member{Tval, Tid}, b::Member{Tval, Tid}) where{Tval, Tid}
  return isless((a.val, a.id), (b.val, b.id))
end

struct Family{Tval, Tid}
  n_columns::MVector{1, Tid}
  n_nz ::MVector{1, Tid}
  n_buffer::MVector{1, Tid}

  # This array gives for contains the ordering. The i-th parent in the
  #daycare has id P[i]
  P::Vector{Tid}

  # This array gives for contains the inverse ordering.
  revP::Vector{Tid}

  # The array that contains the first "child" for every parent
  colptr::Vector{Tid}

  # The array that contains the global id-s of the children
  rowval::Vector{Member{Tval,Tid}}

  # parents[id] contains a member struct specifying the distance of id to its parent and the position in the ordering of the parent
  parents::Vector{Member{Tval,Tid}}
end

function Family{Tval,Tid}(dofs::AbstractArray{<:AbstractDOF}) where {Tval,Tid}
  N = length(dofs)
  return  Family{Float64, Int64}(MVector(zero(Tid)), MVector(zero(Tid)), MVector(N), zeros(Tid, N), zeros(Tid, N), zeros(Tid, N + one(Tid)), Vector{Member{Tval, Tid}}(undef, N), Vector{Member{Tval, Tid}}(undef, N))
end

function Family(dofs::AbstractArray{Point{d,Tval}}) where {d, Tval}
  return  Family{Tval, Int}(dofs)
end

#Function that begins a new parent aka column in daycare
function new_column!(f::Family, Id)
  f.n_columns[1] += 1
  f.P[f.n_columns[1]] = Id
  f.revP[Id] = f.n_columns[1]
  f.colptr[f.n_columns[1]] = f.n_nz[1] + 1
  f.colptr[f.n_columns[1] + 1] = f.n_nz[1] + 1
end

function new_child!(f::Family, new_child)
  # If the buffer is flowing over, increase it
  if f.n_nz[1] >= f.n_buffer[1]
    if f.n_nz[1] <= 1e6
      f.n_buffer[1] = 2 * f.n_buffer[1]
    else
      f.n_buffer[1] = f.n_buffer[1] + 1e6
    end
    resize!( f.rowval, f.n_buffer[1] )
  end
  f.n_nz[1] += 1
  f.colptr[f.n_columns[1] + 1] += 1
  f.rowval[f.n_nz[1]] = new_child
end

function new_children!(f::Family, new_children)
  # If the buffer is flowing over, increase it
  while f.n_nz[1] + size(new_children,1) >= f.n_buffer[1] - 1
    if f.n_nz[1] <= 1e6
      f.n_buffer[1] = 2 * f.n_buffer[1]
    else
      f.n_buffer[1] = f.n_buffer[1] + 1e6
    end
    resize!(f.rowval, f.n_buffer[1])
  end

  f.n_nz[1] += size(new_children,1)
  f.colptr[f.n_columns[1] + 1] += size(new_children,1)
  f.rowval[f.n_nz[1] - size(new_children,1) + 1 : f.n_nz[1]] .= new_children
end

function column_iterator(f, column_index)
  return f.colptr[column_index] : (f.colptr[column_index + 1] - 1)
end


text:
"""
    GS

Type representing Gerchberg-Saxton phase retrieval algorithm.

Requires (in the current realisation) two amplitude constraints or two arrays of equal sizes.
"""
struct GS <: ProjectionsMethod #todo should be sets part of this or added to the step! only?
    a::ConstrainedByAmplitude
    A::ConstrainedByAmplitude
    forward
    backward
end
export GS, PR

function PR(a::Array,A::Array)
    size(a) == size(A) ? (P, PB) = (plan_fft(complex(A)), plan_ifft(complex(A))) : error("Array sizes do not match")
    forward(xf,x) = mul!(xf, P, x)
    backward(xb,x) = mul!(xb, PB, x)
    TwoSetsFP(ConstrainedByAmplitude(a), ConstrainedByAmplitude(A),  forward, backward)
end

function init!(alg::GS,  x⁰)
    a = alg.a
    A = alg.A
    if !( size(a) == size(A) == size(x⁰) )
        println("cannot intialise Gerschberg-Saxton, dimensions do not match")
        # raise error or pad a smaller array with zeroes
        # also can include realingment of the sets to remove any linear pahse  term. But then need to make GS mutable
        return
    end
end

function step!(alg::GS, xᵏ)
    a = alg.a
    A = alg.A

#     # step by step explanation
#     yk = fft(xᵏ)
#     println(yk)
#     zk = project(A, yk)
#     println('z',zk)
#     yk = ifft(zk)
#     println(yk)
#     zk = project(a, yk)
#     println('z',zk)
#

    return project(ifft(project(fft(xᵏ), A)),a)
    #todo replace with plan
end


text:
# Model
include_model("particle")

# Dimensions
nq = 3                    # configuration dimension
nu = 3                    # control dimension
nc = 1                    # number of contact points
nf = 2                    # number of parameters for friction cone
nb = nc * nf
ns = 1

# Parameters
μ = 0.5      # coefficient of friction
mass = 1.0   # mass
g = 9.81

n = 2 * nq
m = nu + nc + nb + ns
d = 0

idx_u = (1:nu)
idx_λ = nu .+ (1:nc)
idx_b = nu + nc .+ (1:nb)
idx_ψ = (1:0) #nu + nc + nb .+ (1:nc)
idx_η = (1:0) #nu + nc + nb + nc .+ (1:nb)
idx_s = nu + nc + nb .+ (1:ns)

model = Particle{Discrete, FixedTime}(n, m, d,
				 mass, g, μ,
				 nq,
				 nu,
				 nc,
				 nf,
				 nb,
				 ns,
				 idx_u,
				 idx_λ,
				 idx_b,
				 idx_ψ,
				 idx_η,
				 idx_s)

function fd(model::Particle{Discrete, FixedTime}, x⁺, x, u, w, h, t)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2⁺ = view(x⁺, 1:model.nq)
	q2⁻ = view(x, model.nq .+ (1:model.nq))
	q1 = view(x, 1:model.nq)
	u_ctrl = view(u, model.idx_u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)

 [q2⁺ - q2⁻;
 ((1.0 / h) * (M_func(model, q1) * (SVector{3}(q2⁺) - SVector{3}(q1))
 - M_func(model, q2⁺) * (SVector{3}(q3) - SVector{3}(q2⁺)))
 + transpose(B_func(model, q3)) * SVector{3}(u_ctrl)
 + transpose(N_func(model, q3)) * SVector{1}(λ)
 + transpose(_P_func(model, q3)) * SVector{2}(b)
 - h * G_func(model, q2⁺))]
end

# Horizon
T = 51

# Time step
tf = 5.0
h = tf / (T-1)

# Bounds
_uu = Inf * ones(model.m)
_uu[model.idx_u] .= 0.0
_ul = zeros(model.m)
_ul[model.idx_b] .= -Inf
ul, uu = control_bounds(model, T, _ul, _uu)

# Initial and final states
q1 = [0.0, 0.0, 1.0]
v1 = [3.0, 5.0, 0.0]
v2 = v1 - G_func(model,q1) * h
q2 = q1 + 0.5 * h * (v1 + v2)

x1 = [q1; q2]

xl, xu = state_bounds(model, T, x1 = x1)

# Objective
obj = PenaltyObjective(1.0e5, model.m)

# Constraints
include_constraints("contact_differentiable_socp")
con_contact = contact_diff_constraints(model, T)

# Problem
prob = trajectory_optimization_problem(model,
               obj,
               T,
               h = h,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con_contact)

# Trajectory initialization
x0 = [0.01 * rand(model.n) for t = 1:T] #linear_interpolation(x1, x1, T) # linear interpolation on state
u0 = [0.001 * rand(model.m) for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem
include_snopt()
@time z̄ , info = solve(prob, copy(z0),
	tol = 1.0e-2, c_tol = 1.0e-2,
	nlp = :ipopt)

check_slack(z̄, prob)
x̄, ū = unpack(z̄, prob)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, state_to_configuration(x̄), Δt = h)

friction_socp(model, x̄[2], ū[1], h)


text:
let
    roadlength = 50.0
    lane_width = 5.0
    roadway,junction = gen_intersection(roadlength = roadlength,lane_width = lane_width)
    scene = Frame(Entity{VehicleState, BicycleModel, Int},100)
    id=1
    push!(scene,Entity(VehicleState(VecSE2(0.0,0.0,0.0), roadway, 10.0),
        BicycleModel(VehicleDef(AgentClass.CAR, 4.826, 1.81)),id))
    id += 1
    push!(scene,Entity(VehicleState(VecSE2(10.0,0.0,0.0), roadway, 10.0),
        BicycleModel(VehicleDef(AgentClass.CAR, 4.826, 1.81)),id))
    id += 1
    push!(scene,Entity(VehicleState(VecSE2(20.0,0.0,0.0), roadway, 9.0),
        BicycleModel(VehicleDef(AgentClass.CAR, 4.826, 1.81)),id))
    connect_two_seg!(roadway.segments[2],roadway.segments[3],roadway)
    connect_two_seg!(roadway.segments[4],roadway.segments[5],roadway)
    connect_two_seg!(roadway.segments[6],roadway.segments[7],roadway)
    connect_two_seg!(roadway.segments[8],roadway.segments[1],roadway)
    framerate = 24
    timeStep = 0.05
    models = Dict{Int, DriverModel}()

    models[1]=UrbanDriver(timeStep,direction=1,a_cp_max=0.7*9.8)
    models[2]=UrbanDriver(timeStep,direction=2,a_cp_max=0.6*9.8)
    models[3]=MultiPtsTurningDriver(timeStep,v_max=20.0,v_min=0.0,turning_direction=3)
    excute_action!(models[3],0.0,0,scene, roadway, 1)
    actions = get_actions!(Array{Any}(undef, length(scene)), scene, roadway, models)
    for i=1:100
        get_actions!(actions, scene, roadway, models)
        tick!(scene, roadway, actions, timeStep)
    end
    actions = get_actions!(Array{Any}(undef, length(scene)), scene, roadway, models)
    @test abs(actions[1].a_lat-0.0)<0.01
    @test abs(actions[1].a_lon-(-0.05641834275191615))<0.01
    @test actions[1].direction==1
    @test abs(actions[2].a_lat-0.0)<0.01
    @test abs(actions[2].a_lon-(-0.0023799116323655056))<0.1
    @test actions[2].direction==2
    @test abs(actions[3].a-0.07816860295178636)<0.01
    @test abs(actions[3].δ-0.00531000497036248)<0.01
    @test abs(actions[3].direction-3 )<0.01
end

max_stars_repo_name: sisl/UrbanDrivingSimulation

max_issues_repo_name: sisl/UrbanDrivingSimulation

max_forks_repo_name: sisl/UrbanDrivingSimulation


text:
# Solver for homogeneous sites
@doc raw"""
    sumBassHomogeneousSolver(p,q,m,c0,alpha,F,T,n,Pi)

This function maximizes the sum of homogeneous Bass functions and returns the maximizer
and maximum. Note that any permutation of the maximizer is also a maximizer.
"""
function sumBassHomogeneousSolver(p::Real,q::Real,m::Real,c0::Real,alpha::Real,F::Real,T::Real,n::Int,Pi::Real)
    if n*F > Pi
        throw(DomainError(Pi,"the problem is unfeasible because the capacity is too small."))
    end
    if n*T <= Pi # Trivial solution
        x = fill(float(T),n)
        v = alpha*n*Bass(T,p,q,m,c0)
        return x,v
    end

    # Case 1
    vstar = 0.0
    xstar = 0
    for i in 1:n # This loop ensures that case does not violate the constraint x<T
        if F+(Pi-F*n)/i < T
            objective = [j * (Bass(F + (Pi-F*n)/j,p,q,m,c0) - Bass(F,p,q,m,c0)) for j in i:n]
            xstar = argmax(objective) + i - 1
            vstar = alpha*(xstar*Bass(F + (Pi-F*n)/xstar,p,q,m,c0) + (n-xstar)*Bass(F,p,q,m,c0))
            break
        end
    end

    # Case 2
    xstar2 = floor(Int,(Pi-n*F)/(T-F))
    vstar2 = alpha*(xstar2*Bass(T,p,q,m,c0) + Bass(Pi-xstar2*T-(n-xstar2-1)*F,p,q,m,c0) +
        (n-xstar2-1)*Bass(F,p,q,m,c0))

    x = Vector{Float64}(undef,n)
    if vstar >= vstar2
        x[1:xstar] .= F + (Pi-F*n)/xstar
        x[(xstar+1):n] .= F
        return x,vstar
    else
        x[1:xstar2] .= T
        x[xstar2+1] = Pi-xstar2*T-(n-xstar2-1)*F
        x[xstar2+2:n] .= F
        return x,vstar2
    end

end

max_stars_repo_name: andres-alban/res-alloc-bass-demands

max_issues_repo_name: andres-alban/res-alloc-bass-demands

max_forks_repo_name: andres-alban/res-alloc-bass-demands


text:
### Module for Simulating Pose Graphs ###
### Ted Steiner, MIT AeroAstro/CSAIL ###
### August 2014 - June 2015 ###

module PoseGraphSimulation

import Graphs
import Winston
import Geodesy
import OpenStreetMap

export Pose2D,
       Meas2D,
       BetweenFactor,
       PosegraphData

export makePoseGraph
export getPoses2D, getPoseOdometry, getMeasurements, getLoopFactors, getPriorMeasurements
export getWaypoints, computeRoute, interpolateRoute!, distanceFromStart
export posegraphError, epsilonTrajectory, errorRatio, errorRatios, alignPoses!
export plotPoseGraph, plotPoseGraphTruth, plotEpsilons
export readiSAMOutput, readiSAMTruth, readiSAMOutputCov
export writeMeasurements, writeTruthFile
export isam

include("posegraph_types.jl")
include("posegraph_routing.jl")
include("posegraph_gen.jl")
include("posegraph_meas.jl")
include("posegraph_io.jl")
include("posegraph_error.jl")
include("posegraph_plot.jl")
include("posegraph_solve.jl")

end


text:
#!/usr/bin/env julia

# A Julia script that reproduces expected results of the test module
# 'test/intfunctionTest.cpp'.
#
# From a shell, run the script as:
#   julia /path/to/intfunction.jl
#
# From Julia, the script can be run as:
#   include("path/to/intfunction.jl")

function intfunctionTest()
    nums = [ 12, 100, 37423 ]
    for n = nums
       println("floor(sqrt(", n, ")) = ", convert(Integer, floor(sqrt(n))) )
    end

    println()
    for n = nums
       println("ceil(sqrt(", n, ")) = ", convert(Integer, ceil(sqrt(n))) )
    end

    println()
    nums = [ 1, 8, 63, 64, 1024, 1025, 65000 ]
    for n = nums
       println("floor(log2(", n, ")) = ", convert(Integer, floor(log2(n))) )
    end

    println()
    for n = nums
       println("ceil(log2(", n, ")) = ", convert(Integer, ceil(log2(n))) )
    end

    println()
    nums = [ 1, 16, 65 ]
    for n = nums
       println("floor(log2(", n, ")) = ", convert(Integer, floor(log2(n))) )
    end

    println()
    for n = nums
       println("ceil(log2(", n, ")) = ", convert(Integer, ceil(log2(n))) )
    end
end

intfunctionTest()

max_stars_repo_name: jkovacic/math

max_issues_repo_name: jkovacic/math

max_forks_repo_name: jkovacic/math


text:
using Dates

struct DiscountingSwapEngine{Y <: YieldTermStructure} <: PricingEngine
  yts::Y
  includeSettlementDateFlows::Bool

  # function call(::Type{DiscountingSwapEngine})
  #   new{YieldTermStructure}()
  # end
  #
  # function call{Y}(::Type{DiscountingSwapEngine}, yts::Y, includeSettlementDateFlows::Bool = true)
  #   new{Y}(yts, includeSettlementDateFlows)
  # end
  # DiscountingSwapEngine(yts::Y, includeSettlementDateFlows) where {Y} = new(yts, includeSettlementDateFlows)
  DiscountingSwapEngine{Y}(yts::Y, includeSettlementDateFlows::Bool = true) where {Y} = new{Y}(yts, includeSettlementDateFlows)
end

DiscountingSwapEngine() = DiscountingSwapEngine{NullYieldTermStructure}(NullYieldTermStructure(), true)

function _calculate!(pe::DiscountingSwapEngine, swap::Swap)
  # stuff
  # println("NEW ONE=============================================================================")
  # if swap.rate.value > 0.0323
  #   error("DUIE")
  # end
  swap.results.value = 0.0
  yts = pe.yts

  ref_date = reference_date(yts)
  swap.results.npvDateDiscount = discount(yts, ref_date)
  # for (i, leg) in enumerate(swap.legs)
  for i = 1:length(swap.legs)
    leg = swap.legs[i]
    swap.results.legNPV[i], swap.results.legBPS[i] = npvbps(leg, yts, ref_date, ref_date, pe.includeSettlementDateFlows)
    swap.results.legNPV[i] *= swap.payer[i]
    swap.results.legBPS[i] *= swap.payer[i]

    d1 = accrual_start_date(leg.coupons[1])::Date
    if d1 >= ref_date
      swap.results.startDiscounts[i] = discount(yts, d1)
    end

    d2 = accrual_end_date(leg.coupons[end])::Date
    if (d2 >= ref_date)
      swap.results.endDiscounts[i] = discount(yts, d2)
    end

    swap.results.value += swap.results.legNPV[i]
  end

  if swap.results.legBPS[1] != 0.0
    # println("fixedRate: $(swap.fixedRate)")
    # println("NPV: $(swap.results.value)")
    # println("legBPS: $(swap.results.legBPS[1])")
    swap.results.fairRate = swap.fixedRate - swap.results.value / (swap.results.legBPS[1] / basisPoint)
  end

  return swap
end

clone(pe::DiscountingSwapEngine, ts::Y) where {Y <: YieldTermStructure} = DiscountingSwapEngine{Y}(ts, pe.includeSettlementDateFlows)


text:
export AABox
export measure, maxima, minima, xmin, ymin, zmin, xmax, ymax, zmax, Δx, Δy, Δz,
       facets, ridges, peaks, faces, edges, vertices

"""
    AABox(minima::Point{Dim,T}, maxima::Point{Dim,T})

A `Dim`-dimensional axis-aligned bounding box using two `Dim`-dimensional
points, representing the `minima` and `maxima` of the box.
"""
struct AABox{Dim,T}
    minima::Point{Dim,T}
    maxima::Point{Dim,T}
    function AABox{Dim,T}(minima::Point{Dim,T}, maxima::Point{Dim,T}) where {Dim,T}
        if any(i->maxima[i] ≤ minima[i], 1:Dim)
            error("AABox maxima must be greater that the minima")
        end
        return new{Dim,T}(minima, maxima)
    end
end

# constructors
function AABox(minima::Point{Dim,T}, maxima::Point{Dim,T}) where {Dim,T}
    return AABox{Dim,T}(minima, maxima)
end
AABox(minima, maxima) = AABox(Point(minima), Point(maxima))

maxima(aab::AABox) = aab.maxima
minima(aab::AABox) = aab.minima
Base.isapprox(aab₁::AABox, aab₂::AABox) = minima(aab₁) ≈ minima(aab₂) &&
                                          maxima(aab₁) ≈ maxima(aab₂)
xmin(aab::AABox   ) = aab.minima[1]
ymin(aab::AABox   ) = aab.minima[2]
zmin(aab::AABox{3}) = aab.minima[3]
xmax(aab::AABox   ) = aab.maxima[1]
ymax(aab::AABox   ) = aab.maxima[2]
zmax(aab::AABox{3}) = aab.maxima[3]
Δx(aab::AABox   ) = xmax(aab) - xmin(aab)
Δy(aab::AABox   ) = ymax(aab) - ymin(aab)
Δz(aab::AABox{3}) = zmax(aab) - zmin(aab)
measure(aab::AABox) = prod(maxima(aab) - minima(aab))
ridges(aab::AABox{2}) = vertices(aab)
facets(aab::AABox{2}) = edges(aab)
peaks( aab::AABox{3}) = vertices(aab)
ridges(aab::AABox{3}) = edges(aab)
facets(aab::AABox{3}) = faces(aab)

function vertices(aab::AABox{2})
    # Ordered CCW
    return Vec(Point(xmin(aab), ymin(aab)),
               Point(xmax(aab), ymin(aab)),
               Point(xmax(aab), ymax(aab)),
               Point(xmin(aab), ymax(aab))
              )
end

function edges(aab::AABox{2})
    v = vertices(aab)
    return Vec(LineSegment(v[1], v[2]),
               LineSegment(v[2], v[3]),
               LineSegment(v[3], v[4]),
               LineSegment(v[4], v[1])
              )
end

function vertices(aab::AABox{3})
    # in CCW order, low z then high z
    #      y
    #      ^
    #      |
    #      |
    #      |------> x
    #     /
    #    /
    #   𝘷
    #  z
    #
    #   4----3
    #  /    /|
    # 8----7 |
    # |    | 2
    # |    |/
    # 5----6
    return Vec(Point(xmin(aab), ymin(aab), zmin(aab)),
               Point(xmax(aab), ymin(aab), zmin(aab)),
               Point(xmax(aab), ymax(aab), zmin(aab)),
               Point(xmin(aab), ymax(aab), zmin(aab)),
               Point(xmin(aab), ymin(aab), zmax(aab)),
               Point(xmax(aab), ymin(aab), zmax(aab)),
               Point(xmax(aab), ymax(aab), zmax(aab)),
               Point(xmin(aab), ymax(aab), zmax(aab))
              )
end

function edges(aab::AABox{3})
    # in CCW order, low z, then high z, then the segments that attach low and
    # high z in CCW order.
    #      y
    #      ^
    #      |
    #      |
    #      |------> x
    #     /
    #    /
    #   𝘷
    #  z
    #       3
    #    +----+
    # 11/ 7  /|
    #  +----+ | 2
    # 8|    |6+
    #  |    |/ 9
    #  +----+
    #     5
    v = vertices(aab)
    return Vec(LineSegment(v[1], v[2]), # lower z
               LineSegment(v[2], v[3]),
               LineSegment(v[3], v[4]),
               LineSegment(v[4], v[1]),
               LineSegment(v[5], v[6]), # upper z
               LineSegment(v[6], v[7]),
               LineSegment(v[7], v[8]),
               LineSegment(v[8], v[5]),
               LineSegment(v[1], v[5]), # lower, upper connections
               LineSegment(v[2], v[6]),
               LineSegment(v[3], v[7]),
               LineSegment(v[4], v[8])
              )
end

function faces(aab::AABox{3})
     v = vertices(aab)
     return Vec(Quadrilateral(v[1], v[2], v[3], v[4]),
                Quadrilateral(v[5], v[6], v[7], v[8]),
                Quadrilateral(v[1], v[2], v[6], v[5]),
                Quadrilateral(v[2], v[3], v[7], v[6]),
                Quadrilateral(v[3], v[4], v[8], v[7]),
                Quadrilateral(v[4], v[1], v[5], v[8]),
               )
end

function Base.show(io::IO, aab::AABox)
    print(io, "AABox($(minima(aab)), $(maxima(aab)))")
end

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
include("findroots.jl")
include("dual_promotion.jl")


text:
# Chebyshev regression: least-square fits of data
# to multidimensional Chebyshev polynomials.

function chebregression(x::AbstractVector{SVector{N,Td}}, y::AbstractVector{T},
                        lb::SVector{N,Td}, ub::SVector{N,Td}, order::NTuple{N,Int}) where {N,Td<:Real,T<:Union{SVector,Number}}
    length(x) == length(y) || throw(DimensionMismatch())
    length(x) ≥ prod(order .+ 1) || throw(ArgumentError("not enough data points $(length(x)) to fit to order $order"))

    # assemble rhs as matrix
    Y = Array{float(eltype(T))}(undef, length(y), length(first(y)))
    for j = 1:length(y)
        Y[j,:] .= y[j + (firstindex(y)-1)]
    end

    # assemble lhs matrix
    # (TODO: this algorithm is O(length(x) * length(c)²),
    #  but it should be possible to do it in linear time.
    #  However, the A \ Y step is also O(mn²), so this
    #  only affects the constant factor in the complexity.)
    c = ChebPoly{N,Td,Td}(zeros(Td, order .+ 1), lb, ub)
    A = Array{Td}(undef, length(x), length(c.coefs))
    for i = 1:length(c.coefs)
        c.coefs[i] = 1 # basis function
        for j = 1:length(x)
            A[j,i] = c(x[j + (firstindex(x)-1)])
        end
        c.coefs[i] = 0 # reset
    end

    # least-square solution
    C = A \ Y

    # rearrange C into a ChebPoly
    Tc = typeof(zero(T) * one(eltype(Y)))
    coefs = Array{Tc,N}(reshape(reinterpret(Tc, vec(transpose(C))), order .+ 1))
    return ChebPoly{N,Tc,Td}(coefs, lb, ub)
end

# convert arrays to vectors of svectors or scalars
to_svectors(x::AbstractVector{<:Number}) = x
to_svectors(x::AbstractVector{<:Number}, ::Val{1}) = SVector{1}.(x)
to_svectors(x::AbstractVector{<:SVector{N}}) where {N} = x
to_svectors(x::AbstractVector{<:SVector{N}}, ::Val{N}) where {N} = x
to_svectors(x::AbstractVector{<:AbstractVector{T}}, ::Val{N}=Val(length(first(x)))) where {T<:Number,N} =
    SVector{N,T}.(x)
to_svectors(x::AbstractMatrix{T}, ::Val{N}=Val(size(x,2))) where {T<:Number,N} =
    SVector{N,T}[row for row in eachrow(x)]

# normalize x and y arguments to vectors of svectors or scalars
chebregression(x::AbstractVecOrMat, y::AbstractVecOrMat, lb::AbstractVector, ub::AbstractVector, order::NTuple{N}) where {N} =
    chebregression(to_svectors(x, Val{N}()), to_svectors(y), SVector{N}(lb), SVector{N}(ub), order)

chebregression(x::AbstractVecOrMat, y::AbstractVecOrMat, order::NTuple{N}) where {N} =
    chebregression(to_svectors(x, Val{N}()), to_svectors(y), order)

# accept scalar bounds and order in 1d case
chebregression(x::AbstractVector{<:Real}, y::AbstractVecOrMat, lb::Real, ub::Real, order::Integer) =
    chebregression(x, y, SVector(lb), SVector(ub), (order,))
chebregression(x::AbstractVector{<:Real}, y::AbstractVecOrMat, order::Integer) =
    chebregression(x, y, minimum(x), maximum(x), order)

# construct lb and ub if omitted
chebregression(x::AbstractVector{<:SVector{N}}, y::AbstractVector, order::NTuple{N}) where {N} =
    chebregression(x, y, reduce((a,b) -> min.(a,b), x), reduce((a,b) -> max.(a,b), x), order)

# promote arguments to common types
function chebregression(x::AbstractVector{SVector{N,Tx}}, y::AbstractVector{Ty},
    lb::SVector{N,Tlb}, ub::SVector{N,Tub}, order::NTuple{N,<:Integer}) where {N,Tx<:Real,Tlb<:Real,Tub<:Real,Ty<:Union{SVector,Number}}
    Td = float(promote_type(Tx,Tub,Tlb))
    return chebregression(AbstractVector{SVector{N,Td}}(x), y, SVector{N,Td}(lb), SVector{N,Td}(ub), NTuple{N,Int}(order))
end

"""
    chebregression(x, y, [lb, ub,] order)

Return a Chebyshev polynomial (`ChebPoly`) constructed by
performing a least-square fit of Chebyshev polnomials of the
given `order`, where `x` are the coordinates of the data
points `y`.  `lb` and `ub` are the lower and upper bounds,
respectively, of the Chebyshev domain; these should normally
enclose all of the points in `x`, and default to the minimum
and maximum coordinates in `x` if they are omitted.

In the 1d case, `x` is an array of scalars, `lb < ub`
are scalars, and `order` is an integers.   In the `N`-dimensional
case, `order` is an `N`-tuple of integers (the order in each
dimension), `lb` and `ub` are `N`-component vectors, and
`x` is an array of `N`-component vectors (or a matrix with
`N` columns, interpreted as the vector components).

`y` can be a vector or numbers or a vector of vectors (for vector-
valued Chebyshev fits).  The latter case can also be input
as a matrix whose columns are the vector componnents.  `size(x,1)`
and `size(y,1)` must match, and must exceed `prod(order .+ 1)`.
"""
function chebregression end


text:
using ImageFiltering, ImageCore
using PkgBenchmark
using BenchmarkTools
using Statistics: quantile, mean, median!

function makeimages(sz)
    imgF32      = rand(Float32, sz)
    imgN0f8     = Array(rand(N0f8, sz))
    imggrayF32  = Array(rand(Gray{Float32}, sz))
    imggrayN0f8 = Array(rand(Gray{N0f8}, sz))
    imgrgbF32   = Array(rand(RGB{Float32}, sz))
    imgrgbN0f8  = Array(rand(RGB{N0f8}, sz))
    return ("F32"=>imgF32, "N0f8"=>imgN0f8, "GrayF32"=>imggrayF32,
            "GrayN0f8"=>imggrayN0f8, "RGBF32"=>imgrgbF32, "RGBN0f8"=>imgrgbN0f8)
end

sz2str(sz) = join(map(string, sz), '×')

SUITE = BenchmarkGroup()
SUITE["mapwindow"] = BenchmarkGroup()

let grp = SUITE["mapwindow"]
    img1d = randn(1000)
    img2d = randn(30,30)
    img3d = randn(10,11,12)
    grp["cheap f, tiny window"] = @benchmarkable mapwindow(first, $img1d, (1,))
    grp["extrema"] = @benchmarkable mapwindow(extrema, $img2d, (5,5))
    grp["maximum"] = @benchmarkable mapwindow(maximum, $img2d, (5,5))
    grp["minimum"] = @benchmarkable mapwindow(minimum, $img2d, (5,5))
    grp["median!"] = @benchmarkable mapwindow(median!, $img2d, (5,5))
    grp["mean, small window"] = @benchmarkable mapwindow(mean, $img1d, (3,))
    grp["mean, large window"] = @benchmarkable mapwindow(mean, $img3d, (5,5,5))
    grp["expensive f"] = @benchmarkable mapwindow(x -> quantile(vec(x), 0.7), $img3d, (3,3,3))
end

SUITE["imfilter"] = BenchmarkGroup()
let grp = SUITE["imfilter"]
    kerninsep   = (centered([-1, 0, 1]),
                   centered([ 1/5  1/4  1/7;
                              1/2  1/3 -1/11;
                             -1/25 1/9 -1/7]),   # has full rank so won't be factored
                   centered(rand(3, 3, 3)))
    for sz in ((100, 100), (2048, 2048), (2048,), (100, 100, 100))
        for (aname, img) in makeimages(sz)
            trues = map(i->true, sz)
            twos  = map(i->2, sz)
            szstr = sz2str(sz)
            kerniir = KernelFactors.IIRGaussian(map(i->10.0f0, sz))
            kerng   = KernelFactors.gaussian(map(i->10.0f0, sz))
            for (kname, kern) in zip(("densesmall", "denselarge", "factoredsmall", "factoredlarge"),
                                      ((kerninsep[length(sz)],), (Kernel.DoG(twos),), KernelFactors.sobel(trues, 1), kerng))
                grp[kname*"_"*aname*"_"*szstr]  = @benchmarkable imfilter($img, $kern, "replicate", ImageFiltering.FIR())
            end
            grp["IIRGaussian_"*aname*"_"*szstr] = @benchmarkable imfilter($img, $kerniir, "replicate", ImageFiltering.IIR())
            grp["FFT_"*aname*"_"*szstr]         = @benchmarkable imfilter($img, $(Kernel.DoG(twos),), "replicate", ImageFiltering.FFT())
        end
    end
end


text:
using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CovidSurvey
using CSV
using DataFrames
using Dates
# using RCall
# using Turing
using StatsPlots
using Plots.PlotMeasures
using Random
using BSON
using PrettyTables
using Turing
plotlyjs()
## ===========================================================================
# load data and model
# ============================================================================
fdir = projectdir("reports/parametric-model-on-2nd-wave-newdata/")
# fname = "cases=2020-06-01_chains=10_hospit=true_model=parametric-cases_observ=2021-01-13_preds=CF,CC,CR,CS_sero=true_steps=2000_warmup=1000"
fname = "cases=2020-06-01_chains=10_hospit=true_model=parametric-cases_observ=2021-01-13_preds=CF,CC,CR,CS_sero=true_steps=2000_warmup=1000"
PREFIX = "RAW-"
epimodel = CovidSurvey.National.model_parametric_cases
#-----------------------------------------------------------------------------
# read params
_, ps, _ = parse_savename(fname, parsetypes = (Int, Float64, Bool))
!("preds" in keys(ps)) && (ps["preds"] = nothing)
# ps = NamedTuple{Tuple(Symbol.(keys(ps)))}(values(ps))

#------------------------------------------------------------------------------
# read data
data = National.load_data(
      ps["observ"]
    , ps["preds"]|> CovidSurvey.parse_predictors
    , ps["cases"]
    , ps["hospit"]
    , ps["sero"]
    , iar_step = 1
    , covariates_kwargs = Dict(
        :fname => projectdir("data", "inferred_contact_rates.csv"),
        #:fname => projectdir("data", "smoothed_contact_rates.csv"), #normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" ),
        # :fname => normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" ),
        :shift => -1,
        :startdate => "2020-11-10",
        # :enddate => "2021-01-13"
    )
)
turing_data = data.turing_data;

#-----------------------------------------------------------------------------
# read chain
function not(c::Chains, i)
    n = size(c, 3)
    s = filter(x->x∉i, 1:n)
    return c[:,:,s]
end
chain = read(fdir*PREFIX*"CHAIN_"*fname*".jls", Chains)
chain = chain[1:10:end,:,:]
chain = not(chain, [2,4])
# fname = normpath( fdir, savename(PREFIX*"CHAIN", ps, "jls") )
# safesave( fname, chain)

#-----------------------------------------------------------------------------
# plot chain
n = filter( x->!occursin(r"(?<!α)\[", x), String.(names(chain)))
let
    p = plot(chain[n])
    saveto = normpath( fdir, PREFIX*"CHAINSPLOT_"*fname*".html")
    savefig(p, saveto )
    run(`firefox $(saveto)`, wait=false)
end

## ===========================================================================
# plot results
# ============================================================================
# make predictions
m_pred = epimodel(turing_data..., true; link=KLogistic(3.), invlink=KLogit(3.))
gq = Turing.generated_quantities(m_pred, chain)
generated_posterior = vectup2tupvec(reshape(gq, length(gq)))

posterior = let
    m_pred = epimodel(turing_data..., true; link=KLogistic(3.), invlink=KLogit(3.) )
    gq = Turing.generated_quantities(m_pred, chain)
    vectup2tupvec( reshape(gq, length(gq)) )
end;
# observed_daily_cases, expected_daily_deaths, expected_daily_hospit, Rt, expected_seropositive, iar = posterior;

#-----------------------------------------------------------------------------
@info "store generated quantities"
let
    dic = Dict( zip( keys(posterior), values(posterior) ) )
    saveto = normpath( fdir, PREFIX*"GENERATED-QUANTITIES_"*fname*".bson")
    bson( saveto ,  dic )
end
#-----------------------------------------------------------------------------
@info "store reproduction number"
rt = let
    rt = let
        ib = findfirst(==(Date("2020-05-15")), data.dates)
        ie = findfirst(==(Date(ps["observ"])), data.dates)
        Rt_array = hcat(posterior.Rt...)[ib:ie,:]

        qs = [quantile(v, [0.025, 0.25, 0.5, 0.75, 0.975]) for v in eachrow(Rt_array)]
        llq, lq, mq, uq, uuq = (eachrow(hcat(qs...))..., )

        date = data.dates[ib:ie]
        DataFrame((;date, llq, lq, mq, uq, uuq))
    end
    saveto = normpath( fdir, PREFIX*"Rt_"*fname*".csv")
    save(saveto, rt)
    rt
end
#-----------------------------------------------------------------------------
@info "plot results"
let
    p = National.plot_results(data, posterior...);
    saveto = normpath( fdir, PREFIX*"PREDICTION_"*fname*".html")
    savefig(p, saveto )
    # savefig(p, projectdir("figures", fname*".png") )
    run(`firefox $(saveto)`, wait=false)
end
#-----------------------------------------------------------------------------
# perform diagnostics
if ps["chains"] > 1
    @info "gelman diagnostics"
    diagnostics = gelmandiag(chain)
    saveto = normpath( fdir, PREFIX*"GELMANDIAG_"*fname*".csv")
    safesave( saveto, diagnostics)
    pretty_table(diagnostics; crop=:none)
end
#-----------------------------------------------------------------------------
@info "meanplot"
let
    p = meanplot(chain[n]);
    saveto = normpath( fdir, PREFIX*"MEANPLOT_"*fname*".html")
    savefig(p, saveto )
    run(`firefox $(saveto)`, wait=false)
end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
using POMDPs
using RockSample
using SARSOP # load a  POMDP Solver
using POMDPGifs # to make gifs
import Cairo, Fontconfig

pomdp = RockSamplePOMDP(rocks_positions=[(2,3), (4,4), (4,2)],
                        sensor_efficiency=10.0,
                        discount_factor=0.95,
                        good_rock_reward = 20.0)

solver = SARSOPSolver(precision=1e-3)

policy = solve(solver, pomdp)

sim = GifSimulator(filename="out.gif", max_steps=30)
simulate(sim, pomdp, policy);


text:
#mutable struct UncertainVariableRef <: RobustVariableRef
#end

struct Uncertain <: RobustVariableType
    stage::Int
end

struct UncertainVariable<: RobustVariable
    info::JuMP.VariableInfo
    stage::Int
end
export Uncertain, UncertainVariable


text:
function texture_id!(fg::FrameGraph, arg::Texture, pass)::UInt32
    # let's create the resource eagerly for now
    img = image(fg, arg.name)
    (; frame, resource_graph) = fg
    if !isnothing(arg.sampling)
        # combined image sampler
        sampling = arg.sampling
        sampler = Vk.Sampler(device(fg), sampling)
        # preserve sampler
        register(frame, gensym(:sampler), sampler; persistent = false)
        combined_image_sampler_state = frame.gd.resources.gset.state[Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER]
        push!(combined_image_sampler_state, Vk.DescriptorImageInfo(sampler, View(img), image_layout(resource_graph, fg.resources[arg.name], pass)))
        length(combined_image_sampler_state) - 1
    else
        # sampled image
        sampler = empty_handle(Vk.Sampler)
        sampled_image_state = frame.gd.resources.gset.state[Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE]
        push!(sampled_image_state, Vk.DescriptorImageInfo(sampler, View(img), image_layout(resource_graph, fg.resources[arg.name], pass)))
        length(sampled_image_state) - 1
    end
end

function sampler_id!(fg::FrameGraph, arg::Sampling)
    (; frame) = fg
    view = empty_handle(Vk.ImageView)
    sampler = Vk.Sampler(device(fg), sampling)
    # preserve sampler
    register(frame, gensym(:sampler), sampler; persistent = false)
    sampler_state = frame.gd.resources.gset.state[Vk.DESCRIPTOR_TYPE_SAMPLER]
    push!(sampler_state, Vk.DescriptorImageInfo(sampler, view, Vk.IMAGE_LAYOUT_UNDEFINED))
    length(sampler_state) - 1
end

#=

"""
Encode a lifetime context that allows the deletion of staged resources once expired.
"""
abstract type LifeTimeContext end

"""
Period of time that extends for a whole frame.
Work between frames is assumed to be independent.
"""
struct FrameContext <: LifeTimeContext end

struct PersistentContext <: LifeTimeContext end

struct Resources
    descriptors::ResourceDescriptors

    "Number of frames during which resources were not used."
    unused::Dictionary{Int,Vector{Int}}
    "Persistent resources, never to be garbage-collected."
    persistent::Vector{Int}
end

"""
Transition from one frame to another.
"""
function transition(resources::Resources)
    for resource in resources
end

struct Frame
    fg::FrameGraph
    data::GlobalData
    resources::Resources
    fences::Vector{Vk.Fence}
end

=#


text:
using Test
using ClimaCore.DataLayouts
using StaticArrays
using ClimaCore.DataLayouts: get_struct, set_struct!

TestFloatTypes = (Float32, Float64)

@testset "VF" begin
    for FT in TestFloatTypes
        S = Tuple{Complex{FT}, FT}
        array = rand(FT, 4, 3)

        data = VF{S}(array)
        @test getfield(data.:1, :array) == @view(array[:, 1:2])

        # test tuple assignment
        data[1] = (Complex{FT}(-1.0, -2.0), FT(-3.0))
        @test array[1, 1] == -1.0
        @test array[1, 2] == -2.0
        @test array[1, 3] == -3.0

        # sum of all the first field elements
        @test sum(data.:1) ≈ Complex{FT}(sum(array[:, 1]), sum(array[:, 2])) atol =
            10eps()

        @test sum(x -> x[2], data) ≈ sum(array[:, 3]) atol = 10eps()
    end
end

@testset "VF type safety" begin
    Nv = 1 # number of vertical levels

    # check that types of the same bitstype throw a conversion error
    SA = (a = 1.0, b = 2.0)
    SB = (c = 1.0, d = 2.0)

    array = zeros(Float64, Nv, 2)
    data = VF{typeof(SA)}(array)

    data[1] = SA
    @test data[1] isa typeof(SA)
    @test_throws MethodError data[1] = SB
end

@testset "broadcasting between 1D data objects and scalars" begin
    for FT in TestFloatTypes
        data1 = ones(FT, 2, 2)
        S = Complex{FT}
        data1 = VF{S}(data1)
        res = data1 .+ 1
        @test res isa VF
        @test parent(res) == FT[2.0 1.0; 2.0 1.0]
        @test sum(res) == Complex{FT}(4.0, 2.0)
        @test sum(Base.Broadcast.broadcasted(+, data1, 1)) ==
              Complex{FT}(4.0, 2.0)
    end
end

@testset "broadcasting 1D assignment from scalar" begin
    for FT in TestFloatTypes
        S = Complex{FT}
        data = VF{S}(Array{FT}, 3)
        data .= Complex{FT}(1.0, 2.0)
        @test parent(data) == FT[1.0 2.0; 1.0 2.0; 1.0 2.0]
        data .= 1
        @test parent(data) == FT[1.0 0.0; 1.0 0.0; 1.0 0.0]
    end
end

@testset "broadcasting between 1D data objects" begin
    for FT in TestFloatTypes
        data1 = ones(FT, 2, 2)
        data2 = ones(FT, 2, 1)
        S1 = Complex{FT}
        S2 = FT
        data1 = VF{S1}(data1)
        data2 = VF{S2}(data2)
        res = data1 .+ data2
        @test res isa VF{S1}
        @test parent(res) == FT[2.0 1.0; 2.0 1.0]
        @test sum(res) == Complex{FT}(4.0, 2.0)
    end
end

#= TODO
@testset "broadcasting 1D data complicated function" begin
    for FT in TestFloatTypes
        S1 = NamedTuple{(:a, :b), Tuple{Complex{FT}, FT}}
        data1 = ones(FT, 2, 2)
        S2 = NamedTuple{(:c,), Tuple{FT,}}
        data2 = ones(FT, 2, 1)
        data1 = VF{S1}(data1)
        data2 = VF{S2}(data2)

        f(a1, a2) = a1.a.re * a2.c + a1.b
        res = f.(data1, data2)
        @test res isa VF[FT]
    end
end
=#


text:
using JuMP, GLPK
using DataFrames, CSV
using LinearAlgebra
using JSON
using LightGraphs

cd(@__DIR__)
include("TSP_FUNCTIONS.jl")

# Read Data Input from training stage

#route_opt_seq_df = DataFrame(CSV.File("../submission/data/build_route_with_seq.csv"));
route_opt_seq_df = DataFrame(CSV.File("../data/model_apply_outputs/model_apply_output/build_route_with_seq.csv"));
zone_travel_time_matrix = JSON.parsefile("../data/model_apply_outputs/model_apply_output/zone_mean_travel_times.json")
stop_travel_time_matrix = JSON.parsefile("../data/model_apply_inputs/new_travel_times.json");
zone_min_tt_nearest_df = DataFrame(CSV.File("../data/model_apply_outputs/model_apply_output/zone_min_tt_df_nearest_9.csv"));
zone_ecu_tt_nearest_df = DataFrame(CSV.File("../data/model_apply_outputs/model_apply_output/zone_ecu_dist_df_nearest_9.csv"));
route_id_list = unique(route_opt_seq_df.route_id);

β = 3.8; γ = 2.5; α = 1.04

# Calculate Zone sequence
Tour_TSP_zone_seq = Dict()
for route_id in route_id_list

    route_df = route_opt_seq_df[route_opt_seq_df.route_id .== route_id, :]
    zone_list = unique(route_df.zone_id)
    zone_distance_matrix = zone_travel_time_matrix[route_id]

    route_min_tt_nearest = zone_min_tt_nearest_df[zone_min_tt_nearest_df.route_id .== route_id, :]
    route_ecu_tt_nearest = zone_ecu_tt_nearest_df[zone_ecu_tt_nearest_df.route_id .== route_id, :]
    adj_zone_distance_matrix = adjust_zone_dist_matrix(zone_distance_matrix, zone_list, route_min_tt_nearest, route_ecu_tt_nearest, β, γ, α)

    tour_tsp_zone_seq = calculate_tsp_seq(route_df, adj_zone_distance_matrix)
    Tour_TSP_zone_seq[route_id] = tour_tsp_zone_seq
end

# Output best zone sequence
open("../data/model_apply_outputs/model_apply_output/mean_dist/opt_zone_seq_tour.json","w") do f
    JSON.print(f, Tour_TSP_zone_seq)
end

tsp_opt_seq = Dict()
for route_id in route_id_list

    route_df = route_opt_seq_df[route_opt_seq_df.route_id .== route_id, :]
    stop_distance_matrix = stop_travel_time_matrix[route_id]

    TOUR_TSP_zone_sequence = Tour_TSP_zone_seq[route_id]
    TOUR_TSP_path = generate_complete_seq_multiple(TOUR_TSP_zone_sequence, route_df, stop_distance_matrix)

    tsp_opt_seq[route_id] = TOUR_TSP_path
end

# Output best complete sequence
open("../data/model_apply_outputs/model_apply_output/mean_dist/opt_complete_seq_tour.json","w") do f
    JSON.print(f, tsp_opt_seq)
end

println("Successfully generate proposed sequences, into the post-processing stge")


text:
struct PTDF{Ax, L <: NTuple{2, Dict}} <: PowerNetworkMatrix{Float64}
    data::Array{Float64, 2}
    axes::Ax
    lookup::L
end

function _buildptdf(branches, nodes, dist_slack::Array{Float64} = [0.1])

    buscount = length(nodes)
    linecount = length(branches)
    num_bus = Dict{Int32, Int32}()

    for (ix, b) in enumerate(nodes)
        num_bus[get_number(b)] = ix
    end

    A = zeros(Float64, buscount, linecount)
    inv_X = zeros(Float64, linecount, linecount)

    #build incidence matrix
    #incidence_matrix = A

    for (ix, b) in enumerate(branches)
        if isa(b, DCBranch)
            @warn("PTDF construction ignores DC-Lines")
            continue
        end

        A[num_bus[get_number(get_from(get_arc(b)))], ix] = 1

        A[num_bus[get_number(get_to(get_arc(b)))], ix] = -1

        if isa(b, Transformer2W)
            inv_X[ix, ix] = 1 / get_x(b)

        elseif isa(b, TapTransformer)
            inv_X[ix, ix] = 1 / (get_x(b) * get_tap(b))

        elseif isa(b, Line)
            inv_X[ix, ix] = 1 / get_x(b)

        elseif isa(b, PhaseShiftingTransformer)
            y = 1 / (get_r(b) + get_x(b) * 1im)
            y_a = y / (get_tap(b) * exp(get_α(b) * 1im * (π / 180)))
            inv_X[ix, ix] = 1 / imag(y_a)
        end

    end

    slacks = [num_bus[get_number(n)] for n in nodes if get_bustype(n) == BusTypes.REF]
    slack_position = slacks[1]
    B = gemm(
        'N',
        'T',
        gemm('N', 'N', A[setdiff(1:end, slack_position), 1:end], inv_X),
        A[setdiff(1:end, slack_position), 1:end],
    )
    if dist_slack[1] == 0.1 && length(dist_slack) == 1
        (B, bipiv, binfo) = getrf!(B)
        S_ = gemm(
            'N',
            'N',
            gemm('N', 'T', inv_X, A[setdiff(1:end, slack_position), :]),
            getri!(B, bipiv),
        )
        S = hcat(S_[:, 1:(slack_position - 1)], zeros(linecount), S_[:, slack_position:end])

    elseif dist_slack[1] != 0.1 && length(dist_slack) == buscount
        @info "Distributed bus"
        (B, bipiv, binfo) = getrf!(B)
        S_ = gemm(
            'N',
            'N',
            gemm('N', 'T', inv_X, A[setdiff(1:end, slack_position), :]),
            getri!(B, bipiv),
        )
        S = hcat(S_[:, 1:(slack_position - 1)], zeros(linecount), S_[:, slack_position:end])
        slack_array = dist_slack / sum(dist_slack)
        slack_array = reshape(slack_array, buscount, 1)
        S = S - gemm('N', 'N', gemm('N', 'N', S, slack_array), ones(1, buscount))

    elseif length(slack_position) == 0
        @warn("Slack bus not identified in the Bus/Nodes list, can't build PTDF")
        S = Array{Float64, 2}(undef, linecount, buscount)
    end

    return S, A

end

function PTDF(branches, nodes, dist_slack::Array{Float64} = [0.1])

    #Get axis names
    line_ax = [get_name(branch) for branch in branches]
    bus_ax = [get_number(bus) for bus in nodes]
    S, A = _buildptdf(branches, nodes, dist_slack)

    axes = (line_ax, bus_ax)
    look_up = (_make_ax_ref(line_ax), _make_ax_ref(bus_ax))

    return PTDF(S, axes, look_up)

end

function PTDF(sys::System, dist_slack::Array{Float64} = [0.1])
    branches = get_components(ACBranch, sys)
    nodes = get_components(Bus, sys)

    return PTDF(branches, nodes, dist_slack)

end


text:
# cat
A1data, A2data = [1 3; 2 4], [5 7; 6 8]

A1 = AxisArray(A1data, Axis{:Row}([:First, :Second]), Axis{:Col}([:A, :B]))
A2 = AxisArray(A2data, Axis{:Row}([:Third, :Fourth]), Axis{:Col}([:A, :B]))
@test isa(cat(A1, A2, dims=1), AxisArray)
@test cat(A1, A2, dims=1) == AxisArray(vcat(A1data, A2data),
                                  Axis{:Row}([:First, :Second, :Third, :Fourth]), Axis{:Col}([:A, :B]))

A2 = AxisArray(A2data, Axis{:Row}([:First, :Second]), Axis{:Col}([:C, :D]))
@test isa(cat(A1, A2, dims=2), AxisArray)
@test cat(A1, A2, dims=2) == AxisArray(hcat(A1data, A2data),
                                  Axis{:Row}([:First, :Second]), Axis{:Col}([:A, :B, :C, :D]))

A2 = AxisArray(A2data, Axis{:Row}([:First, :Second]), Axis{:Col}([:A, :B]))
@test isa(cat(A1, A2, dims=3), AxisArray)
@test cat(A1, A2, dims=3) == AxisArray(cat(A1data, A2data, dims=3),
                                       Axis{:Row}([:First, :Second]), Axis{:Col}([:A, :B]),
                                       Axis{:page}(1:2))

A1 = AxisArray(A1data, :Row, :Col)
A2 = AxisArray(A2data, :Row, :Col)
@test_throws ArgumentError cat(A1, A2, dims=2)
@test cat(A1, A2, dims=3) == AxisArray(cat(A1data, A2data, dims=3), :Row, :Col)

# merge
Adata, Bdata, Cdata = randn(4,4,2), randn(4,4,2), randn(4,4,2)
A = AxisArray(Adata, Axis{:X}([1,2,3,4]), Axis{:Y}([10.,20,30,40]), Axis{:Z}([:First, :Second]))
B = AxisArray(Bdata, Axis{:X}([3,4,5,6]), Axis{:Y}([30.,40,50,60]), Axis{:Z}([:First, :Second]))

ABdata = zeros(6,6,2)
ABdata[1:4,1:4,:] = Adata
ABdata[3:6,3:6,:] = Bdata
@test merge(A,B) == AxisArray(ABdata, Axis{:X}([1,2,3,4,5,6]), Axis{:Y}([10.,20,30,40,50,60]), Axis{:Z}([:First, :Second]))
@test_throws ErrorException AxisArrays.axismerge(:notouter, Axis{:X}([1,2,3,4]), Axis{:X}([1,2,3,4]))

AC = AxisArray(cat(Adata, Cdata, dims=3), :X, :Y, :Z)
B2 = AxisArray(Bdata, :X, :Y, :Z)
@test merge(AC,B2) == AxisArray(cat(Bdata, Cdata, dims=3), :X, :Y, :Z)

# join
ABdata = zeros(6,6,2,2)
ABdata[1:4,1:4,:,1] = Adata
ABdata[3:6,3:6,:,2] = Bdata
@test join(A,B) == AxisArray(ABdata, Axis{:X}([1,2,3,4,5,6]), Axis{:Y}([10.,20,30,40,50,60]), Axis{:Z}([:First, :Second]))
@test join(A,B, newaxis=Axis{:JoinAxis}([:A, :B])) == AxisArray(ABdata, Axis{:X}([1,2,3,4,5,6]), Axis{:Y}([10.,20,30,40,50,60]), Axis{:Z}([:First, :Second]), Axis{:JoinAxis}([:A, :B]))
@test join(A,B,method=:inner) == AxisArray(ABdata[3:4, 3:4, :, :], Axis{:X}([3,4]), Axis{:Y}([30.,40]), Axis{:Z}([:First, :Second]))
@test join(A,B,method=:left) == AxisArray(ABdata[1:4, 1:4, :, :], A.axes...)
@test join(A,B,method=:right) == AxisArray(ABdata[3:6, 3:6, :, :], B.axes...)
@test join(A,B,method=:outer) == join(A,B)

# collapse
A1 = AxisArray(A1data, Axis{:X}(1:2), Axis{:Y}(1:2))
A2 = AxisArray(reshape(A2data, size(A2data)..., 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:Z}([:foo]))

@test @inferred(collapse(Val(2), A1, A2)) == AxisArray(cat(A1data, A2data, dims=3), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:collapsed}(AxisArrays.CategoricalVector([(1,), (2, :foo)])))
@test @inferred(collapse(Val(2), A1)) == AxisArray(reshape(A1, 2, 2, 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:collapsed}(AxisArrays.CategoricalVector([(1,)])))
@test @inferred(collapse(Val(2), A1)) == AxisArray(reshape(A1.data, size(A1)..., 1), AxisArrays.axes(A1)..., Axis{:collapsed}(AxisArrays.CategoricalVector([(1,)])))

@test @inferred(collapse(Val(2), (:A1, :A2), A1, A2)) == AxisArray(cat(A1data, A2data, dims=3), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:collapsed}(AxisArrays.CategoricalVector([(:A1,), (:A2, :foo)])))
@test @inferred(collapse(Val(2), (:foo,), A1)) == AxisArray(reshape(A1, 2, 2, 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:collapsed}(AxisArrays.CategoricalVector([(:foo,)])))
@test @inferred(collapse(Val(2), (:a,), A1)) == AxisArray(reshape(A1.data, size(A1)..., 1), AxisArrays.axes(A1)..., Axis{:collapsed}(AxisArrays.CategoricalVector([(:a,)])))

@test @inferred(collapse(Val(0), A1)) == AxisArray(vec(A1data), Axis{:collapsed}(AxisArrays.CategoricalVector(vec(collect(Iterators.product((1,), axisvalues(A1)...))))))
@test @inferred(collapse(Val(1), A1)) == AxisArray(A1data, Axis{:row}(1:2), Axis{:collapsed}(AxisArrays.CategoricalVector(vec(collect(Iterators.product((1,), axisvalues(A1)[2]))))))
@test @inferred(collapse(Val(1), (1,), A1)) == collapse(Val(1), A1)
@test @inferred(collapse(Val(1), Array{Int, 2}, A1)) == collapse(Val(1), A1)
@test @inferred(collapse(Val(1), Array{Int, 2}, (1,), A1)) == collapse(Val(1), A1)

@test_throws ArgumentError collapse(Val(-1), A1)
@test_throws ArgumentError collapse(Val(10), A1)

A1ᵀ = transpose(A1)
@test_throws ArgumentError collapse(Val(-1), A1, A1ᵀ)
@test_throws ArgumentError collapse(Val(1), A1, A1ᵀ)
@test_throws ArgumentError collapse(Val(10), A1, A1ᵀ)


text:
"""
    StateSpaceFunction(f!, h; kwargs...)

# Examples
```julia
x = Float64[1, 2]
dx = similar(x)
p = 5

f!(dx, x, p, t; u, err) = (dx .= p*x .+ err)
h(x, p, t; u=1000, err) = x .- 3u
ssf = StateSpaceFunction(f!, h; u=(x,p,t)->cos(t), err=0.0)
```
"""
struct StateSpaceFunction{F,H,NT} <: Function
    f!::F
    h::H
    inputs::NT
end
StateSpaceFunction(f!, h; inputs...) = StateSpaceFunction(f!, h, (;inputs...))

function (ssf::StateSpaceFunction)(dx, args...; kwargs...)
    applied_kwargs = map(f -> _maybe_apply(f, args...), (;ssf.inputs..., kwargs...))
    ssf.f!(dx, args...; applied_kwargs...)
    return ssf.h(args...; applied_kwargs...)
end


text:
function copper_plate(psi_container::PSIContainer, expression::Symbol, bus_count::Int)

    time_steps = model_time_steps(psi_container)
    remove_undef!(psi_container.expressions[expression])

    constraint_val = JuMPConstraintArray(undef, time_steps)
    assign_constraint!(psi_container, "CopperPlateBalance", constraint_val)

    for t in time_steps
        constraint_val[t] = JuMP.@constraint(
            psi_container.JuMPmodel,
            sum(psi_container.expressions[expression].data[1:bus_count, t]) == 0
        )
    end

    return

end


text:
# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

"""
    rad2sec(rad) -> seconds

### Purpose ###

Convert from radians to seconds.

### Argument ###

* `rad`: number of radians.

### Output ###

The number of seconds corresponding to `rad`.

### Example ###

```jldoctest
julia> using AstroLib

julia> rad2sec(1)
206264.80624709636
```

### Notes ###

Use `sec2rad` to convert seconds to radians.
"""
rad2sec(rad::Real) = rad2deg(rad) * 3600


text:
using MelGeneralizedCepstrums
using SynthesisFilters
using Compat
using Base.Test

@static is_linux() ? include("sptk.jl") : nothing

type TestSynthesisFilter <: SynthesisFilter
end

let
    @test_throws ErrorException to_filtcoef(TestSynthesisFilter(), rand(10))
    @test_throws ErrorException filt!(TestSynthesisFilter(), 1.0, rand(10))
end

function test_synthesis_one_frame(f::SynthesisFilter, order)
    srand(98765)
    excite = rand(1024)
    previous_b = rand(order+1)
    current_b = rand(order+1)

    r = synthesis_one_frame!(f, excite, previous_b, current_b)
    @test any(!isnan(r))
end

function test_poledf_synthesis(order::Int, hopsize::Int)
    srand(98765)
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    l = estimate(LinearPredictionCoef(order), x)

    f = AllPoleDF(order)
    r = synthesis!(f, excite, l, hopsize)
    @test any(!isnan(r))
end

function test_synthesis(def::SpectralParam, hopsize)
    srand(98765)
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))

    order = param_order(def)

    if isa(def, LineSpectralPair)
        state = lpc2lsp(estimate(LinearPredictionCoef(order), x))
    elseif isa(def, PartialAutoCorrelation)
        state = lpc2par(estimate(LinearPredictionCoef(order), x))
    else
        state = estimate(def, x)
    end

    r = synthesis(excite, state, hopsize)
    @test any(!isnan(r))
end

function test_poledf_exception()
    srand(98765)
    N = 512
    order = 25
    hopsize = 80
    excite = rand(N)
    x = rand(N, div(N, hopsize))

    l = estimate(LinearPredictionCoef(order), x)
    f = AllPoleDF(order)
    synthesis!(f, excite, l, hopsize)

    for def in [LinearCepstrum(order),
                GeneralizedCepstrum(order, -0.1),
                MelCepstrum(order, 0.41),
                MelGeneralizedCepstrum(order, 0.41, -0.1)
                ]
        l = estimate(def, x)
        @test_throws ErrorException synthesis!(f, excite, l, hopsize)
    end
end

function test_ltcdf_synthesis(order::Int, hopsize::Int)
    srand(98765)
    N = 1024
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    par = lpc2par(estimate(LinearPredictionCoef(order), x))

    f = AllPoleLatticeDF(order)
    r = synthesis!(f, excite, par, hopsize)
    @test any(!isnan(r))
end

function test_lspdf_synthesis(order::Int, hopsize::Int)
    srand(98765)
    N = 1024
    excite = rand(N)
    x = rand(order+1, div(N, hopsize))
    lsp = lpc2lsp(estimate(LinearPredictionCoef(order), x))

    f = LSPDF(order)
    r = synthesis!(f, excite, lsp, hopsize)
    @test any(!isnan(r))
end

function test_lmadf_synthesis(order::Int, pade::Int, hopsize::Int)
    srand(98765)
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    mgc = estimate(LinearCepstrum(order), x)

    f = LMADF(order; pade=pade)
    r = synthesis!(f, excite, mgc, hopsize)
    @test any(!isnan(r))
end

function test_lmadf_exception()
    srand(98765)
    N = 512
    order = 25
    hopsize = 80
    excite = rand(N)
    x = rand(N, div(N, hopsize))

    mgc = estimate(LinearCepstrum(order), x)
    f = LMADF(order)
    synthesis!(f, excite, mgc, hopsize)

    for def in [
                MelCepstrum(order, 0.41),
                GeneralizedCepstrum(order, -0.01),
                MelGeneralizedCepstrum(order, 0.41, -0.01),
                AllPoleCepstrum(order)
                ]
        mgc = estimate(def, x)
        @test_throws ErrorException synthesis!(f, excite, mgc, hopsize)
    end

    ## pade approximation
    @test_throws ErrorException LMADF(order, pade=3)
    try LMADF(order; pade=4); catch @test false; end
    try LMADF(order; pade=5); catch @test false; end
    @test_throws ErrorException LMADF(order, pade=6)
end

function test_mlsadf_synthesis(order::Int, α::Float64, pade::Int, hopsize::Int)
    srand(98765)
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    mgc = estimate(MelCepstrum(order, α), x)

    f = MLSADF(order, α; pade=pade)
    r = synthesis!(f, excite, mgc, hopsize)
    @test any(!isnan(r))
end

function test_mlsadf_exception()
    srand(98765)
    N = 512
    order = 25
    hopsize = 80
    excite = rand(N)
    x = rand(N, div(N, hopsize))

    # should accept LinearCepstrum
    mgc = estimate(LinearCepstrum(order), x)
    f = MLSADF(order, 0.0)
    synthesis!(f, excite, mgc, hopsize)

    mgc = estimate(MelCepstrum(order, 0.41), x)
    f = MLSADF(order, 0.41)
    synthesis!(f, excite, mgc, hopsize)

    for def in [
                GeneralizedCepstrum(order, -0.01),
                MelGeneralizedCepstrum(order, 0.41, -0.01),
                AllPoleCepstrum(order)
                ]
        mgc = estimate(def, x)
        @test_throws ErrorException synthesis!(f, excite, mgc, hopsize)
    end

    ## pade approximation
    @test_throws ErrorException MLSADF(order, 0.41; pade=3)
    try MLSADF(order, 0.41; pade=4); catch @test false; end
    try MLSADF(order, 0.41; pade=5); catch @test false; end
    @test_throws ErrorException MLSADF(order, 0.41; pade=6)
end

function test_mglsadf_synthesis(order::Int, α::Float64, ns::Int, hopsize::Int)
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    mgc = estimate(MelGeneralizedCepstrum(order, α, -1/ns), x)

    f = MGLSADF(order, α, ns)
    r = synthesis!(f, excite, mgc, hopsize)
    @test any(!isnan(r))
end

function test_mglsadf_exceptions()
    N = 512
    order = 25
    hopsize = 80
    excite = rand(N)
    x = rand(N, div(N, hopsize))

    ns = 10.5
    α = 0.41
    mgc = estimate(MelGeneralizedCepstrum(order, α, -1/ns), x)
    # should raise exception for non-integer ns
    @test_throws ErrorException synthesis(excite, mgc, hopsize)
end

function test_synthesis()
    N = 512
    excite = rand(N)
    x = rand(N, div(N, hopsize))
    mgc = estimate(MelGeneralizedCepstrum(order, α, -1/ns), x)
end

### Synthesis with AllPoleDF ###

test_poledf_exception()

for order in 20:5:30
    println("poledf_synthesis_one_frame: testing with order=$order")
    f = AllPoleDF(order)
    test_synthesis_one_frame(f, order)
end

for order in 20:5:30
    for hopsize in [80, 160]
        println("poledf_synthesis: testing with order=$order, hopsize=$hopsize")
        test_poledf_synthesis(order, hopsize)
    end
end

### Synthesis with AllPoleLatticeDF ###

for order in 20:5:30
    println("ltcdf_synthesis_one_frame: testing with order=$order")
    f = AllPoleLatticeDF(order)
    test_synthesis_one_frame(f, order)
end

for order in 20:5:30
    for hopsize in [80, 160]
        println("ltcdf_synthesis: testing with order=$order, hopsize=$hopsize")
        test_ltcdf_synthesis(order, hopsize)
    end
end

### Synthesis with LSPDF ###

for order in 20:5:30
    println("lspdf_synthesis_one_frame: testing with order=$order")
    f = LSPDF(order)
    test_synthesis_one_frame(f, order)
end

for order in 20:5:30
    for hopsize in [80, 160]
        println("lspdf_synthesis: testing with order=$order, hopsize=$hopsize")
        test_lspdf_synthesis(order, hopsize)
    end
end

### Synthesis with LMADF ###

let
    f = LMADF(20)
    @test SynthesisFilters.allpass_alpha(f) == 0.0
    @test SynthesisFilters.glog_gamma(f) == 0.0
end

test_lmadf_exception()

for order in 20:5:30
    for pade in [4, 5]
        println("lamdf_synthesis_one_frame: testing with order=$order, pade=$pade")
        f = LMADF(order, pade=pade)
        test_synthesis_one_frame(f, order)
    end
end

for order in 20:5:30
    for pade in [4, 5]
        for hopsize in [80, 160]
            println("lmadf_synthesis: testing with order=$order, pade=$pade, hopsize=$hopsize")
            test_lmadf_synthesis(order, pade, hopsize)
        end
    end
end

### Synthesis with MLSADF ###

let
    f = MLSADF(20, 0.41)
    @test SynthesisFilters.allpass_alpha(f) == 0.41
    @test SynthesisFilters.glog_gamma(f) == 0.0
end

test_mlsadf_exception()

for order in 20:5:30
    for α in [0.35, 0.544]
        for pade in [4, 5]
            println("mlsadf_synthesis_one_frame: testing with order=$order, α=$α, pade=$pade")
            f = MLSADF(order, α, pade=pade)
            test_synthesis_one_frame(f, order)
        end
    end
end

for order in 20:5:30
    for α in [0.35, 0.544]
        for pade in [4, 5]
            for hopsize in [80, 160]
                println("mlsadf_synthesis: testing with order=$order, α=$α, pade=$pade, hopsize=$hopsize")
                test_mlsadf_synthesis(order, α, pade, hopsize)
            end
        end
    end
end

### Synthesis with MGLSADF ###

let
    nstage = 10
    γ = -1/nstage
    f = MGLSADF(20, 0.41, nstage)
    @test SynthesisFilters.allpass_alpha(f) == 0.41
    @test SynthesisFilters.glog_gamma(f) == γ
end

for order in 20:5:30
    for α in [0.35, 0.544]
        for ns in 2:10
            println("mglsadf_synthesis_one_frame: testing with order=$order, α=$α, nstage=$ns, γ=$(-1.0/ns)")
            f = MGLSADF(order, α, ns)
            test_synthesis_one_frame(f, order)
        end
    end
end

for order in 20:5:30
    for α in [0.35, 0.544]
        for ns in 2:10
            for hopsize in [80, 160]
                println("mglsadf_synthesis: testing with order=$order, α=$α, nstage=$ns, γ=$(-1.0/ns), hopsize=$hopsize")
                test_mglsadf_synthesis(order, α, ns, hopsize)
            end
        end
    end
end

test_mglsadf_exceptions()

### Test Syntehsis ###

for def in [
            LinearPredictionCoef(20),
            LineSpectralPair(20),
            PartialAutoCorrelation(20),
            LinearCepstrum(20),
            MelCepstrum(20, 0.41),
            GeneralizedCepstrum(20, -1/10),
            MelGeneralizedCepstrum(20, 0.41, -1/10)
            ]
    println("test_synthesis: testing with $(def)")
    test_synthesis(def, 80)
end


text:
using DrillMudsThermalProps
using Test

@testset "Entire DrillMudsThermalProps.jl" begin
    # Write your tests here.
    include("WaterTests.jl")
    include("Solids.jl")
    include("MixtureDensity.jl")

    @test isa(try
                rho(101325.0,298.15,TestFluid())
            catch e
                e
            end,Exception)

end


text:
function solve(problem::AllenCahnProblem1D{NeumannBC, NeumannBC, SparseArrays.SparseMatrixCSC{Float64, Int64}}, stepping_method::TTSM) where {TTSM <: BackwardEulerMethod}
    u = problem.u
    t = 0.
    t_trajectory = [t]; u_trajectory = [u]
    n = 1

    matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) - problem.Δt*problem.ϵ*problem.A
    boundary_correction = zeros(problem.nₓ + 2)
    boundary_correction[1] = -2*problem.ϵ*problem.left_bc.boundary_condition/problem.Δx
    boundary_correction[end] = 2*problem.ϵ*problem.right_bc.boundary_condition/problem.Δx

    while n <= problem.nₜ
        problem.rhs .-= problem.rhs
        problem.rhs .+= problem.f.(u)
        problem.rhs .+= boundary_correction
        b = u + problem.Δt*problem.rhs
        t += problem.Δt; push!(t_trajectory, t)
        u = matrix\b; push!(u_trajectory, u)
        n += 1
    end

    return t_trajectory, u_trajectory
end

function solve(problem::AllenCahnProblem1D{PeriodicBC, PeriodicBC, SparseArrays.SparseMatrixCSC{Float64, Int64}}, stepping_method::TTSM) where {TTSM <: BackwardEulerMethod}
    u = problem.u
    t = 0.
    t_trajectory = [t]; u_trajectory = [u]
    n = 1

    matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) - problem.Δt*problem.ϵ*problem.A

    while n <= problem.nₜ
        problem.rhs .-= problem.rhs
        problem.rhs .+= problem.f.(u)
        b = u + problem.Δt*problem.rhs
        t += problem.Δt; push!(t_trajectory, t)
        u = matrix\b; push!(u_trajectory, u)
        n += 1
    end

    push!(problem.x, problem.a + (problem.nₓ + 1)*problem.Δx)
    for i ∈ 1:length(u_trajectory)
        push!(u_trajectory[i], u_trajectory[i][1])
    end

    return t_trajectory, u_trajectory
end

function solve(problem::AllenCahnProblem1D{NeumannBC, NeumannBC, SparseArrays.SparseMatrixCSC{Float64, Int64}}, stepping_method::TTSM) where {TTSM <: CrankNicolsonMethod}
    u = problem.u
    t = 0.
    t_trajectory = [t]; u_trajectory = [u]
    n = 1

    left_matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) - (problem.Δt*problem.ϵ/2)*problem.A
    right_matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) + (problem.Δt*problem.ϵ/2)*problem.A

    while n <= problem.nₜ
        problem.rhs .-= problem.rhs
        problem.rhs .+= problem.f.(u)
        b = right_matrix*u + problem.Δt*problem.rhs
        t += problem.Δt; push!(t_trajectory, t)
        u = left_matrix\b; push!(u_trajectory, u)
        n += 1
    end

    return t_trajectory, u_trajectory
end

function solve(problem::AllenCahnProblem1D{PeriodicBC, PeriodicBC, SparseArrays.SparseMatrixCSC{Float64, Int64}}, stepping_method::TTSM) where {TTSM <: CrankNicolsonMethod}
    u = problem.u
    t = 0.
    t_trajectory = [t]; u_trajectory = [u]
    n = 1

    left_matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) - (problem.Δt*problem.ϵ/2)*problem.A
    right_matrix = sparse(Matrix(1.0I, size(problem.A)[1], size(problem.A)[2])) + (problem.Δt*problem.ϵ/2)*problem.A

    while n <= problem.nₜ
        problem.rhs .-= problem.rhs
        problem.rhs .+= problem.f.(u)
        b = right_matrix*u + problem.Δt*problem.rhs
        t += problem.Δt; push!(t_trajectory, t)
        u = left_matrix\b; push!(u_trajectory, u)
        n += 1
    end

    push!(problem.x, problem.a + (problem.nₓ + 1)*problem.Δx)
    for i ∈ 1:length(u_trajectory)
        push!(u_trajectory[i], u_trajectory[i][1])
    end

    return t_trajectory, u_trajectory
end