text:
module LinearRegressionKit

export regress, predict_in_sample, predict_out_of_sample, linRegRes, kfold, ridge, ridgeRegRes

using Base: Tuple, Int64, Float64, Bool
using StatsBase:eltype, isapprox, length, coefnames, push!, append!
using Distributions, HypothesisTests, LinearAlgebra
using Printf, NamedArrays, FreqTables # FreqTables for check_cardinality
using StatsBase, Random, StatsModels, DataFrames
using VegaLite

include("sweep_operator.jl")
include("utilities.jl")
include("vl_utilities.jl")
include("newey_west.jl")
include("kfold.jl")
include("ridge.jl")

"""
    struct linRegRes

    Store results from the regression
"""
struct linRegRes
    extended_inverse::Matrix                # Store the extended inverse matrix
    coefs::Vector                           # Store the coefficients of the fitted model
    white_types::Union{Nothing,Vector}      # Store the type of White's covariance estimator(s) used
    hac_types::Union{Nothing,Vector}        # Store the type of White's covariance estimator(s) used
    stderrors::Union{Nothing,Vector}        # Store the standard errors for the fitted model
    white_stderrors::Union{Nothing,Vector}  # Store the standard errors modified for the White's covariance estimator
    hac_stderrors::Union{Nothing,Vector}    # Store the standard errors modified for the Newey-West covariance estimator
    t_values::Union{Nothing,Vector}         # Store the t values for the fitted model
    white_t_values::Union{Nothing,Vector}   # Store the t values modified for the White's covariance estimator
    hac_t_values::Union{Nothing,Vector}     # Store the t values modified for the Newey-West covariance estimator
    p::Int64                                # Store the number of parameters (including the intercept as a parameter)
    MSE::Union{Nothing,Float64}             # Store the Mean squared error for the fitted model
    intercept::Bool                         # Indicate if the model has an intercept
    R2::Union{Nothing,Float64}              # Store the R-squared value for the fitted model
    ADJR2::Union{Nothing,Float64}           # Store the adjusted R-squared value for the fitted model
    RMSE::Union{Nothing,Float64}            # Store the Root mean square error for the fitted model
    AIC::Union{Nothing,Float64}             # Store the Akaike information criterion for the fitted model
    σ̂²::Union{Nothing,Float64}              # Store the σ̂² for the fitted model
    p_values::Union{Nothing,Vector}         # Store the p values for the fitted model
    white_p_values::Union{Nothing,Vector}   # Store the p values modified for the White's covariance estimator
    hac_p_values::Union{Nothing,Vector}     # Store the p values modified for the Newey-West covariance estimator
    ci_up::Union{Nothing,Vector}            # Store the upper values confidence interval of the coefficients
    ci_low::Union{Nothing,Vector}           # Store the lower values confidence interval of the coefficients
    white_ci_up::Union{Nothing,Vector}      # Store the upper values confidence interval of the coefficients for White covariance estimators
    white_ci_low::Union{Nothing,Vector}     # Store the upper values confidence interval of the coefficients for White covariance estimators
    hac_ci_up::Union{Nothing,Vector}        # Store the upper values confidence interval of the coefficients for Newey-West covariance estimators
    hac_ci_low::Union{Nothing,Vector}       # Store the upper values confidence interval of the coefficients for Newey-West covariance estimators
    observations                            # Store the number of observations used in the model
    t_statistic::Union{Nothing,Float64}     # Store the t statistic
    VIF::Union{Nothing,Vector}              # Store the Variance inflation factor
    Type1SS::Union{Nothing,Vector}          # Store the Type 1 Sum of Squares
    Type2SS::Union{Nothing,Vector}          # Store the Type 2 Sum of Squares
    pcorr1::Union{Nothing,Vector{Union{Missing, Float64}}}           # Store the squared partial correlation coefficients using Type1SS
    pcorr2::Union{Nothing,Vector{Union{Missing, Float64}}}           # Store the squared partial correlation coefficients using Type2SS
    scorr1::Union{Nothing,Vector{Union{Missing, Float64}}}           # Store the squared semi-partial correlation coefficient using Type1SS
    scorr2::Union{Nothing,Vector{Union{Missing, Float64}}}           # Store the squared semi-partial correlation coefficient using Type2SS
    modelformula                            # Store the model formula
    dataschema                              # Store the dataschema
    updformula                              # Store the updated model formula (after the dataschema has been applied)
    alpha                                   # Store the alpha used to compute the confidence interval of the coefficients
    KS_test::Union{Nothing,String}          # Store results of the Kolmogorov-Smirnov test
    AD_test::Union{Nothing,String}          # Store results of the Anderson–Darling test
    JB_test::Union{Nothing,String}          # Store results of the Jarque-Bera test
    White_test::Union{Nothing,String}       # Store results of the White test
    BP_test::Union{Nothing,String}          # Store results of the Breusch-Pagan test
    weighted::Bool                          # Indicates if this is a weighted regression
    weights::Union{Nothing,String}          # Indicates which column of the dataframe contains the analytical weights
    PRESS::Union{Nothing,Float64}           # Store the PRESS statistic
    cond::Union{Nothing,Float64}            # Store Condition number of the design matrix
end

"""
    function Base.show(io::IO, lr::linRegRes)

    Display information about the fitted model
"""
function Base.show(io::IO, lr::linRegRes)
    if lr.weighted
        println(io, "Weighted regression")
    end
    println(io, "Model definition:\t", lr.modelformula)
    println(io, "Used observations:\t", lr.observations)
    if !isnothing(lr.cond)
        println(io, "Condition number:\t", lr.cond)
    end

    println(io, "Model statistics:")
    # Display stats when available
    if !isnothing(lr.R2) && !isnothing(lr.ADJR2)
        @printf(io, "  R²: %g\t\t\tAdjusted R²: %g\n", lr.R2, lr.ADJR2)
    elseif !isnothing(lr.R2)
        @printf(io, "  R²: %g\n", lr.R2)
    end

    if !isnothing(lr.MSE) && !isnothing(lr.RMSE)
        @printf(io, "  MSE: %g\t\t\tRMSE: %g\n", lr.MSE, lr.RMSE)
    elseif !isnothing(lr.MSE)
        @printf(io, "  MSE: %g\n", lr.MSE)
    end
    if !isnothing(lr.PRESS)
        @printf(io, "  PRESS: %g\n", lr.PRESS)
    end

    if length(lr.white_types) + length(lr.hac_types) == 0
        if !isnothing(lr.σ̂²) && !isnothing(lr.AIC)
            @printf(io, "  σ̂²: %g\t\t\tAIC: %g\n", lr.σ̂², lr.AIC)
        elseif !isnothing(lr.σ̂²)
            @printf(io, "  σ̂²: %g\n", lr.σ̂²)
        elseif !isnothing(lr.AIC)
            @printf(io, "  AIC: %g\n", lr.AIC)
        end
    end

    if !isnothing(lr.ci_low) || !isnothing(lr.ci_up)
        @printf(io, "Confidence interval: %g%%\n", (1 - lr.alpha) * 100 )
    end

    vec_stats_title = ["Coefs", "Std err", "t", "Pr(>|t|)", "low ci", "high ci", "VIF",
            "Type1 SS", "Type2 SS", "PCorr1", "PCorr2",
            "SCorr1", "SCorr2"]

    if length(lr.white_types) + length(lr.hac_types) == 0
        helper_print_table(io, "Coefficients statistics:",
            [lr.coefs, lr.stderrors, lr.t_values, lr.p_values, lr.ci_low, lr.ci_up, lr.VIF,
                lr.Type1SS, lr.Type2SS, lr.pcorr1, lr.pcorr2, lr.scorr1, lr.scorr2],
            vec_stats_title,
            lr.updformula)
    end

    if length(lr.white_types) > 0
        for (cur_i, cur_type) in enumerate(lr.white_types)
            helper_print_table(io, "White's covariance estimator ($(Base.Unicode.uppercase(string(cur_type)))):",
                [lr.coefs, lr.white_stderrors[cur_i], lr.white_t_values[cur_i], lr.white_p_values[cur_i],
                    lr.white_ci_low[cur_i], lr.white_ci_up[cur_i], lr.VIF, lr.Type1SS, lr.Type2SS,
                    lr.pcorr1, lr.pcorr2, lr.scorr1, lr.scorr2],
                vec_stats_title,
                lr.updformula)
        end
    end

    if length(lr.hac_types) > 0
        for (cur_i, cur_type) in enumerate(lr.hac_types)
            helper_print_table(io, "Newey-West's covariance estimator:",
                [lr.coefs, lr.hac_stderrors[cur_i], lr.hac_t_values[cur_i], lr.hac_p_values[cur_i],
                    lr.hac_ci_low[cur_i], lr.hac_ci_up[cur_i], lr.VIF, lr.Type1SS, lr.Type2SS,
                    lr.pcorr1, lr.pcorr2, lr.scorr1, lr.scorr2],
                vec_stats_title,
                lr.updformula)
        end
    end

    if !isnothing(lr.KS_test) || !isnothing(lr.AD_test) || !isnothing(lr.JB_test) || !isnothing(lr.White_test) || !isnothing(lr.BP_test)
        println(io, "\nDiagnostic Tests:\n")
        !isnothing(lr.KS_test) && print(io, lr.KS_test)
        !isnothing(lr.AD_test) && print(io, lr.AD_test)
        !isnothing(lr.JB_test) && print(io, lr.JB_test)
        !isnothing(lr.White_test) && print(io, lr.White_test)
        !isnothing(lr.BP_test) && print(io, lr.BP_test)
    end
end

"""
    function getVIF(x, intercept, p)

    (internal) Calculates the VIF, Variance Inflation Factor, for a given regression (the X).
    When the regression has an intercept uses the simplified formula. When there is no intercept uses the classical formula.
"""
function getVIF(x, intercept, p)
    if intercept
        if p == 1
            return [0., 1.]
        end
        return vcat(0, diag(inv(cor(@view(x[:, 2:end])))))
    else
        if p == 1
            return [0.]
        end
        return diag(inv(cor(x)))
    end
end

"""
    function getSST(y, intercept)

    (internal) Calculates "total sum of squares" see link for description.
    https://en.wikipedia.org/wiki/Total_sum_of_squares
    When the mode has no intercept the SST becomes the sum of squares of y
"""
function getSST(y, intercept)
    SST = zero(eltype(y))
    if intercept
        ȳ = mean(y)
        SST = sum(abs2.(y .- ȳ))
    else
        SST = sum(abs2.(y))
    end
    return SST
end

"""
    function getSST(y, intercept, weights, ridge=false)

    (internal) Calculates "total sum of squares" for weighted regression see link for description.
    https://en.wikipedia.org/wiki/Total_sum_of_squares
    When the mode has no intercept the SST becomes the sum of squares of y.
    When called from ridge regression the ys are not weighted, when called from regression the ys are already weighted.
"""
function getSST(y, intercept, weights, ridge=false)
    SST = zero(eltype(y))
    unweightedys = nothing
    if ridge
        unweightedys = y
    else
        unweightedys = y ./ sqrt.(weights)
    end
    if intercept
        ȳ = mean(unweightedys, aweights(weights))
        SST = sum(weights .* abs2.(unweightedys .- ȳ))
    else
        SST = sum(weights .* abs2.(unweightedys))
    end
    return SST
end

"""
    function lr_predict(xs, coefs, intercept::Bool)

    (internal) Predict the ŷ given the x(s) and the coefficients of the linear regression.
"""
function lr_predict(xs, coefs, intercept::Bool)
    if intercept
        return muladd(@view(xs[:, 2:end]), @view(coefs[2:end]), coefs[1])
    else
        return muladd(xs, coefs, zero(eltype(coefs)))
    end
end

"""
    function hasintercept!(f::StatsModels.FormulaTerm)
    (internal) return a tuple with the first item being true when the formula has an intercept term, the second item being the potentially updated formula.
    If there is no intercept indicated add one.
    If the intercept is specified as absent (y ~ 0 + x) then do not change.
"""
function hasintercept!(f::StatsModels.FormulaTerm)
    intercept = true
    if f.rhs isa ConstantTerm{Int64}
        intercept = convert(Bool, f.rhs.n)
        return intercept, f
    elseif f.rhs isa Tuple
        for t in f.rhs
            if t isa ConstantTerm{Int64}
                intercept = convert(Bool, t.n)
                return intercept, f
            end
        end
    end
    f = FormulaTerm(f.lhs, InterceptTerm{true}() + f.rhs)
    return intercept, f
end

"""
    function get_pcorr(typess, sse, intercept)

    (internal) Get squared partial correlation coefficient given a TYPE1SS or Type2SS.
"""
function get_pcorr(typess, sse, intercept)
    pcorr = Vector{Union{Missing, Float64}}(undef, length(typess))
    if intercept
        @inbounds pcorr[1] = missing
        @inbounds for i in 2:length(typess)
            pcorr[i] = typess[i] / (typess[i] + sse)
        end
    else
        @inbounds for i in 1:length(typess)
            pcorr[i] = typess[i] / (typess[i] + sse)
        end
    end

    return pcorr
end

"""
    function get_scorr(typess, sst, intercept)

    (internal) Get squared semi-partial correlation coefficient given a TYPE1SS or Type2SS.
"""
function get_scorr(typess, sst, intercept)
    scorr = Vector{Union{Missing, Float64}}(undef, length(typess))
    if intercept
        @inbounds scorr[1] = missing
        @inbounds for i in 2:length(typess)
            scorr[i] = typess[i] / sst[i]
        end
    else
        @inbounds for i in 1:length(typess)
            scorr[i] = typess[i] / sst[i]
        end
    end

    return scorr
end

"""
function design_matrix!(f::StatsModels.FormulaTerm, df::DataFrames.AbstractDataFrame;
    weights::Union{Nothing,String}=nothing, remove_missing=false, contrasts=nothing)

    (Internal) Give a design matrix (with x and y separated) given a formula and dataframe.
    Uses weights, and contrasts when given.
    Updates the formula and to removes ambiguity about the intercept term.
    Potentially provide a new dataframe
"""
function design_matrix!(f::StatsModels.FormulaTerm, df::DataFrames.AbstractDataFrame;
    weights::Union{Nothing,String}=nothing, remove_missing=false, contrasts=nothing, ridge=false)

    intercept, f = hasintercept!(f)

    copieddf = df
    if remove_missing
        copieddf = copy(df[: , Symbol.(keys(schema(f, df).schema))])
        dropmissing!(copieddf)
    end

    if isa(weights, String)
        if !in(Symbol(weights), propertynames(copieddf))
            println("Weights have been specified being the column $(weights) however such colum does not exist in the dataframe provided. Regression will be done without weights")
            weights = nothing
        else
            if remove_missing
                copieddf[!, weights] = df[!, weights]
            end
            allowmissing!(copieddf, weights)
            copieddf[!, weights][copieddf[!, weights] .<= 0] .= missing
            dropmissing!(copieddf)
        end
    end
    isweighted = !isnothing(weights)

    if isnothing(contrasts)
        dataschema = schema(f, copieddf)
    else
        dataschema = schema(f, copieddf, contrasts)
    end
    updatedformula = apply_schema(f, dataschema)

    y, x = modelcols(updatedformula, copieddf)
    n, p = size(x)
    if isweighted && ridge == false
        x = x .* sqrt.(copieddf[!, weights])
        y = y .* sqrt.(copieddf[!, weights])
    end

    return x, y, n, p, intercept, f, copieddf, updatedformula, isweighted, dataschema

end

"""
    function regress(f::StatsModels.FormulaTerm, df::AbstractDataFrame, req_plots; α::Float64=0.05, req_stats=["default"], weights::Union{Nothing,String}=nothing, remove_missing=false, cov=[:none], contrasts=nothing, plot_args=Dict("plot_width" => 400, "loess_bw" => 0.6, "residuals_with_density" => false))

    Estimate the coefficients of the regression, given a dataset and a formula. and provide the requested plot(s).
    A dictionary of the generated plots indexed by the descritption of the plots.

    It is possible to indicate the width of the plots, and the bandwidth of the Loess smoother.
"""
function regress(f::StatsModels.FormulaTerm, df::AbstractDataFrame, req_plots; α::Float64=0.05, req_stats=["default"], weights::Union{Nothing,String}=nothing, remove_missing=false, cov=[:none], contrasts=nothing, plot_args=Dict("plot_width" => 400, "loess_bw" => 0.6, "residuals_with_density" => false))

    all_plots = Dict{String,VegaLite.VLSpec}()
    neededplots = get_needed_plots(req_plots)
    lm = regress(f, df, α=α, req_stats=req_stats, remove_missing=remove_missing, cov=cov,
                contrasts=contrasts, weights=weights)
    results = predict_in_sample(lm, df, req_stats="all")

    if :fit in neededplots
        fitplot!(all_plots, results, lm, plot_args)
    end
    if :residuals in neededplots
        residuals_plots!(all_plots, results, lm, plot_args)
    end
    if :normal_checks in neededplots
        normality_plots!(all_plots, results, lm, plot_args)
    end
    if :homoscedasticity in neededplots
        scalelocation_plot!(all_plots, results, lm, plot_args)
    end
    if :cooksd in neededplots
        cooksd_plot!(all_plots, results, lm, plot_args)
    end
    if :leverage in neededplots
        leverage_plot!(all_plots, results, lm, plot_args)
    end
    return (lm, all_plots)
end

"""
    function regress(f::StatsModels.FormulaTerm, df::DataFrames.AbstractDataFrame; α::Float64=0.05, req_stats=["default"], weights::Union{Nothing,String}=nothing,
    remove_missing=false, cov=[:none], contrasts=nothing)

    Estimate the coefficients of the regression, given a dataset and a formula.

    The formula details are provided in the StatsModels package and the behaviour aims to be similar as what the Julia GLM package provides.
    The data shall be provided as a DataFrame without missing data.
    If remove_missing is set to true a copy of the dataframe will be made and the row with missing data will be removed.
    Some robust covariance estimator(s) can be requested through the `cov` argument.
    Default contrast is dummy coding, other contrasts can be requested through the `contrasts` argument.
    For a weighted regression, the name of column containing the analytical weights shall be identified by the `weights` argument.
"""
function regress(f::StatsModels.FormulaTerm, df::DataFrames.AbstractDataFrame; α::Float64=0.05, req_stats=["default"], weights::Union{Nothing,String}=nothing,
                remove_missing=false, cov=[:none], contrasts=nothing)
    (α > 0. && α < 1.) || throw(ArgumentError("α must be between 0 and 1"))

    needed_stats = get_needed_model_stats(req_stats)
    # stats initialization
        total_scalar_stats = Set([:sse, :mse, :sst, :r2, :adjr2, :rmse, :aic, :sigma, :t_statistic, :press, :cond ])
        total_vector_stats = Set([:coefs, :stderror, :t_values, :p_values, :ci, :vif, :t1ss, :t2ss, :pcorr1, :pcorr2, :scorr1, :scorr2])
        total_diag_stats = Set([:diag_ks, :diag_ad, :diag_jb, :diag_white, :diag_bp])

        scalar_stats = Dict{Symbol,Union{Nothing,Float64}}(intersect(total_scalar_stats, needed_stats) .=> nothing)
        vector_stats = Dict{Symbol,Union{Nothing,Vector}}(intersect(total_vector_stats, needed_stats) .=> nothing)
        diag_stats = Dict{Symbol,Union{Nothing,String}}(intersect(total_diag_stats, needed_stats) .=> nothing)
    sse = nothing

    x, y, n, p, intercept, f, copieddf, updatedformula, isweighted, dataschema = design_matrix!(f, df, weights=weights, remove_missing=remove_missing, contrasts=contrasts)
    xy = [x y]

    if :cond in needed_stats
        scalar_stats[:cond] = cond(xy)
    end

    xytxy = xy' * xy

    try
        if :t1ss in needed_stats
            sse, vector_stats[:t1ss] = sweep_op_fullT1SS!(xytxy)
        else
            sse = sweep_op_full!(xytxy)
        end
        catch ae
            throw(ae)
        finally
            check_cardinality(copieddf, updatedformula)
    end
    coefs = xytxy[1:p, end]
    mse = xytxy[p + 1, p + 1] / (n - p)

    if :sst in needed_stats
        if isweighted
            scalar_stats[:sst] = getSST(y, intercept, copieddf[!, weights])
        else
            scalar_stats[:sst] = getSST(y, intercept)
        end
    end
    if :r2 in needed_stats
        scalar_stats[:r2] = 1. - (sse / scalar_stats[:sst])
    end
    if :adjr2 in needed_stats
        scalar_stats[:adjr2] = 1. - ((n - convert(Int64, intercept)) * (1. - scalar_stats[:r2])) / (n - p)
    end
    if :rmse in needed_stats
        scalar_stats[:rmse] = real_sqrt(mse)
    end
    if :aic in needed_stats
        scalar_stats[:aic] =  n * log(sse / n) + 2p
    end
    if :sigma in needed_stats
        scalar_stats[:sigma] = mse
    end
    if :t_statistic in needed_stats
        scalar_stats[:t_statistic] = quantile(TDist(n - p), 1 - α / 2)
    end
    if :t2ss in needed_stats
        vector_stats[:t2ss] = get_TypeIISS(xytxy)
    end
    if :pcorr1 in needed_stats
        vector_stats[:pcorr1] = get_pcorr(vector_stats[:t1ss], sse, intercept)
    end
    if :pcorr2 in needed_stats
        vector_stats[:pcorr2] = get_pcorr(vector_stats[:t2ss], sse, intercept)
    end
    if :scorr1 in needed_stats
        vector_stats[:scorr1] = get_scorr(vector_stats[:t1ss], scalar_stats[:sst], intercept)
    end
    if :scorr2 in needed_stats
        vector_stats[:scorr2] = get_scorr(vector_stats[:t2ss], scalar_stats[:sst], intercept)
    end
    if :stderror in needed_stats
        vector_stats[:stderror] = real_sqrt.(diag(mse * @view(xytxy[1:end - 1, 1:end - 1])))
    end
    if :t_values in needed_stats
        vector_stats[:t_values] = coefs ./ vector_stats[:stderror]
    end
    if :p_values in needed_stats
        vector_stats[:p_values] = ccdf.(Ref(FDist(1., (n - p))), abs2.(vector_stats[:t_values]))
    end
    if :ci in needed_stats
        vector_stats[:ci] = vector_stats[:stderror] * scalar_stats[:t_statistic]
    end
    if :vif in needed_stats
        vector_stats[:vif] = getVIF(x, intercept, p)
    end

    if length(intersect(needed_stats, Set([:diag_ks, :diag_ad, :diag_jb, :diag_white, :diag_bp]))) > 0
        residuals = y - lr_predict(x, coefs, intercept)
        if :diag_ks in needed_stats
            diag_stats[:diag_ks] = present_kolmogorov_smirnov_test(residuals, α)
        end
        if :diag_ad in needed_stats
            diag_stats[:diag_ad] = present_anderson_darling_test(residuals, α)
        end
        if :diag_jb in needed_stats
            diag_stats[:diag_jb] = present_jarque_bera_test(residuals, α)
        end
        if :diag_white in needed_stats
            if intercept && !isweighted
                diag_stats[:diag_white] = present_white_test(x, residuals, α)
            else
                println("White test diagnostic for heteroscedasticity was requested but it requires a non-weighted model with intercept")
            end
        end
        if :diag_bp in needed_stats
            if intercept && !isweighted
                diag_stats[:diag_bp] = present_breusch_pagan_test(x, residuals, α)
            else
                println("Breusch-Pagan test diagnostic for heteroscedasticity was requested but it requires a non weighted model with intercept")
            end
        end
    end

needed_white, needed_hac = get_needed_robust_cov_stats(cov)
# robust estimators stats
white_types = Vector{Symbol}()
white_stds = Vector{Vector}()
white_t_vals = Vector{Vector}()
white_p_vals = Vector{Vector}()
white_ci_up = Vector{Vector}()
white_ci_low = Vector{Vector}()
hac_types = Vector{Symbol}()
hac_stds = Vector{Vector}()
hac_t_vals = Vector{Vector}()
hac_p_vals = Vector{Vector}()
hac_ci_up = Vector{Vector}()
hac_ci_low = Vector{Vector}()

# statistics requiring predictions (robust estimator and PRESS)
if length(needed_white) > 0 || length(needed_hac) > 0 || :press in needed_stats
    predict_results = predict_internal(copieddf, f, updatedformula, isweighted, weights, xytxy, coefs, intercept,
        length(needed_white) > 0, length(needed_hac) > 0, mse, scalar_stats[:t_statistic], p, n;
        α= α, req_stats=[:residuals, :press], dropmissingvalues = false)
    residuals = predict_results.residuals
    presses = predict_results.press
    scalar_stats[:press] = sum(presses.^2)

    if length(needed_white) > 0
        for t in needed_white
            if t in white_types
                continue
            end
            cur_type, cur_std = heteroscedasticity(t, x, y, residuals, n, p, xytxy)
            push!(white_types, cur_type)
            push!(white_stds, cur_std)

            if !isnothing(get(vector_stats, :t_values, nothing))
                cur_t_vals = coefs ./ cur_std
                push!(white_t_vals, cur_t_vals)
            else
                white_t_vals = nothing
            end

            if !isnothing(get(vector_stats, :p_values, nothing))
                cur_p_vals = ccdf.(Ref(FDist(1., (n - p))), abs2.(cur_t_vals))
                push!(white_p_vals, cur_p_vals)
            else
                white_p_vals = nothing
            end
            if !isnothing(get(vector_stats, :ci, nothing))
                cur_ci = cur_std * scalar_stats[:t_statistic]
                cur_ci_up = coefs .+ cur_ci
                cur_ci_low = coefs .- cur_ci
                push!(white_ci_up, cur_ci_up)
                push!(white_ci_low, cur_ci_low)
            else
                white_ci_up = nothing
                white_ci_low = nothing
            end
        end
    end

    if length(needed_hac) > 0
        for t in needed_hac
            if t in hac_types
                continue
            end
            cur_type, cur_std = HAC(t, x, y, residuals, n, p)
            push!(hac_types, cur_type)
            push!(hac_stds, cur_std)

            if !isnothing(get(vector_stats, :t_values, nothing))
                cur_t_vals = coefs ./ cur_std
                push!(hac_t_vals, cur_t_vals)
            else
                hac_t_vals = nothing
            end
            if !isnothing(get(vector_stats, :p_values, nothing))
                cur_p_vals = ccdf.(Ref(FDist(1., (n - p))), abs2.(cur_t_vals))
                push!(hac_p_vals, cur_p_vals)
            else
                hac_p_vals = nothing
            end
            if !isnothing(get(vector_stats, :ci, nothing))
                cur_ci = cur_std * scalar_stats[:t_statistic]
                cur_ci_up = coefs .+ cur_ci
                cur_ci_low = coefs .- cur_ci
                push!(hac_ci_up, cur_ci_up)
                push!(hac_ci_low, cur_ci_low)
            else
                hac_ci_up = nothing
                hac_ci_low = nothing
            end
        end
    end

end

    sres = linRegRes(xytxy, coefs,
        white_types, hac_types,
        get(vector_stats, :stderror, nothing), white_stds, hac_stds,
        get(vector_stats, :t_values, nothing), white_t_vals, hac_t_vals,
        p, mse, intercept, get(scalar_stats, :r2, nothing),
        get(scalar_stats, :adjr2, nothing), get(scalar_stats, :rmse, nothing), get(scalar_stats, :aic, nothing), get(scalar_stats, :sigma, nothing),
        get(vector_stats, :p_values, nothing), white_p_vals, hac_p_vals,
        haskey(vector_stats, :ci) ? coefs .+ vector_stats[:ci] : nothing,
        haskey(vector_stats, :ci) ? coefs .- vector_stats[:ci] : nothing,
        white_ci_up, white_ci_low,
        hac_ci_up, hac_ci_low,
        n, get(scalar_stats, :t_statistic, nothing), get(vector_stats, :vif, nothing),
        get(vector_stats, :t1ss, nothing), get(vector_stats, :t2ss, nothing),
        get(vector_stats, :pcorr1, nothing), get(vector_stats, :pcorr2, nothing),
        get(vector_stats, :scorr1, nothing), get(vector_stats, :scorr2, nothing),
        f, dataschema, updatedformula, α,
        get(diag_stats, :diag_ks, nothing), get(diag_stats, :diag_ad, nothing), get(diag_stats, :diag_jb, nothing),
        get(diag_stats, :diag_white, nothing),  get(diag_stats, :diag_bp, nothing),
        isweighted, weights, get(scalar_stats, :press, nothing),
        get(scalar_stats, :cond, nothing)
        )

    return sres
end

"""
    function HAC(t::Symbol, x, y, residuals, n, p)

    (Internal) Return the relevant HAC (heteroskedasticity and autocorrelation consistent) estimator.
    In the current version only Newey-West is implemented.
"""
function HAC(t::Symbol, x, y, residuals, n, p)
    inv_xtx = inv(x' * x)
    xe = x .* residuals
    return (t, sqrt.(diag(n * inv_xtx * newey_west(xe) * inv_xtx)))
end

"""
    function heteroscedasticity(t::Symbol, x, y, residuals, n, p, xytxy)

    (Internal) Compute the standard errors modified for the White's covariance estimator.
    Currently support HC0, HC1, HC2 and HC3. When :white is passed, select HC3 when the number of observation is below 250 otherwise select HC0.
"""
function heteroscedasticity(t::Symbol, x, y, residuals, n, p, xytxy)
        inv_xtx = inv(x' * x)
    XX = @view(xytxy[1:end - 1, 1:end - 1])
    xe = x .* residuals

    if t == :white && n <= 250
        t = :hc3
    elseif t == :white && n > 250
        t = :hc0
    end

    if t == :hc0
        xetxe = xe' * xe
        return (:hc0, real_sqrt.(diag(XX * xetxe * XX)))
    elseif t == :hc1
        scale = (n / (n - p))
        xetxe = xe' * xe
        return (:hc1, real_sqrt.(diag(XX * xetxe * XX .* scale)))
    elseif t == :hc2
        leverage = diag(x * inv(x'x) * x')
        scale = @.( 1. / (1. - leverage))
        xe = @.(xe .* real_sqrt(scale))
        xetxe = xe' * xe
        return (t, sqrt.(diag(XX * xetxe * XX)))
    elseif t == :hc3
        leverage = diag(x * inv(x'x) * x')
        scale = @.( 1. / (1. - leverage)^2)
        xe = @.(xe .* real_sqrt(scale))
        xetxe = xe' * xe
        return (t, sqrt.(diag(XX * xetxe * XX)))
    else
        throw(error("Unknown symbol ($(t)) used as the White's covariance estimator"))
    end

end

"""
    function predict_internal(df::AbstractDataFrame, modelformula, updatedformula, weighted, weights, extended_inverse,
    coefs, intercept, needed_white, needed_hac, σ̂², t_statistic, p, n, oos=false;
    α=0.05, req_stats=["none"], dropmissingvalues=true)

    Internal, users should use `predict_in_sample` or `predict_out_of_sample`. This should be used only when the `struct linRegRes` is not constructed yet.
"""
function predict_internal(df::AbstractDataFrame, modelformula, updatedformula, weighted, weights, extended_inverse,
            coefs, intercept, needed_white, needed_hac, σ̂², t_statistic, p, n, oos=false;
            α=0.05, req_stats=["none"], dropmissingvalues=true)

    copieddf = df
    if oos
        copieddf = df[: , Symbol.(keys(schema(modelformula.rhs, df).schema))]
    else
        copieddf = df[: , Symbol.(keys(schema(modelformula, df).schema))]
    end
    if dropmissingvalues == true
        dropmissing!(copieddf)
    end

    if weighted
        if !in(Symbol(weights), propertynames(df))
            println(io, "Weights have been specified being the column $(weights) however such colum does not exist in the dataframe provided. Regression will be done without weights")
            weights = nothing
        else
            copieddf[!, weights] = df[!, weights]
            allowmissing!(copieddf, weights)
            copieddf[!, weights][copieddf[!, weights] .<= 0] .= missing
            dropmissing!(copieddf)
        end
    end
    y = nothing
    x = nothing

    if oos
        x = modelcols(updatedformula.rhs, copieddf)
    else
        y, x = modelcols(updatedformula, copieddf)
    end

    needed, present = get_prediction_stats(req_stats)
    needed_stats = Dict{Symbol,Vector}()
    for sym in needed
        needed_stats[sym] = zeros(length(n))
    end
    if :leverage in needed
        pinverse = @view(extended_inverse[1:end - 1, 1:end - 1])
        if weighted
            needed_stats[:leverage] = copieddf[!, weights] .* diag(x * pinverse * x')
        else
            needed_stats[:leverage] = diag(x * pinverse * x')
        end
    end
    if :predicted in needed
        needed_stats[:predicted] = lr_predict(x, coefs, intercept)
    end
    if :residuals in needed && oos == false
        needed_stats[:residuals] = y .- needed_stats[:predicted]
    end
    if :stdp in needed
        if isnothing(σ̂²)
            throw(ArgumentError(":stdp requires that the σ̂² (:sigma) was previously calculated through the regression"))
        end
        warn_sigma(needed_white, needed_hac, :stdp)
        if weighted
            needed_stats[:stdp] = real_sqrt.(needed_stats[:leverage] .* σ̂² ./ copieddf[!, weights])
        else
            needed_stats[:stdp] = real_sqrt.(needed_stats[:leverage] .* σ̂²)
        end
    end
    if :stdi in needed
        if isnothing(σ̂²)
            throw(ArgumentError(":stdi requires that the σ̂² (:sigma) was previously calculated through the regression"))
        end
        warn_sigma(needed_white, needed_hac, :stdi)
        if weighted
            needed_stats[:stdi] = real_sqrt.((1. .+ needed_stats[:leverage]) .* σ̂² ./ copieddf[!, weights])
        else
            needed_stats[:stdi] = real_sqrt.((1. .+ needed_stats[:leverage]) .* σ̂²)
        end
    end
    if :stdr in needed
        if isnothing(σ̂²)
            throw(ArgumentError(":stdr requires that the σ̂² (:sigma) was previously calculated through the regression"))
        end
        warn_sigma(needed_white, needed_hac, :stdr)
        if weighted
            needed_stats[:stdr] = real_sqrt.((1. .- needed_stats[:leverage]) .* σ̂² ./ copieddf[!, weights] )
        else
            needed_stats[:stdr] = real_sqrt.((1. .- needed_stats[:leverage]) .* σ̂²)
        end
    end
    if :student in needed && oos == false
        warn_sigma(needed_white, needed_hac, :student)
        needed_stats[:student] = needed_stats[:residuals] ./ needed_stats[:stdr]
    end
    if :rstudent in needed && oos == false
        warn_sigma(needed_white, needed_hac, :rstudent)
        needed_stats[:rstudent] = needed_stats[:student] .* real_sqrt.( (n .- p .- 1 ) ./ (n .- p .- needed_stats[:student].^2 ) )
    end
    if :lcli in needed
        warn_sigma(needed_white, needed_hac, :lcli)
        needed_stats[:lcli] = needed_stats[:predicted] .- (t_statistic .* needed_stats[:stdi])
    end
    if :ucli in needed
        warn_sigma(needed_white, needed_hac, :ucli)
        needed_stats[:ucli] = needed_stats[:predicted] .+ (t_statistic .* needed_stats[:stdi])
    end
    if :lclp in needed
        warn_sigma(needed_white, needed_hac, :lclp)
        needed_stats[:lclp] = needed_stats[:predicted] .- (t_statistic .* needed_stats[:stdp])
    end
    if :uclp in needed
        warn_sigma(needed_white, needed_hac, :uclp)
        needed_stats[:uclp] = needed_stats[:predicted] .+ (t_statistic .* needed_stats[:stdp])
    end
    if :press in needed && oos == false
        needed_stats[:press] = needed_stats[:residuals] ./ (1. .- needed_stats[:leverage])
    end
    if :cooksd in needed && oos == false
        warn_sigma(needed_white, needed_hac, :cooksd)
        needed_stats[:cooksd] = needed_stats[:stdp].^2 ./  needed_stats[:stdr].^2 .* needed_stats[:student].^2 .* (1 / p)
    end

    for sym in present
        copieddf[!, sym] = needed_stats[sym]
    end

    return copieddf
end

"""
    function predict_in_sample(lr::linRegRes, df::AbstractDataFrame; α=0.05, req_stats=["none"], dropmissingvalues=true)

    Using the estimated coefficients from the regression make predictions, and calculate related statistics.
"""
function predict_in_sample(lr::linRegRes, df::AbstractDataFrame; α=0.05, req_stats=["none"], dropmissingvalues=true)
    predict_internal(df, lr.modelformula, lr.updformula, lr.weighted, lr.weights, lr.extended_inverse, lr.coefs, lr.intercept,
        lr.white_types, lr.hac_types, lr.σ̂², lr.t_statistic, lr.p, lr.observations;
        α=α, req_stats=req_stats, dropmissingvalues=dropmissingvalues)
end

"""
    function predict_out_of_sample(lr::linRegRes, df::AbstractDataFrame; α=0.05, req_stats=["none"], dropmissingvalues=true)

    Similar to `predict_in_sample` although it does not expect a response variable nor produce statistics requiring a response variable.
"""
function predict_out_of_sample(lr::linRegRes, df::AbstractDataFrame; α=0.05, req_stats=["none"], dropmissingvalues=true)
    predict_internal(df, lr.modelformula, lr.updformula, lr.weighted, lr.weights, lr.extended_inverse, lr.coefs, lr.intercept,
        lr.white_types, lr.hac_types, lr.σ̂², lr.t_statistic, lr.p, lr.observations, true;
        α=α, req_stats=req_stats, dropmissingvalues=dropmissingvalues)
end

end # end of module definition


text:
KitBase.sample_maxwell(0.8, 1.0)
KitBase.sample_maxwell([1.0, 0.0, 1.0])
KitBase.sample_maxwell([1.0, 0.0, 0.0, 1.0])
KitBase.sample_maxwell([1.0, 0.0, 0.0, 0.0, 1.0])
KitBase.sample_maxwell(1.0, -2.0, 2.0)
KitBase.sample_maxwell(1.0, 0.3, -2.0, 2.0)
KitBase.sample_maxwell([1.0, 0.0, 1.0], -1.0, 1.0)
KitBase.sample_maxwell([1.0, 0.0, 0.0, 1.0], -1.0, 1.0)
KitBase.sample_maxwell([1.0, 0.0, 0.0, 0.0, 1.0], -1.0, 1.0)

KitBase.next_collision_time(1.0)

KitBase.boundary_time(1.0, 1.0, 0.5, 1.5)
KitBase.boundary_time(1.0, -1.0, 0.5, 1.5)
KitBase.boundary_time(1.0, 0.0, 0.5, 1.5)
KitBase.boundary_time(1.0, [1.0, 0.0, 0.0], 0.5, 1.5)
KitBase.boundary_time(1.0, [-1.0, 0.0, 0.0], 0.5, 1.5)
KitBase.boundary_time(1.0, [0.0, 0.0, 0.0], 0.5, 1.5)

ptc1 = KitBase.Particle1D(1e-3, 0.0, randn(3), 0.1, 1)
ptc2 = KitBase.Particle2D(1e-3, 0.0, 0.0, randn(3), 0.1, 1, 1)

KitBase.sample_particle!(ptc1, 1e-4, rand(), randn(3), rand(), 2, 0, 0.1)
KitBase.sample_particle!(ptc1, 1e-4, rand(), rand(), [1.0, 0.0, 1.0], 2, 1e-3, 0.72, 0)
KitBase.sample_particle!(
    ptc1,
    1e-4,
    rand(),
    rand(),
    [1.0, 0.0, 1.0],
    -1.0,
    1.0,
    2,
    1e-3,
    0.72,
    0,
)

cd(@__DIR__)
ks, ctr, face, simTime = KitBase.initialize("config.txt")

KitBase.sample_particle!(ptc1, ks, ctr[1], 1, 2.0, 0.05)

KitBase.Particle(
    ones(50) .* 1e-3,
    randn(50),
    randn(50, 3),
    rand(50),
    collect(1:50),
    zeros(Int64, 50),
    zeros(50),
)
KitBase.ControlVolumeParticle1D(
    rand(),
    rand(),
    KitBase.prim_conserve([1.0, 0.0, 1.0], 1.67),
    [1.0, 0.0, 1.0],
)
KitBase.ControlVolumeParticle2D(
    rand(),
    rand(),
    rand(),
    rand(),
    KitBase.prim_conserve([1.0, 0.0, 0.0, 1.0], 1.67),
    [1.0, 0.0, 0.0, 1.0],
)

using OffsetArrays

begin
    cd(@__DIR__)
    D = KitBase.read_dict("poiseuille.txt")
    for key in keys(D)
        s = Symbol(key)
        @eval $s = $(D[key])
    end

    γ = KitBase.heat_capacity_ratio(inK, 1)
    set = KitBase.Setup(
        matter,
        case,
        space,
        flux,
        collision,
        nSpecies,
        interpOrder,
        limiter,
        boundary,
        cfl,
        maxTime,
    )
    pSpace = KitBase.PSpace1D(x0, x1, nx, nxg)
    vSpace = KitBase.VSpace1D(umin, umax, nu, vMeshType, nug)
    μᵣ = KitBase.ref_vhs_vis(knudsen, alphaRef, omegaRef)
    gas = KitBase.Gas(
        knudsen,
        mach,
        prandtl,
        inK,
        γ,
        omega,
        alphaRef,
        omegaRef,
        μᵣ,
        mass,
        0,
        (),
    )

    primL = [1.0, 0.0, -1.0, 1.0] # left wall
    primR = [1.0, 0.0, 1.0, 1.0] # right wall
    wL = KitBase.prim_conserve(primL, γ)
    wR = KitBase.prim_conserve(primR, γ)

    fw = function (x, p)
        if x <= (pSpace.x0 + pSpace.x1) / 2
            wL
        else
            wR
        end
    end
    bc = function (x, p)
        if x <= (pSpace.x0 + pSpace.x1) / 2
            primL
        else
            primR
        end
    end
    ib = KitBase.IB(fw, bc)

    ks = KitBase.SolverSet(set, pSpace, vSpace, gas, ib, pwd())
end

begin
    ctr = OffsetArray{KitBase.ControlVolumeParticle1D}(undef, eachindex(ks.pSpace.x))
    face = Array{KitBase.Interface1D}(undef, ks.pSpace.nx + 1)
    for i in eachindex(ctr)
        prim = [1.0, 0.0, primL[3] + 2.0 * (ks.pSpace.x[i] - ks.pSpace.x0), 1.0]

        ctr[i] = KitBase.ControlVolumeParticle1D(
            ks.pSpace.x[i],
            ks.pSpace.dx[i],
            KitBase.prim_conserve(prim, ks.gas.γ),
            prim,
            KitBase.vhs_collision_time(prim, ks.gas.μᵣ, ks.gas.ω),
        )
    end

    for i = 1:ks.pSpace.nx+1
        face[i] = KitBase.Interface1D(ctr[1].w)
    end

    t = 0.0
    dt = KitBase.timestep(ks, ctr, t)
    res = zeros(4)
end

ptc = KitBase.init_ptc!(ks, ctr)

for iter = 1:1
    KitBase.free_transport!(ks, ptc.x, ptc.v, ptc.flag, dt)
    KitBase.boundary!(ks, ctr, ptc, face, dt, :maxwell)
    KitBase.sort!(ks, ctr, ptc.x, ptc.idx, ptc.ref)
    KitBase.dsmc!(ks, ctr, ptc.ref, ptc.v, dt)
    KitBase.stat!(ks, ctr, ptc)
end

ext: jl


text:
function FieldTuple(arg0::Vector{RealFieldElement})
    return FieldTuple((Vector{RealFieldElement},), arg0)
end

function abs(obj::FieldTuple)
    return jcall(obj, "abs", FieldTuple, ())
end

function acos(obj::FieldTuple)
    return jcall(obj, "acos", FieldTuple, ())
end

function acosh(obj::FieldTuple)
    return jcall(obj, "acosh", FieldTuple, ())
end

function add(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "add", FieldTuple, (FieldTuple,), arg0)
end

function add(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "add", FieldTuple, (jdouble,), arg0)
end

function asin(obj::FieldTuple)
    return jcall(obj, "asin", FieldTuple, ())
end

function asinh(obj::FieldTuple)
    return jcall(obj, "asinh", FieldTuple, ())
end

function atan(obj::FieldTuple)
    return jcall(obj, "atan", FieldTuple, ())
end

function atan2(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "atan2", FieldTuple, (FieldTuple,), arg0)
end

function atanh(obj::FieldTuple)
    return jcall(obj, "atanh", FieldTuple, ())
end

function cbrt(obj::FieldTuple)
    return jcall(obj, "cbrt", FieldTuple, ())
end

function ceil(obj::FieldTuple)
    return jcall(obj, "ceil", FieldTuple, ())
end

function copy_sign(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "copySign", FieldTuple, (FieldTuple,), arg0)
end

function copy_sign(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "copySign", FieldTuple, (jdouble,), arg0)
end

function cos(obj::FieldTuple)
    return jcall(obj, "cos", FieldTuple, ())
end

function cosh(obj::FieldTuple)
    return jcall(obj, "cosh", FieldTuple, ())
end

function divide(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "divide", FieldTuple, (FieldTuple,), arg0)
end

function divide(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "divide", FieldTuple, (jdouble,), arg0)
end

function equals(obj::FieldTuple, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function exp(obj::FieldTuple)
    return jcall(obj, "exp", FieldTuple, ())
end

function expm1(obj::FieldTuple)
    return jcall(obj, "expm1", FieldTuple, ())
end

function floor(obj::FieldTuple)
    return jcall(obj, "floor", FieldTuple, ())
end

function get_component(obj::FieldTuple, arg0::jint)
    return jcall(obj, "getComponent", RealFieldElement, (jint,), arg0)
end

function get_components(obj::FieldTuple)
    return jcall(obj, "getComponents", Vector{RealFieldElement}, ())
end

function get_dimension(obj::FieldTuple)
    return jcall(obj, "getDimension", jint, ())
end

function get_field(obj::FieldTuple)
    return jcall(obj, "getField", Field, ())
end

function get_real(obj::FieldTuple)
    return jcall(obj, "getReal", jdouble, ())
end

function hash_code(obj::FieldTuple)
    return jcall(obj, "hashCode", jint, ())
end

function hypot(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "hypot", FieldTuple, (FieldTuple,), arg0)
end

function is_infinite(obj::CalculusFieldElement)
    return jcall(obj, "isInfinite", jboolean, ())
end

function is_na_n(obj::CalculusFieldElement)
    return jcall(obj, "isNaN", jboolean, ())
end

function linear_combination(obj::FieldTuple, arg0::FieldTuple, arg1::FieldTuple, arg2::FieldTuple, arg3::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (FieldTuple, FieldTuple, FieldTuple, FieldTuple), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::FieldTuple, arg0::FieldTuple, arg1::FieldTuple, arg2::FieldTuple, arg3::FieldTuple, arg4::FieldTuple, arg5::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::FieldTuple, arg0::FieldTuple, arg1::FieldTuple, arg2::FieldTuple, arg3::FieldTuple, arg4::FieldTuple, arg5::FieldTuple, arg6::FieldTuple, arg7::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple, FieldTuple), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function linear_combination(obj::FieldTuple, arg0::Vector{FieldTuple}, arg1::Vector{FieldTuple})
    return jcall(obj, "linearCombination", FieldTuple, (Vector{FieldTuple}, Vector{FieldTuple}), arg0, arg1)
end

function linear_combination(obj::FieldTuple, arg0::Vector{jdouble}, arg1::Vector{FieldTuple})
    return jcall(obj, "linearCombination", FieldTuple, (Vector{jdouble}, Vector{FieldTuple}), arg0, arg1)
end

function linear_combination(obj::FieldTuple, arg0::jdouble, arg1::FieldTuple, arg2::jdouble, arg3::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (jdouble, FieldTuple, jdouble, FieldTuple), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::FieldTuple, arg0::jdouble, arg1::FieldTuple, arg2::jdouble, arg3::FieldTuple, arg4::jdouble, arg5::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (jdouble, FieldTuple, jdouble, FieldTuple, jdouble, FieldTuple), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::FieldTuple, arg0::jdouble, arg1::FieldTuple, arg2::jdouble, arg3::FieldTuple, arg4::jdouble, arg5::FieldTuple, arg6::jdouble, arg7::FieldTuple)
    return jcall(obj, "linearCombination", FieldTuple, (jdouble, FieldTuple, jdouble, FieldTuple, jdouble, FieldTuple, jdouble, FieldTuple), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function log(obj::FieldTuple)
    return jcall(obj, "log", FieldTuple, ())
end

function log10(obj::FieldTuple)
    return jcall(obj, "log10", FieldTuple, ())
end

function log1p(obj::FieldTuple)
    return jcall(obj, "log1p", FieldTuple, ())
end

function multiply(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "multiply", FieldTuple, (FieldTuple,), arg0)
end

function multiply(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "multiply", FieldTuple, (jdouble,), arg0)
end

function multiply(obj::FieldTuple, arg0::jint)
    return jcall(obj, "multiply", FieldTuple, (jint,), arg0)
end

function negate(obj::FieldTuple)
    return jcall(obj, "negate", FieldTuple, ())
end

function new_instance(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "newInstance", FieldTuple, (jdouble,), arg0)
end

function pow(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "pow", FieldTuple, (FieldTuple,), arg0)
end

function pow(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "pow", FieldTuple, (jdouble,), arg0)
end

function pow(obj::FieldTuple, arg0::jint)
    return jcall(obj, "pow", FieldTuple, (jint,), arg0)
end

function reciprocal(obj::FieldTuple)
    return jcall(obj, "reciprocal", FieldTuple, ())
end

function remainder(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "remainder", FieldTuple, (FieldTuple,), arg0)
end

function remainder(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "remainder", FieldTuple, (jdouble,), arg0)
end

function rint(obj::FieldTuple)
    return jcall(obj, "rint", FieldTuple, ())
end

function root_n(obj::FieldTuple, arg0::jint)
    return jcall(obj, "rootN", FieldTuple, (jint,), arg0)
end

function scalb(obj::FieldTuple, arg0::jint)
    return jcall(obj, "scalb", FieldTuple, (jint,), arg0)
end

function signum(obj::FieldTuple)
    return jcall(obj, "signum", FieldTuple, ())
end

function sin(obj::FieldTuple)
    return jcall(obj, "sin", FieldTuple, ())
end

function sin_cos(obj::FieldTuple)
    return jcall(obj, "sinCos", FieldSinCos, ())
end

function sinh(obj::FieldTuple)
    return jcall(obj, "sinh", FieldTuple, ())
end

function sqrt(obj::FieldTuple)
    return jcall(obj, "sqrt", FieldTuple, ())
end

function subtract(obj::FieldTuple, arg0::FieldTuple)
    return jcall(obj, "subtract", FieldTuple, (FieldTuple,), arg0)
end

function subtract(obj::FieldTuple, arg0::jdouble)
    return jcall(obj, "subtract", FieldTuple, (jdouble,), arg0)
end

function tan(obj::FieldTuple)
    return jcall(obj, "tan", FieldTuple, ())
end

function tanh(obj::FieldTuple)
    return jcall(obj, "tanh", FieldTuple, ())
end

function to_degrees(obj::FieldTuple)
    return jcall(obj, "toDegrees", FieldTuple, ())
end

function to_radians(obj::FieldTuple)
    return jcall(obj, "toRadians", FieldTuple, ())
end


text:
function write_output(	file_path::String,
						number_of_samples::Int64,
						d2::Array{Float64, 1},
						S2::Array{Float64, 1},
						d3::Array{Float64, 1},
						theta3::Array{Float64, 1},
						S3::Array{Float64, 3},
						zeta2::Float64,
						t_exec::Float64)

	file_stream::IOStream = open(file_path, "w")

	@printf(file_stream, "%s", "<output_characterization>\n")

	write_key(file_stream, "number_of_samples", number_of_samples)
	write_key(file_stream, "d2", d2)
	write_key(file_stream, "S2", S2)
	write_key(file_stream, "d3", d3)
	write_key(file_stream, "theta3", theta3)
	write_key(file_stream, "S3", S3[:])
	write_key(file_stream, "zeta2", zeta2)
	write_key(file_stream, "execution_time", t_exec)

	@printf(file_stream, "%s", "</output_characterization>")

	close(file_stream)

	nothing
end


text:
#=
ir_dump.jl
2019-03-05 Jeff Fessler
=#

export ir_dump

"""
    ir_dump(x::Any ; io::IO = stdout)
    ir_dump(io::IO, x::Any)

Show all the fields of a structure or `NamedTuple` more nicely than dump() does
"""
function ir_dump(y::Any ; io::IO = stdout, ntuplemax::Int = 3)
	x = typeof(y)
	print(io, x)
	fields = fieldnames(x)
	fieldtypes = x.types
	for (idx,fd) in enumerate(fields)
		println(io)
		print(io, " ", fd, "::")
		ft = fieldtypes[idx]
		print(io, ft)
		if ft <: Number
			print(io, " ", getfield(y, fd))
		end
		if ft == String
			print(io, " \"", getfield(y, fd), "\"")
		end
		if ft <: Tuple # nice for ImageGeom
			tmp = getfield(y, fd)
			if length(tmp) <= ntuplemax
				print(io, " ", tmp)
			else
				print(io, " (")
				for it = 1:ntuplemax
					print(io, tmp[it], ", ")
				end
				print(io, "…)")
			end
		end
		if ft <: AbstractArray{Bool} # nice for ImageGeom
			tmp = getfield(y, fd)
			print(io, " {", sum(tmp), " of ", length(tmp), "}")
		end
	end
	println(io)
	nothing
end

ir_dump(io::IO, x::Any) = ir_dump(x ; io=io)


text:
using AWSS3
using Test
using Dates
using AWSCore
using Retry
using HTTP
using FilePathsBase
using FilePathsBase: /, join
using FilePathsBase.TestPaths
using UUIDs: uuid4

AWSCore.set_debug_level(0)

@testset "AWSS3.jl" begin
    include("s3path.jl")
    include("awss3.jl")
end


text:
"""
Struct containing all the preallocated arrays for the calculation of horizontal diffusion.
"""
struct HorizontalDiffusion{NF<:AbstractFloat}   # Number format NF
    dmp::Array{NF,2}                # Damping coefficient for temperature and vorticity (explicit)
    dmpd::Array{NF,2}               # Damping coefficient for divergence (explicit)
    dmps::Array{NF,2}               # Damping coefficient for extra diffusion in the stratosphere (explicit)

    dmp1::Array{NF,2}               # Damping coefficient for temperature and vorticity (implicit)
    dmp1d::Array{NF,2}              # Damping coefficient for divergence (implicit)
    dmp1s::Array{NF,2}              # Damping coefficient for extra diffusion in the stratosphere (implicit)

    tcorv::Array{NF,1}              # Vertical component of orographic correction for temperature
    qcorv::Array{NF,1}              # Vertical component of orographic correction for humidity

    tcorh::Array{Complex{NF},2}     # Horizontal component of orographic correction for temperature
    qcorh::Array{Complex{NF},2}     # Horizontal component of orographic correction for humidity
end

"""
Generator function for a HorizontalDiffusion struct.
"""
function HorizontalDiffusion(   P::Parameters,      # Parameter struct
                                G::GeoSpectral,     # Geometry and spectral struct
                                B::Boundaries)      # Boundaries struct

    @unpack nlev, σ_full = G.geometry
    @unpack trunc, mx, nx = G.spectral
    @unpack g, R, γ, hscale, hshum, rh_ref = P
    @unpack npowhd, thd, thdd, thds, tdrs = P
    @unpack ϕ0trunc = B

    # Damping frequencies [1/s]
    hdiff = 1.0/(3600.0*thd)            # Spectral damping for temperature and vorticity
    hdifd = 1.0/(3600.0*thdd)           # Spectral damping coefficient for divergence
    hdifs = 1.0/(3600.0*thds)           # Spectral damping coefficient for stratosphere
    rlap  = 1.0/(trunc*(trunc + 1))     # ?

    # preallocate (high precision, conversion to NF later)
    dmp = zeros(mx,nx)                  # Damping coefficient for temperature and vorticity (explicit)
    dmpd = zeros(mx,nx)                 # Damping coefficient for divergence (explicit)
    dmps = zeros(mx,nx)                 # Damping coefficient for extra diffusion in the stratosphere (explicit)

    dmp1 = zeros(mx,nx)                 # Damping coefficient for temperature and vorticity (implicit)
    dmp1d = zeros(mx,nx)                # Damping coefficient for divergence (implicit)
    dmp1s = zeros(mx,nx)                # Damping coefficient for extra diffusion in the stratosphere (implicit)

    for j in 1:nx
        for i in 1:mx
            N = i+j-2
            elap = (N*(N + 1.0)*rlap)
            elapn = elap^npowhd
            dmp[i,j]  = hdiff*elapn
            dmpd[i,j] = hdifd*elapn
            dmps[i,j] = hdifs*elap

            dmp1[i,j]  = 1/(1+dmp*Δt)
            dmp1d[i,j] = 1/(1+dmpd*Δt)
            dmp1s[i,j] = 1/(1+dmps*Δt)
        end
    end

    # Orographic correction terms for temperature and humidity (vertical component)
    γ_g = γ/(1000.0g)
    rgam = R*γ_g
    qexp = hscale/hshum

    # preallocate (high precision, conversion to NF later)
    tcorv = zeros(nlev)    # Vertical component of orographic correction for temperature
    qcorv = zeros(nlev)    # Vertical component of orographic correction for humidity

    for k in 2:nlev
        tcorv[k] = σ_full[k]^rgam
        if k > 2
            qcorv[k] = σ_full[k]^qexp
        end
    end

    # Orographic correction term for temperature (horizontal component)
    corh = zeros(nlon, nlat)            # in grid-point space

    for j in 1:nlat
        for i = 1:nlon
            corh[i,j] = γ_g*ϕ0trunc[i,j]
        end
    end

    tcorh = spectral(corh,G) # correction in spectral space

    # Orographic correction terms for humidity (horizontal component)
    corh .= rh_ref                  # relative humidity reference value
    qcorh = spectral(corh,G)

    # convert to number format NF here
    return HorizontalDiffusion{P.NF}(   dmp,dmpd,dmps,
                                        dmp1,dmpd1,dmps1,
                                        tcorv,qcorv,tcorh,qcorh)
end

"""Apply horizontal diffusion to 2D field in spectral space."""
function horizontal_diffusion!( A::AbstractArray{Complex{NF},2},        # spectral horizontal field
                                tendency::AbstractArray{Complex{NF},2}, # its tendency
                                dmp::AbstractArray{NF,2},               # damping coefficients (explicit)
                                dmp1::AbstractArray{NF,2}               # damping coefficients (implicit)
                                ) where {NF<:AbstractFloat}

    @boundscheck size(A) == size(tendency) || throw(BoundsError())
    @boundscheck size(A) == size(dmp) || throw(BoundsError())
    @boundscheck size(A) == size(dmp1) || throw(BoundsError())

    @inbounds for i in eachindex(A)
        tendency[i] = (tendency[i] - dmp[i]*A[i])*dmp1[i]
    end
end

"""Apply horizontal diffusion to 3D field layer by layer in spectral space."""
function horizontal_diffusion!( A::AbstractArray{Complex{NF},3},        # spectral horizontal field
                                tendency::AbstractArray{Complex{NF},3}, # its tendency
                                dmp::AbstractArray{NF,2},               # damping coefficients (explicit)
                                dmp1::AbstractArray{NF,2}               # damping coefficients (implicit)
                                ) where {NF<:AbstractFloat}             # number format NF
    _,_,nlev = size(A)
    @boundscheck size(A) == size(tendency) || throw(BoundsError())

    for k in 1:nlev
        A_layer = view(A,:,:,k)
        tendency_layer = view(tendency,:,:,k)
        horizontal_diffusion!(A_layer, tendency_layer, dmp, dmp1)
    end
end

"""Zonal drag in the stratosphere (uppermost layer) for vorticity and divergence."""
function stratospheric_zonal_drag!( A::AbstractArray{Complex{NF},4},        # spectral vorticity or divergence
                                    tendency::AbstractArray{Complex{NF},3}, # its tendency
                                    sdrag::NF                               # drag coefficient [1/s]
                                    ) where {NF<:AbstractFloat}             # number format NF

    mx,nx,nlev,nleapfrog = size(A)
    @boundscheck (mx,nx,nlev) == size(tendency) || throw(BoundsError())
    @boundscheck nleapfrog == 2 || throw(BoundsError())

    @inbounds for j in 1:nx
        # size(A) = mx x nx x nlev, nlev = 1 is uppermost model level
        # apply drag only to largest zonal wavenumber (mx = 1)
        tendency[1,j,1] = tendency[1,j,1] - sdrag*A[1,j,1,1]
    end
end

# """Orographic temperature correction for absolute temperature to be applied before the horizontal diffusion."""
# function temperature_correction!(   Tabs_corrected::AbstractArray{Complex{NF},3},   # Corrected abs temperature T
#                                     Tabs::AbstractArray{Complex{NF},4},             # Absolute temperature
#                                     l::Int,                                         # leapfrog index
#                                     HD::HorizontalDiffusion{NF}                     # struct for correction arrays
#                                     ) where NF

#     @unpack tcorh, tcorv = HD   # load horizontal (tcorh) and vertical (tcorv) correction arrays

#     mx,nx,nlev,nleapfrog = size(Tabs)
#     @boundscheck (mx,nx,nlev) == size(Tabs_corrected) || throw(BoundsError())
#     @boundscheck (mx,nx) == size(tcorh) || throw(BoundsError())
#     @boundscheck (nlev,) == size(tcorv) || throw(BoundsError())
#     @boundscheck nleapfrog == 2 || throw(BoundsError())
#     @boundscheck l in [1,2] || throw(BoundsError())

#     @inbounds for k in 1:nlev
#         for j in 1:nx
#             for i in 1:mx
#                 Tabs_corrected[i,j,k] = Tabs[i,j,k,l] + tcorh[i,j]*tcorv[k]
#             end
#         end
#     end
# end

# """Orographic temperature correction for absolute temperature to be applied before the horizontal diffusion."""
# function humidity_correction!(  humid_corrected::AbstractArray{Complex{NF},3},  # Corrected abs temperature T
#                                 humid::AbstractArray{Complex{NF},4},            # Absolute temperature
#                                 l::Int,                                         # leapfrog index
#                                 HD::HorizontalDiffusion{NF}                     # struct for correction arrays
#                                 ) where NF

#     @unpack qcorh, qcorv = HD       # load horizontal (qcorh) and vertical (qcorv) correction arrays

#     mx,nx,nlev,nleapfrog = size(humid)
#     @boundscheck (mx,nx,nlev) == size(humid_corrected) || throw(BoundsError())
#     @boundscheck (mx,nx) == size(qcorh) || throw(BoundsError())
#     @boundscheck (nlev,) == size(qcorv) || throw(BoundsError())
#     @boundscheck nleapfrog == 2 || throw(BoundsError())
#     @boundscheck l in [1,2] || throw(BoundsError())

#     @inbounds for k in 1:nlev
#         for j in 1:nx
#             for i in 1:mx
#                 humid_corrected[i,j,k] = humid[i,j,k,l] + qcorh[i,j]*qcorv[k]
#             end
#         end
#     end
# end

"""Orographic temperature correction for absolute temperature to be applied before the horizontal diffusion."""
function orographic_correction!(A_corrected::AbstractArray{Complex{NF},3},  # Corrected variable
                                A::AbstractArray{Complex{NF},4},            # Variable (temperature or humidity)
                                l::Int,                                     # leapfrog index
                                hori_correction::AbstractArray,             # horizontal correction array
                                vert_correction::AbstractArray,             # vertical correction array
                                ) where NF

    mx,nx,nlev,nleapfrog = size(A)
    @boundscheck (mx,nx,nlev) == size(A_corrected) || throw(BoundsError())
    @boundscheck (mx,nx) == size(hori_correction) || throw(BoundsError())
    @boundscheck (nlev,) == size(vert_correction) || throw(BoundsError())
    @boundscheck nleapfrog == 2 || throw(BoundsError())
    @boundscheck l in [1,2] || throw(BoundsError())

    @inbounds for k in 1:nlev
        for j in 1:nx
            for i in 1:mx
                A_corrected[i,j,k] = A[i,j,k,l] + hori_correction[i,j]*vert_correction[k]
            end
        end
    end
end


text:
# module Jive

"""
    Mock

The Mock store for mocked modules.
"""
module Mock
end # Jive.Mock

function using_symbols(modul::Module, name::Symbol, used::Vector{Symbol}, funcs::Vector{Symbol})
    filter(names(modul, all=true)) do sym
        if sym in used
            false
        elseif sym in funcs
            true
        else
            first(String(sym)) != '#' && !(sym in (name, :eval, :include)) && !(getfield(modul, sym) isa Function)
        end
    end
end

function diff_symbols(modul::Module, nofuncmodul::Module)
    modulnames = names(modul, all=true)
    nofuncnames = names(nofuncmodul, all=true)
    filter(setdiff(modulnames, nofuncnames)) do sym
        first(String(sym)) != '#' && (getfield(modul, sym) isa Function)
    end
end

function mockup_a_module(modul::Module, expr::Expr, used::Vector{Symbol}, funcs::Vector{Symbol})
    name = expr.args[2]
    body = expr.args[3]
    syms = using_symbols(modul, name, used, funcs)
    using_exprs = Vector{Expr}()
    for sym in syms
        ex = Expr(:using, Expr(:(:),
            Expr(:., :., :., :., name),
            Expr(:., sym)))
        push!(used, sym)
        push!(using_exprs, ex)
    end
    pushfirst!(body.args, using_exprs...)
    toplevel = Expr(:toplevel,
        Expr(:module, true, :Mock, Expr(:block,
            Expr(:module, true, name, body)),
    ))
    toplevel
end

"""
    @mockup

used to produce a replica from the other module.
"""
macro mockup(expr::Expr)
    name = expr.args[2]
    modul = getfield(__module__, name)
    local mockmodul
    local oldstderr = stderr
    local errread, errwrite, errstream
    try
        (errread, errwrite) = redirect_stderr()
        used = Vector{Symbol}()
        ex = mockup_a_module(modul, expr, used, Symbol[])
        nofuncmodul = getfield(__module__.eval(ex), name)
        funcs = diff_symbols(modul, nofuncmodul)
        mockex = mockup_a_module(modul, expr, used, funcs)
        # @info :mockex mockex
        mockmodul = getfield(__module__.eval(mockex), name)
        errstream = @async read(errread, String)
    finally
        redirect_stderr(oldstderr)
        close(errwrite)
    end
    errmsg = fetch(errstream)
    if !isempty(errmsg)
        errmsg = replace(errmsg, "WARNING: replacing module Mock.\n" => "")
        print(stderr, errmsg)
    end
    mockmodul
end

# module Jive


text:
function run_all_tests()
    include("test_plots.jl")
end

# Test file (temporary) placement from Luxor.jl
if get(ENV, "MECHGLUEPLOTS_KEEP_TEST_RESULTS", false) == "true"
    opwd = pwd()
    cd(mktempdir(cleanup=false))
    @info("...Keeping the test plots in: $(pwd())")
    if Sys.iswindows()
        run(`cmd /c explorer .`, wait=false)
    end
    try
        run_all_tests()
    finally
        println("Switch back to original directory ")
        cd(opwd)
    end
    @info("Test plots were saved in: $(pwd())")
else
    mktempdir() do tmpdir
        cd(tmpdir) do
            @info("...Running tests in: $(pwd())")
            @info("but not keeping the results")
            @info("because you didn't do: ENV[\"MECHGLUEPLOTS_KEEP_TEST_RESULTS\"] = \"true\"")
            run_all_tests()
            @info("Test images weren't saved. To see the test images, next time do this before running:")
            @info(" ENV[\"MECHGLUEPLOTS_KEEP_TEST_RESULTS\"] = \"true\"")
        end
    end
end


text:
# Copyright 2015, 2016 Martin Holters
# See accompanying license file.

export resistor, potentiometer, capacitor, inductor, transformer,
       voltagesource, currentsource,
       voltageprobe, currentprobe, diode, bjt, opamp

resistor(r) = Element(mv=-1, mi=r)

potentiometer(r, pos) = Element(mv=-eye(2), mi=[r*pos 0; 0 r*(1-pos)],
                                pins=[1, 2, 2, 3])
potentiometer(r) =
    Element(mv=[eye(2); zeros(3, 2)], mi=[zeros(2, 2); eye(2); zeros(1, 2)],
            mq=-eye(5), mu=[zeros(4, 1); -1],
            nonlinear_eq = quote
                let v1=q[1], v2=q[2], i1=q[3], i2=q[4], pos=q[5]
                    res[1] = v1 - $(r)*pos*i1
                    res[2] = v2 - $(r)*(1-pos)*i2
                    J[1,1] = 1
                    J[1,2] = 0
                    J[1,3] = $(-r)*pos
                    J[1,4] = 0
                    J[1,5] = $(-r)*i1
                    J[2,1] = 0
                    J[2,2] = 1
                    J[2,3] = 0
                    J[2,4] = $(-r)*(1-pos)
                    J[2,5] = $(-r)*i2
                end
            end,
            pins=[1, 2, 2, 3])

capacitor(c) = Element(mv=[c;0], mi=[0; 1], mx=[-1;0], mxd=[0;-1])
inductor(l) = Element(mv=[1;0], mi=[0; l], mx=[0;-1], mxd=[-1;0])
transformer(l1, l2; coupling_coefficient=1,
            mutual_coupling=coupling_coefficient*sqrt(l1*l2)) =
    Element(mv=[1 0; 0 1; 0 0; 0 0],
            mi=[0 0; 0 0; l1 mutual_coupling; mutual_coupling l2],
            mx=[0 0; 0 0; -1 0; 0 -1], mxd=[-1 0; 0 -1; 0 0; 0 0],
            pins=[:primary1; :primary2; :secondary1; :secondary2])

function transformer(::Type{Val{:JA}}; D=2.4e-2, A=4.54e-5, ns=[],
                     a=14.1, α=5e-5, c=0.55, k=17.8, Ms=2.75e5)
    const μ0 = 1.2566370614e-6
    nonlinear_eq = quote
        coth_q1 = coth(q[1])
        a_q1 = abs(q[1])
        L_q1 = a_q1 < 1e-4 ? q[1]/3 : coth_q1 - 1/q[1]
        Ld_q1 = a_q1 < 1e-4 ? 1/3 : 1/q[1]^2 - coth_q1^2 + 1
        Ld2_q1 = a_q1 < 1e-3 ? -2/15*q[1] : 2*coth_q1*(coth_q1^2 - 1) - 2/q[1]^3
        δ = q[3] > 0 ? 1.0 : -1.0

        Man = $(Ms)*L_q1
        δM = sign(q[3]) == sign(Man - q[2]) ? 1.0 : 0.0

        den = δ*$(k*(1-c))-$(α)*(Man-q[2])
        # at present, the error needs to be scaled to be comparable to those of
        # the other elements, hence the factor 1e-4/Ms
        res[1] = $(1e-4/Ms) * ($(1-c) * δM*(Man-q[2])/den * q[3]
                               + $(c*Ms/a)*(q[3]+$(α)*q[4])*Ld_q1 - q[4])
        J[1,1] = $(1e-4/Ms) * ($((1-c)^2*k*Ms) * δM*Ld_q1*δ/den^2 * q[3]
                               + $(c*Ms/a)*(q[3]+$(α)*q[4])*Ld2_q1)
        J[1,2] = $(1e-4/Ms) * -$((1-c)^2*k) * δM*δ/den^2 * q[3]
        J[1,3] = $(1e-4/Ms) * ($(1-c) * δM*(Man-q[2])/den + $(c*Ms/a)*Ld_q1)
        J[1,4] = $(1e-4/Ms) * ($(c * Ms/a * α)*Ld_q1 - 1)
    end
    Element(mv=[speye(length(ns)); spzeros(5, length(ns))],
            mi=[spzeros(length(ns), length(ns)); ns.'; spzeros(4, length(ns))],
            mx=[spzeros(length(ns), 2); -π*D 0; -1/a -α/a; 0 -1; 0 0; 0 0],
            mxd=[-μ0*A*ns -μ0*ns*A; 0 0; 0 0; 0 0; -1 0; 0 -1],
            mq=[zeros(length(ns)+1,4); eye(4)], nonlinear_eq = nonlinear_eq)
end

inductor(::Type{Val{:JA}}; n=230, args...) =
    transformer(Val{:JA}; ns=[n], args...)

voltagesource(v) = Element(mv=1, u0=v, pins=[:+; :-])
voltagesource() = Element(mv=1, mu=1, pins=[:+; :-])
currentsource(i) = Element(mi=-1, u0=i, pins=[:+; :-])
currentsource() = Element(mi=-1, mu=1, pins=[:+; :-])

voltageprobe() = Element(mi=1, pv=1, pins=[:+; :-])
currentprobe() = Element(mv=1, pi=1, pins=[:+; :-])

diode(;is::Real=1e-12, η::Real = 1) =
  Element(mv=[1;0], mi=[0;1], mq=[-1 0; 0 -1], pins=[:+; :-], nonlinear_eq =
    quote
      let v = q[1], i = q[2], ex = exp(v*$(1./(25e-3 * η)))
        res[1] = $(is) * (ex - 1) - i
        J[1,1] = $(is/(25e-3 * η)) * ex
        J[1,2] = -1
      end
    end
  )

function bjt(typ; is=1e-12, η=1, isc=is, ise=is, ηc=η, ηe=η, βf=1000, βr=10)
    local polarity
    if typ == :npn
        polarity = 1
    elseif typ == :pnp
        polarity = -1
    else
        throw(ArgumentError(string("Unknown bjt type ", typ,
                                   ", must be :npn or :pnp")))
    end
    αf = -βf/(1+βf)
    αr = -βr/(1+βr)
    nonlinear_eq =
        quote
            let vE = q[1], vC = q[2], iE = q[3], iC = q[4],
                expE=exp(vE*$(1/(25e-3*ηe))), expC=exp(vC*$(1/(25e-3*ηc)))

                res[1] = $(ise) * (expE-1) + $(αr*isc) * (expC-1) - iE
                res[2] = $(αf*ise) * (expE-1) + $(isc) * (expC-1) - iC
                J[1,1] = $(ise/(25e-3*ηe)) * expE
                J[1,2] = $(αr*isc/(25e-3*ηc)) * expC
                J[1,3] = -1
                J[1,4] = 0
                J[2,1] = $(αf*ise/(25e-3*ηe)) * expE
                J[2,2] = $(isc/(25e-3*ηc)) * expC
                J[2,3] = 0
                J[2,4] = -1
            end
        end
    return Element(mv=[1 0; 0 1; 0 0; 0 0], mi = [0 0; 0 0; 1 0; 0 1],
                   mq = -polarity*speye(4), nonlinear_eq = nonlinear_eq,
                   pins = [:base; :emitter; :base; :collector])
end

opamp() = Element(mv=[0 0; 1 0], mi=[1 0; 0 0],
                  pins=["in+", "in-", "out+", "out-"])

function opamp(::Type{Val{:macak}}, gain, vomin, vomax)
    offset = 0.5 * (vomin + vomax)
    scale = 0.5 * (vomax - vomin)
    nonlinear_eq =
        quote
            let vi = q[1], vo = q[2], vi_scaled = vi * $(gain/scale)
                res[1] = tanh(vi_scaled) * $(scale) - vo
                J[1,1] = $(gain) / cosh(vi_scaled)^2
                J[1,2] = -1
            end
        end
    return Element(mv=[0 0; 1 0; 0 1], mi=[1 0; 0 0; 0 0], mq=[0 0; -1 0; 0 -1],
                   u0=[0; 0; offset],
                   nonlinear_eq = nonlinear_eq,
                   pins=["in+", "in-", "out+", "out-"])
end

@Base.deprecate(opamp_ideal, opamp)
@Base.deprecate(opamp_macak(gain, vomin, vomax), opamp(Val{:macak}, gain, vomin, vomax))


text:
function set_nsegs!(n::Neuron,frequency=100.0,d_lambda=.1)

    n.internal_nodes=[Array(Int64,0) for i=1:4]
    n.nodes=Array(Node,0)

    nodesec=zeros(Int64,length(n.secs)+1)
    nseglist=zeros(Int64,length(n.secs))
    nodesec[1]=1

    #find total number of segments
    for i=1:length(n.secs)

        lambda=lambda_f(frequency, n.secs[i], n)
        nseg=floor(Int64, (n.secs[i].length / (d_lambda * lambda) + .9) /2 )*2+1

        nodesec[i+1]=nodesec[i]+nseg

        nseglist[i]=nseg
    end

    first=0
    last=0
    for i=1:length(n.secs)

        if n.secs[i].mtype>1

            for j=1:nseglist[i]

                (area, ri,mypt3d) = r_a_calc(n.secs[i],j,nseglist[i],n.Ra)

                if (j==1)&&(j==nseglist[i])

                    first=length(n.nodes)+1
                    last=length(n.nodes)+2
                    single_node(n,nodesec,area,ri,mypt3d,i)

                elseif j==1

                    first=length(n.nodes)+1
                    first_node(n,nodesec,area,ri,mypt3d,i)

                elseif j==nseglist[i]

                    last=length(n.nodes)+2
                    last_node(n,area,ri,mypt3d,nodesec,i)
                else
                    middle_node(n,area,ri,mypt3d,i)
                end
            end
        else
            #approximate soma as 1 segment with everything connected at middle
            first=length(n.nodes)+1
            last=length(n.nodes)+3

            parent=length(n.nodes)+2
            children=Array(Int64,0)
            edge_node(n,parent,children,length(n.secs[end].pt3d))

            parent=0
            children=edge_children(n,nodesec,i)
            (area, ri, mypt3d) = r_a_calc(n.secs[i],1,1,n.Ra)
            push!(children,length(n.nodes))
            push!(children,length(n.nodes)+2)
            internal_node(n,area,ri,parent,children,mypt3d,i)

            parent=length(n.nodes)
            children=Array(Int64,0)
            edge_node(n,parent,children,1)
        end
        Section!(n.secs,i,first:last)
    end
    nothing
end

function lambda_f(frequency::Float64,sec::Section,neuron::Neuron)

    x1 = 0.0
    d1 = sec.pt3d[1].d
    lam = 0.0

    for i=2:length(sec.pt3d)
        x2=sec.pt3d[i].arc*sec.length
        d2=sec.pt3d[i].d
        lam+= (x2-x1)/sqrt(d1+d2)
        x1=x2
        d1=d2
    end

    lam *= sqrt(2) * 1e-5 * sqrt(4*pi*frequency*neuron.Ra*neuron.Cm)

    sec.length/lam

end

internal_node(n,a,ri,p,c,pt,i)=(nd=Node(a,ri,p,c,pt); add_node(n,nd); add_inode(n,i))

edge_node(n,p,c,pt)=(nd=Node(p,c,pt); add_node(n,nd))

edge_children(n,ns,i)=Int64[ns[n.secs[i].child[k]]+n.secs[i].child[k]-1 for k=1:length(n.secs[i].child)]

first_par(n,ns,i)=ns[n.secs[i].parent+1]+n.secs[i].parent-1

first_childs(n)=Int64[length(n.nodes)+2]

first_node(n,ns,area,ri,pt,i)=(p=first_par(n,ns,i); c=first_childs(n); internal_node(n,area,ri,p,c,pt,i))

middle_node(n,area,ri,pt,i)=(p=length(n.nodes);c=first_childs(n);internal_node(n,area,ri,p,c,pt,i))

function last_node(n,area,ri,pt,ns,i)
    p=length(n.nodes)
    c=edge_children(n,ns,i)
    internal_node(n,area,ri,p,[length(n.nodes)+2],pt,i)
    edge_node(n,length(n.nodes),c,pt[end])
end

function single_node(n,ns,area,ri,pt,i)
    p=first_par(n,ns,i)
    if length(n.secs[i].child)==0
        c=Array(Int64,0)
    else
        c=edge_children(n,ns,i)
    end
    internal_node(n,area,ri,p,[length(n.nodes)+2],pt,i)
    edge_node(n,length(n.nodes),c,pt[end])
end

function r_a_calc(sec::Section,x::Int64,nseg::Int64,Ra::Float64)

    #find location of other nodes on section in normalized units

    cent=(2*x-1)/(2*nseg)
    interval=Float64[(x-1)/nseg, x/nseg]

    #get all of the arc lengths into one data structures
    arc3d=collect(Float64[sec.pt3d[i].arc for i=1:length(sec.pt3d)])

    #Find the first 3d points 1) before the segment, 2) after the middle of the segment and 3) before the end of the segment
    first=findfirst(arc3d.>interval[1])-1
    mid=findfirst(arc3d.>cent)-1
    last=findnext(arc3d.>=interval[2],first)-1

    if length(arc3d)<4

        #If just a few points, whole thing is just modeled as a cylinder
        myarea=2*pi*sec.pt3d[1].d/2*sec.length*(interval[2]-interval[1])

        area=SegArea(myarea/2,myarea/2,myarea)

        r=frustrum_resistance(sec.pt3d[1].d,sec.pt3d[1].d,sec.length*(interval[2]-interval[1])/2,Ra)
        ri=SegRi(.01*r,.01*r)

        first=1
        last=length(arc3d)

    else

        al=0.0
        ar=0.0
        rl=0.0
        rr=0.0

    #=
        Area of a segment is modeled as a series of truncated cones between the 3d points. The diameter of the top and bottom of the frustrum is taken to be the average of the diameters of the nearest 3d points
    =#

        #left side to first point before center
        frac=interp_area(arc3d[first],interval[1],arc3d[first+1])

        diam=frac[1]*sec.pt3d[first].d+frac[2]*sec.pt3d[first+1].d
        height=(arc3d[first+1]-interval[1])*sec.length

        for i=(first+1):mid
            al+=frustrum_area(diam, sec.pt3d[i].d,height)
            rl+=frustrum_resistance(diam,sec.pt3d[i].d,height,Ra)
            diam=sec.pt3d[i].d
            height=(arc3d[i+1]-arc3d[i])*sec.length
        end

        #First point before center to center
        frac=interp_area(arc3d[mid], cent, arc3d[mid+1])

        diam=frac[1]*sec.pt3d[mid].d+frac[2]*sec.pt3d[mid+1].d
        height=(cent-arc3d[mid])*sec.length

        al+=frustrum_area(sec.pt3d[mid].d,diam,height)
        rl+=frustrum_resistance(sec.pt3d[mid].d,diam,height,Ra)

        #first point after center
        height=(arc3d[mid+1]-cent)*sec.length

        ar+=frustrum_area(diam,sec.pt3d[mid+1].d,height)
        rr+=frustrum_resistance(diam,sec.pt3d[mid+1].d,height,Ra)

        diam=sec.pt3d[mid+1].d
        height=(arc3d[mid+2]-arc3d[mid+1])*sec.length

        for i=(mid+2):last
            ar+=frustrum_area(diam,sec.pt3d[i].d,height)
            rr+=frustrum_resistance(diam,sec.pt3d[i].d,height,Ra)
            diam=sec.pt3d[i].d
            height=(arc3d[i+1]-arc3d[i])*sec.length
        end

        #Final point
        frac=interp_area(arc3d[last], interval[2], arc3d[last+1])

        diam=frac[1]*sec.pt3d[last].d+frac[2]*sec.pt3d[last+1].d
        height=(interval[2]-arc3d[last])*sec.length

        ar+=frustrum_area(sec.pt3d[last].d,diam,height)
        rr+=frustrum_resistance(sec.pt3d[last].d,diam,height,Ra)

        area=SegArea(al,ar,al+ar)
        ri=SegRi(.01*rl,.01*rr)
    end

    (area, ri, first:last)

end

frustrum_area(d1::Float64,d2::Float64,h::Float64)=pi*(d1/2+d2/2)*sqrt((d1/2-d2/2)^2+h^2)

frustrum_resistance(d1::Float64,d2::Float64,h::Float64,ri::Float64)=ri*h/(pi*(d1/2)*(d2/2))

interp_area(x1::Float64, x2::Float64, x3::Float64)=((x2-x1)/(x3-x1),(x3-x2)/(x3-x1))

max_stars_repo_name: paulmthompson/JNeuron

max_issues_repo_name: paulmthompson/JNeuron

max_forks_repo_name: paulmthompson/JNeuron


text:
using DelimitedFiles

# n = number of events
# prob = probability of incrementing counter
function random_count(n, prob; a=30, prob_calc=false)
    v = 0
    base = 1+1/a
    output = zeros(n)

    if prob_calc
        prob = 1
        v = 1
    end

    if rand() < prob
        output[1] = 1/prob
    end

    for i = 2:n
        if prob_calc
            prob = 1/((a*(base^(v+1)-1))-output[i-1])

            if rand() <= prob
                v += 1
            end
            output[i] = a*(base^v-1)
        else
            if rand() <= prob
                output[i] = output[i-1]+1/prob
            else
                output[i] = output[i-1]
            end
        end
    end

    return output
end

# m = number of counting trials
# l = number of saved trials
function multi_count(n, m, l, prob; a=30, prob_calc=false, file_mod="",
                     stops = [100000, 500000, 1000000])
    out = zeros(n, l)
    extremes = zeros(n, 2)

    chosen_set = zeros(Int, l)

    for i = 1:l
        chosen_number = rand(1:m)
        while chosen_number in chosen_set
            chosen_number = rand(1:m)
        end
        chosen_set[i] = chosen_number
    end

    histograms = zeros(Float64, m, length(stops))

    out_count = 1
    for i = 1:m
        current_dist = random_count(n, prob; a, prob_calc)
        if i == 1
            extremes[:,1] .= current_dist
            extremes[:,2] .= current_dist
        else
            for j = 1:n
                if current_dist[j] < extremes[j,1]
                    extremes[j,1] = current_dist[j]
                end
                if current_dist[j] > extremes[j,2]
                    extremes[j,2] = current_dist[j]
                end

            end
        end

        for j = 1:length(stops)
            histograms[i,j] = current_dist[stops[j]]
        end

        if i in chosen_set
            out[:,out_count] = current_dist
            out_count += 1
        end
    end

    output_file = open("out"*file_mod*".dat", "w")
    writedlm(output_file, out)
    close(output_file)

    extreme_output_file = open("extremes"*file_mod*".dat", "w")
    writedlm(extreme_output_file, extremes)
    close(extreme_output_file)

    formatted_histograms = [[],[],[]]
    # Going through the histogram data to put it into the right format
    for j = 1:length(stops)

        max = floor(Int,maximum(histograms[:,j]))
        min = floor(Int,minimum(histograms[:,j]))

        println(min, '\t', max, '\t', sum(histograms[:,j])/m)

        temp_array = zeros(max - min+1)

        for i = 1:m

            temp_array[floor(Int,histograms[i,j])-min+1] += 1

        end

        formatted_histograms[j] = temp_array
    end

    # output histograms into different files for each one

    for i = 1:length(stops)
        histogram_file = open("histogram_" * string(i)*file_mod* ".dat", "w")
        writedlm(histogram_file, formatted_histograms[i])
        close(histogram_file)
    end
    return formatted_histograms
end


text:
include_backtrace() = false

function Try.enable_errortrace()
    old = include_backtrace()
    @eval include_backtrace() = true
    return old
end

function Try.disable_errortrace()
    old = include_backtrace()
    @eval include_backtrace() = false
    return old
end

function Try.enable_errortrace(yes::Bool)
    include_backtrace() == yes && return yes
    return yes ? Try.enable_errortrace() : Try.disable_errortrace()
end

maybe_backtrace() = include_backtrace() ? backtrace() : nothing

function common_suffix(bt, here = backtrace())
    a = lastindex(bt)
    b = lastindex(here)
    while true
        a < firstindex(bt) && return a  # unreachable?
        b < firstindex(bt) && break
        bt[a] != here[b] && break
        a -= 1
        b -= 1
    end
    return a
end

function simplify_backtrace_impl(bt)
    a = common_suffix(bt)
    no_common_suffix = a == lastindex(bt)
    bt = Base.process_backtrace(bt[1:a])

    j = lastindex(bt)
    if no_common_suffix
        # No common suffix. Evaluated in REPL's display?
        fr, = bt[end]
        if fr.func === :_start && basename(string(fr.file)) == "client.jl"
            j = findlast(((fr, _),) -> !fr.from_c && fr.func === :eval, bt)
            j = max(firstindex(bt), j - 1)
        end
    end

    i = firstindex(bt)
    if bt[i][1].func === :maybe_backtrace
        i += 1
    end
    if bt[i][1].func === :Err
        i += 1
    end

    return bt[i:j]
end

function simplify_backtrace(bt)
    try
        return simplify_backtrace_impl(bt)
    catch err
        @error(
            "Fail to simplify backtrace. Fallback to plain backtrace.",
            exception = (err, catch_backtrace()),
            maxlog = 5,
        )
        return bt
    end
end

struct ErrorTrace <: Exception
    exception::Exception
    backtrace::typeof(Base.backtrace())
end

function Base.showerror(io::IO, errtrace::ErrorTrace)
    print(io, "Original Error: ")
    showerror(io, errtrace.exception)
    println(io)

    # TODO: remove common prefix?
    buffer = IOBuffer()
    Base.show_backtrace(IOContext(buffer, io), simplify_backtrace(errtrace.backtrace))
    seekstart(buffer)
    println(io, "┌ Original: stacktrace")
    for ln in eachline(buffer)
        print(io, "│  ")
        println(io, ln)
    end
    println(io, "└")
end

function Base.show(io::IO, errtrace::ErrorTrace)
    Base.show(io, ErrorTrace)
    print(io, '(')
    Base.show(io, errtrace.exception)
    print(io, ", …)")
end


text:
# This file defines the Transition type. This type defines a transition between
#   states that a personnel member can make.

# The Tansition type requires the State type.
requiredTypes = [ "state" ]

for reqType in requiredTypes
    if !isdefined( Symbol( uppercase( string( reqType[ 1 ] ) ) * reqType[ 2:end ] ) )
        include( joinpath( typePath, reqType * ".jl" ) )
    end  # if !isdefined( Symbol( ...
end  # for reqType in requiredTypes

export Transition
"""
This type defines a transition between two states that a personnel member can
perform, along with all the necessary information  about the conditions of the
transition.

The type contains the following fields:
* `name::string`: the name of the transition.
* `startState::State`: the state that the personnel member is currently in.
* `endState::State`: the state that the personnel member can attain.
* `freq::Float64`: the time between two checks in the transition's schedule.
* `offset::Float64`: the offset of the transition's schedule with respect to the
  start of the simulation.
* `extraConditions::Vector{Condition}`: the extra conditions that must be
  satisfied before the transition can take place.
* `extraChanges::Vector{PersonnelAttribute}`: the extra changes to attributes
  that happen during the transition.
* `probabilityList::Vector{Float64}`: the list of probabilities for this
  transition to occur.
* `maxAttempts::Int`: the maximum number of tries a personnel member has to
  undergo the transition.
* `timeBetweenAttempts::Float64`: the minimum time between two successive
  attempts.
* `isFiredOnFail::Bool`: a flag indicating that the personnel member gets fired
  if he fails to make the transition in the provided number of attempts.
* `maxFlux::Int`: the maximum number of people that can undergo the transition
  at the same time.
* `hasPriority::Bool`: a flag indicating that this transition can override the
  target population of the transition's target state. If the flag is `true`, it
  means that if the max flux of the state is 15, and only 10 spots are available
  in the target state, 15 people will undergo the transition nonetheless. If the
  flag is `false`, 10 persons would.
* `transPriority::Int`: the priority in the simulation on which the transition
  gets executed. This priority will be 0 for transitions with the `hasPriority`
  flag set to `true`, and < 0 otherwise. A priority == 1 means it needs to be
  determined first.
"""
type Transition

    name::String
    startState::State
    endState::State
    freq::Float64
    offset::Float64
    extraConditions::Vector{Condition}
    extraChanges::Vector{PersonnelAttribute}
    probabilityList::Vector{Float64}
    maxAttempts::Int
    timeBetweenAttempts::Float64
    isFiredOnFail::Bool
    maxFlux::Int
    hasPriority::Bool
    transPriority::Int

    # Basic constructor.
    function Transition( name::String, startState::State, endState::State;
        freq::T1 = 1.0, offset::T2 = 0.0, minTime::T3 = 1.0,
        maxAttempts::T4 = 1, maxFlux::T5 = -1, isFiredOnFail::Bool = false ) where T1 <: Real where T2 <: Real where T3 <: Real where T4 <: Integer where T5 <: Integer

        if freq <= 0.0
            error( "Time between two transition checks must be > 0.0" )
        end  # if freq <= 0.0

        if minTime <= 0.0
            error( "Time that personnel member must have initial state must be > 0.0" )
        end  # if minTime <= 0.0

        if maxAttempts < -1
            error( "Maximum number of attempts must be => -1, where -1 stands for infinite" )
        end  # if maxAttempts < -1

        if maxFlux < -1
            error( "Maximum flux must be => -1, where -1 stands for infinite" )
        end

        newTrans = new()
        newTrans.name = name
        newTrans.startState = startState
        newTrans.endState = endState
        newTrans.freq = freq
        newTrans.offset = offset % freq + ( offset < 0.0 ? freq : 0.0 )
        newTrans.extraConditions = Vector{Condition}()
        newTrans.extraChanges = Vector{PersonnelAttribute}()
        newTrans.probabilityList = [ 1.0 ]
        newTrans.maxAttempts = maxAttempts
        newTrans.timeBetweenAttempts = 0.0
        newTrans.isFiredOnFail = isFiredOnFail
        newTrans.maxFlux = -1
        newTrans.hasPriority = false
        newTrans.transPriority = 1
        return newTrans

    end  # Transition( name, startState, endState; freq, offset, minTime )

end  # type Transition


text:
#studies
#############
@with_kw struct MetricStudy
    actiondistr     = UniformActions()
    G               = TwoBumps()
    n_seeds::Int    = 500
    n_iters::Int    = 100
end
struct MetricStudyResult
    cum_regrets
    simple_regrets
end

function generate_sim_q(study::MetricStudy, alg)
    q = []
    for i in 1:study.n_seeds
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                outs=Set([:simple_regret, :cum_regret]))
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::MetricStudy)
    cum_regrets = []
    simple_regrets = []
    for alg in [RandomBandit, PWUCB, SBUCB, GPUCB]
		q = generate_sim_q(study, alg)
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    MetricStudyResult(cum_regrets, simple_regrets)
end
Plots.plot(result::MetricStudyResult, metric::Symbol) = plot(result, Val{metric})
function Plots.plot(result::MetricStudyResult, ::Type{Val{:cum_regret}})
    p = plot(; ylabel="cumulative regret", xlabel="number of samples", legend=:topleft)
    for (alg,cum_regret) in result.cum_regrets
        plot!(p, cum_regret, label=string(alg))
    end
    p
end
function Plots.plot(result::MetricStudyResult, ::Type{Val{:simple_regret}})
    p = plot(; ylabel="simple regret", xlabel="number of samples")
    for (alg,simple_regret) in result.simple_regrets
        plot!(p, simple_regret, label=string(alg))
    end
    p
end

#################
@with_kw struct SweepStudy
    actiondistr     = UniformActions()
    G               = NarrowBump()
    n_seeds::Int     = 100
    n_iters::Int     = 100
    ks              = linspace(0.0, 3.0, 10)
    αs              = linspace(0.0, 2.0, 10)
end

function generate_sim_q(study::SweepStudy, alg)
    q = []
    for i = 1:study.n_seeds
        b = RandomBandit(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                         outs=Set([:simple_regret,:cum_regret]))
        push!(q, BanditSim(b, study.G))
    end
    for i = 1:study.n_seeds
        b = PWUCB(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                         outs=Set([:simple_regret,:cum_regret]))
        push!(q, BanditSim(b, study.G))
    end
    for i = 1:study.n_seeds, k in study.ks, α in study.αs
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters, k=k, α=α,
                         outs=Set([:simple_regret,:cum_regret]))
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::SweepStudy, alg)
    q = generate_sim_q(study, alg)
    data = run_parallel(q) do sim, result
        return vcat(metadata(sim.b),
                    [:mean_cum_regret=>mean(result.cum_regret),
                     :mean_simple_regret=>mean(result.simple_regret)])
    end
    data
end

#############
@with_kw struct GPkMetricStudy
    actiondistr     = UniformActions()
    G               = TwoBumps()
    n_seeds::Int    = 500
    n_iters::Int    = 100
    ks::Vector{Int} = [100, 50, 20, 10, 5, 1]
end
struct GPkMetricStudyResult
    cum_regrets
    simple_regrets
end
function generate_sim_q(study::GPkMetricStudy, alg; kwargs...)
    q = []
    for i in 1:study.n_seeds
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters, outs=Set([:simple_regret, :cum_regret]),
                kwargs...)
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::GPkMetricStudy)
    cum_regrets = []
    simple_regrets = []
    for alg in [RandomBandit, PWUCB, SBUCB, GPUCBGrid]
		q = generate_sim_q(study, alg)
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for k in study.ks
		q = generate_sim_q(study, GPUCB; k=k)
		results = pmap(POMDPs.simulate, q)
        push!(cum_regrets, "$(string(GPUCB))-$k"=>mean(hcat([r.cum_regret for r in results]...),2))
        push!(simple_regrets, "$(string(GPUCB))-$k"=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    GPkMetricStudyResult(cum_regrets, simple_regrets)
end
Plots.plot(result::GPkMetricStudyResult, metric::Symbol) = plot(result, Val{metric})
function Plots.plot(result::GPkMetricStudyResult, ::Type{Val{:cum_regret}})
    p = plot(; xlabel="number of samples", ylabel="cumulative regret", legend=:topleft)
    for (alg,cum_regret) in result.cum_regrets
        plot!(p, cum_regret, label=string(alg))
    end
    p
end
function Plots.plot(result::GPkMetricStudyResult, ::Type{Val{:simple_regret}})
    p = plot(; xlabel="number of samples", ylabel="simple regret")
    for (alg,simple_regret) in result.simple_regrets
        plot!(p, simple_regret, label=string(alg))
    end
    p
end

################
@with_kw struct GPOptimMetricStudy
    actiondistr     = UniformActions()
    G               = TwoBumps()
    n_seeds::Int    = 500
    n_iters::Int    = 100
    mean_inits::Vector{Int} = [8.0,5.0,3.0,2.0,0.0]
end
struct GPOptimMetricStudyResult
    cum_regrets
    simple_regrets
end
function generate_sim_q(study::GPOptimMetricStudy, alg; kwargs...)
    q = []
    for i in 1:study.n_seeds
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                outs=Set([:simple_regret, :cum_regret]), kwargs...)
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::GPOptimMetricStudy)
    cum_regrets = []
    simple_regrets = []
    for alg in [RandomBandit, PWUCB, SBUCB, GPUCBGrid]
		q = generate_sim_q(study, alg)
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for minit in study.mean_inits
		q = generate_sim_q(study, GPUCB; mean_init=minit)
		results = pmap(POMDPs.simulate, q)
        push!(cum_regrets, "$(string(GPUCB))-$minit"=>mean(hcat([r.cum_regret for r in results]...),2))
        push!(simple_regrets, "$(string(GPUCB))-$minit"=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    GPOptimMetricStudyResult(cum_regrets, simple_regrets)
end
Plots.plot(result::GPOptimMetricStudyResult, metric::Symbol) = plot(result, Val{metric})
function Plots.plot(result::GPOptimMetricStudyResult, ::Type{Val{:cum_regret}})
    p = plot(; xlabel="number of samples", ylabel="cumulative regret", legend=:topleft)
    for (alg,cum_regret) in result.cum_regrets
        plot!(p, cum_regret, label=string(alg))
    end
    p
end
function Plots.plot(result::GPOptimMetricStudyResult, ::Type{Val{:simple_regret}})
    p = plot(; xlabel="number of samples", ylabel="simple regret")
    for (alg,simple_regret) in result.simple_regrets
        plot!(p, simple_regret, label=string(alg))
    end
    p
end

####
@with_kw struct MetricStudy2D
    actiondistr     = UniformActions2D()
    G               = WideBump2D()
    n_seeds::Int    = 20
    n_iters::Int    = 200
end
struct MetricStudy2DResult
    cum_regrets
    simple_regrets
end

function generate_sim_q(study::MetricStudy2D, alg; kwargs...)
    q = []
    for i in 1:study.n_seeds
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                outs=Set([:simple_regret, :cum_regret]), kwargs...)
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::MetricStudy2D)
    cum_regrets = []
    simple_regrets = []
    for alg in [RandomBandit, PWUCB, SBUCB]
		q = generate_sim_q(study, alg)
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for alg in [GPUCBGrid, GPUCB]
        q = generate_sim_q(study, alg, log_length_scale=-0.5, grid=GridActions2D())
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for alg in [GPUCBGrid, GPUCB]
        q = generate_sim_q(study, alg, mean_init=study.G.g_max+1.0, log_length_scale=-0.5, grid=GridActions2D())
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, "$(string(alg))-optim"=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, "$(string(alg))-optim"=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    MetricStudy2DResult(cum_regrets, simple_regrets)
end
Plots.plot(result::MetricStudy2DResult, metric::Symbol) = plot(result, Val{metric})
function Plots.plot(result::MetricStudy2DResult, ::Type{Val{:cum_regret}})
    p = plot(; ylabel="cumulative regret", xlabel="number of samples", legend=:topleft)
    for (alg,cum_regret) in result.cum_regrets
        plot!(p, cum_regret, label=string(alg))
    end
    p
end
function Plots.plot(result::MetricStudy2DResult, ::Type{Val{:simple_regret}})
    p = plot(; ylabel="simple regret", xlabel="number of samples")
    for (alg,simple_regret) in result.simple_regrets
        plot!(p, simple_regret, label=string(alg))
    end
    p
end

######
@with_kw struct LengthScaleStudy2D
    actiondistr     = UniformActions2D()
    G               = WideBump2D()
    n_seeds::Int    = 20
    n_iters::Int    = 200
    ls::Vector{Float64} = [0.05, 0.1, 0.3, 0.6, 0.8, 1.0, 1.5]
end
struct LengthScaleStudy2DResult
    cum_regrets
    simple_regrets
end

function generate_sim_q(study::LengthScaleStudy2D, alg; kwargs...)
    q = []
    for i in 1:study.n_seeds
        b = alg(; actiondistr=study.actiondistr, seed=i, n_iters=study.n_iters,
                outs=Set([:simple_regret, :cum_regret]), kwargs...)
        push!(q, BanditSim(b, study.G))
    end
    q
end
function run_study(study::LengthScaleStudy2D)
    cum_regrets = []
    simple_regrets = []
    for alg in [RandomBandit, PWUCB, SBUCB]
		q = generate_sim_q(study, alg)
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for alg in [GPUCBGrid]
        q = generate_sim_q(study, alg, log_length_scale=-0.5, grid=GridActions2D())
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, string(alg)=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, string(alg)=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    for len in study.ls, alg in [GPUCB]
        q = generate_sim_q(study, alg, log_length_scale=log(len), grid=GridActions2D())
		results = pmap(POMDPs.simulate, q)
		push!(cum_regrets, "$(string(alg))-l=$len"=>mean(hcat([r.cum_regret for r in results]...),2))
		push!(simple_regrets, "$(string(alg))-l=$len"=>mean(hcat([r.simple_regret for r in results]...),2))
    end
    LengthScaleStudy2DResult(cum_regrets, simple_regrets)
end
Plots.plot(result::LengthScaleStudy2DResult, metric::Symbol) = plot(result, Val{metric})
function Plots.plot(result::LengthScaleStudy2DResult, ::Type{Val{:cum_regret}})
    p = plot(; ylabel="cumulative regret", xlabel="number of samples", legend=:topleft)
    for (alg,cum_regret) in result.cum_regrets
        plot!(p, cum_regret, label=string(alg))
    end
    p
end
function Plots.plot(result::LengthScaleStudy2DResult, ::Type{Val{:simple_regret}})
    p = plot(; ylabel="simple regret", xlabel="number of samples")
    for (alg,simple_regret) in result.simple_regrets
        plot!(p, simple_regret, label=string(alg))
    end
    p
end