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


using JuMP
using MathOptInterface
using GLPK
using Test
using MathOptSetDistances
const MOI = MathOptInterface
const MOIU = MOI.Utilities
const MOD = MathOptSetDistances

variable_primal(model) = x -> MOI.get(model, MOI.VariablePrimal(), x)

model = Model(GLPK.Optimizer)
@variable(model, x>=1)
@variable(model, y>=0)
@constraint(model,c1, x + y >= 3)
@constraint(model,c2, x + y >= 2)
@objective(model, Min, x + y)
optimize!(model)

varval = variable_primal(backend(model))
f = JuMP.MOI.get(model, JuMP.MOI.ConstraintFunction(), c2)

delete(model,c2)
val  = MOIU.eval_variables(varval, f)

# m = HeuristicOptimizer.heuristicCreate(model)
# HeuristicOptimizer.removeConstrains!(m)
# optimize!(m.model)
# HeuristicOptimizer.updateFunctionValue!(m)
# HeuristicOptimizer.updateViolation!(m)
# @show m.violation

c = f, s


using Ipopt
using PowerModelsPrivacyPreserving
using PowerModels
using Distributions
using Random

const PM = PowerModels
const PMPP = PowerModelsPrivacyPreserving

ipopt = Ipopt.Optimizer
file =  "test/data/matpower/case5.m"
data_unpert = PM.parse_file(file)

for (i,gen) in data_unpert["gen"]
    gen["eta"] = 0.95
end
# data_min_loss = deepcopy(data_unpert)
# data_min_cost = deepcopy(data_unpert)

result_unpert_cost = PMPP.run_opf_bf_dvorkin(data_unpert, BFAPowerModel, ipopt)

##
# https://github.com/mlubin/RobustCCOPFSupplement/
# https://github.com/mlubin/RobustCCOPFSupplement/blob/master/codejl/ccopfmodel_simulation.jl
# Lubin, M., Dvorkin, Y., & Backhaus, S. (2015). A robust approach to chance constrained optimal power flow with renewable generation. IEEE Trans. Power Syst., 11. Cryptography and Security. Retrieved from http://arxiv.org/abs/1504.0601
# https://jumpchance.readthedocs.io/en/latest/
# https://github.com/mlubin/JuMPChance.jl
result_unpert_cost_cc = PMPP.run_opf_bf_dvorkin_cc(data_unpert, BFAPowerModel, ipopt)


#####
##### Divergence operators
#####

"""
    divᶜᶜᶜ(i, j, k, grid, u, v, w)

Calculates the divergence ∇·𝐔 of a vector field 𝐔 = (u, v, w),

    1/V * [δxᶜᵃᵃ(Ax * u) + δxᵃᶜᵃ(Ay * v) + δzᵃᵃᶜ(Az * w)],

which will end up at the cell centers `ccc`.
"""
@inline function divᶜᶜᶜ(i, j, k, grid, u, v, w)
    return 1 / Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u) +
                                      δyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v) +
                                      δzᵃᵃᶜ(i, j, k, grid, Az_qᶜᶜᶠ, w))
end

"""
    div_xyᶜᶜᵃ(i, j, k, grid, u, v)

Returns the discrete `div_xy = ∂x u + ∂y v` of velocity field `u, v` defined as

```
1 / Azᶜᶜᵃ * [δxᶜᵃᵃ(Δyᵃᶜᵃ * u) + δyᵃᶜᵃ(Δxᶜᵃᵃ * v)]
```

at `i, j, k`, where `Azᶜᶜᵃ` is the area of the cell centered on (Center, Center, Any) --- a tracer cell,
`Δy` is the length of the cell centered on (Face, Center, Any) in `y` (a `u` cell),
and `Δx` is the length of the cell centered on (Center, Face, Any) in `x` (a `v` cell).
`div_xyᶜᶜᵃ` ends up at the location `cca`.
"""
@inline function div_xyᶜᶜᶜ(i, j, k, grid, u, v)
    return 1 / Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, u) +
                                       δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, v))
end

# Convention
 index_left(i, ::Center) = i
 index_left(i, ::Face)   = i - 1
index_right(i, ::Center) = i + 1
index_right(i, ::Face)   = i

@inline Base.div(i, j, k, grid::AbstractGrid, loc, q_west, q_east, q_south, q_north, q_bottom, q_top) =
    1 / volume(i, j, k, grid, loc...) * (δx_Ax_q(i, j, k, grid, loc, q_west, q_east) +
                                         δy_Ay_q(i, j, k, grid, loc, q_south, q_north) +
                                         δz_Az_q(i, j, k, grid, loc, q_bottom, q_top))

@inline function δx_Ax_q(i, j, k, grid, (LX, LY, LZ), qᵂ, qᴱ)
    iᵂ = index_left(i, LX)
    Axᵂ = Ax(iᵂ, j, k, grid, LX, LY, LZ)

    iᴱ = index_right(i, LX)
    Axᴱ = Ax(iᴱ, j, k, grid, LX, LY, LZ)

    return Axᴱ * qᴱ - Axᵂ * qᵂ
end

@inline function δy_Ay_q(i, j, k, grid, (LX, LY, LZ), qˢ, qᴺ)
    jˢ = index_left(j, LY)
    Ayˢ = Ay(i, jˢ, k, grid, LX, LY, LZ)

    jᴺ = index_right(j, LY)
    Ayᴺ = Ay(i, jᴺ, k, grid, LX, LY, LZ)

    return Ayᴺ * qᴺ - Ayˢ * qˢ
end

@inline function δz_Az_q(i, j, k, grid, (LX, LY, LZ), qᴮ, qᵀ)
    kᴮ = index_left(k, LZ)
    Azᴮ = Az(i, j, kᴮ, grid, LX, LY, LZ)

    kᵀ = index_right(k, LZ)
    Azᵀ = Az(i, j, kᵀ, grid, LX, LY, LZ)

    return Azᵀ * qᵀ - Azᴮ * qᴮ
end

# And flat!

@inline δx_Ax_q(i, j, k, grid::XFlatGrid, args...) = zero(grid)
@inline δy_Ay_q(i, j, k, grid::YFlatGrid, args...) = zero(grid)
@inline δz_Az_q(i, j, k, grid::ZFlatGrid, args...) = zero(grid)


function isarmstrong(n::Int)
    digitbydigit = digits(n)
    len = length(digitbydigit)
    n == sum(digitbydigit.^len)
end


using LowRank, Test
using Random, Distributions

N = 50
d = 5
n = 10
k = 2

# non-zero error
A = zeros(d, N)
U = rand(Uniform(-1.0, 1.0), d-1, k)
V = rand(Uniform(-1.0, 1.0), k, n)
A[1:d-1, 1:n] = U*V
A[d,1:n] = rand(Uniform(-1.0, 1.0), 1, n)
A[:, n+1:N] = rand(Uniform(-1.0, 1.0), d , N-n)
error, _ = LowRank.truncated_svd(A[:,1:n], k)
data_set = collect(eachcol(A))

@test find_low_rank_subset_sample_rep(data_set,n,k)[1] <= 4*error
@test FastLowRank(data_set, n, k)[1] <= error + 0.0001


using JuMP, EAGO, Gurobi, CSV, DataFrames, Ipopt

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[1, 2, 3, 4, 5, 6, 7, 8]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[1], 0.1)
JuMP.set_upper_bound(x[1], 10.0)
JuMP.set_lower_bound(x[2], 0.1)
JuMP.set_upper_bound(x[2], 10.0)
JuMP.set_lower_bound(x[3], 0.1)
JuMP.set_upper_bound(x[3], 10.0)
JuMP.set_lower_bound(x[4], 0.1)
JuMP.set_upper_bound(x[4], 10.0)
JuMP.set_lower_bound(x[5], 0.1)
JuMP.set_upper_bound(x[5], 10.0)
JuMP.set_lower_bound(x[6], 0.1)
JuMP.set_upper_bound(x[6], 10.0)
JuMP.set_lower_bound(x[7], 0.1)
JuMP.set_upper_bound(x[7], 10.0)
JuMP.set_lower_bound(x[8], 0.1)
JuMP.set_upper_bound(x[8], 10.0)

# ----- Constraints ----- #
@NLconstraint(m, e2, 0.0588*x[5]*x[7]+0.1*x[1] <= 1.0)
@NLconstraint(m, e3, 0.0588*x[6]*x[8]+0.1*x[1]+0.1*x[2] <= 1.0)
@NLconstraint(m, e4, 4*x[3]/x[5]+2/(x[3]^0.71*x[5])+0.0588*x[7]/x[3]^1.3 <= 1.0)
@NLconstraint(m, e5, 4*x[4]/x[6]+2/(x[4]^0.71*x[6])+0.0588*x[4]^1.3*x[8] <= 1.0)

# ----- Objective ----- #
@NLobjective(m, Min,(0.4*x[1]^0.67/x[7]^0.67+0.4*x[2]^0.67/x[8]^0.67-x[1]-x[2]) + 10.0)

JuMP.optimize!(m)

backend_opt = JuMP.backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


"""
    niw_hyperparams(κ::Float32, m::AbstractArray{Float32}, ν::Float32, ψ::AbstractArray{Float32})

[Normal Inverse Wishart](https://en.wikipedia.org/wiki/Normal-inverse-Wishart_distribution)
"""
struct niw_hyperparams <: distribution_hyper_params
    κ::Float32
    m::AbstractArray{Float32}
    ν::Float32
    ψ::AbstractArray{Float32}
end

mutable struct niw_sufficient_statistics <: sufficient_statistics
    N::Float32
    points_sum::AbstractArray{Float32,1}
    S::AbstractArray{Float32,2}
end

function calc_posterior(prior:: niw_hyperparams, suff_statistics::niw_sufficient_statistics)
    if suff_statistics.N == 0
        return prior
    end
    κ = prior.κ + suff_statistics.N
    ν = prior.ν + suff_statistics.N
    m = (prior.m.*prior.κ + suff_statistics.points_sum) / κ
    ψ = (prior.ν * prior.ψ + prior.κ*prior.m*prior.m' -κ*m*m'+ suff_statistics.S) / ν
    ψ = Matrix(Hermitian(ψ))
    return niw_hyperparams(κ,m,ν,ψ)
end

function sample_distribution(hyperparams::niw_hyperparams)
    Σ = rand(Distributions.InverseWishart(hyperparams.ν, hyperparams.ν* hyperparams.ψ))
    μ = rand(Distributions.MvNormal(hyperparams.m, Σ/hyperparams.κ))
    invΣ = inv(Σ)
    chol = cholesky(Hermitian(invΣ))
    return mv_gaussian(μ,Σ,invΣ,logdet(Σ),chol.U)
end

function create_sufficient_statistics(hyper::niw_hyperparams,posterior::niw_hyperparams,points::AbstractArray{Float32,2}, pts_to_group = 0)
    if size(points,2) == 0
        return niw_sufficient_statistics(size(points,2),zeros(Float32,length(hyper.m)),zeros(Float32,length(hyper.m),length(hyper.m)))
    end
    pts = copy(points)
    points_sum = sum(pts, dims = 2)[:]
    S = pts * pts'
    return niw_sufficient_statistics(size(points,2),points_sum,S)
end

function log_marginal_likelihood(hyper::niw_hyperparams, posterior_hyper::niw_hyperparams, suff_stats::niw_sufficient_statistics)
    D = length(suff_stats.points_sum)
    logpi = log(pi)
    return -suff_stats.N*D*0.5*logpi +
        log_multivariate_gamma(posterior_hyper.ν/2, D)-
        log_multivariate_gamma(hyper.ν/2, D) +
         (hyper.ν/2)*(D*log(hyper.ν)+logdet(hyper.ψ))-
         (posterior_hyper.ν/2)*(D*log(posterior_hyper.ν) + logdet(posterior_hyper.ψ)) +
         (D/2)*(log(hyper.κ/posterior_hyper.κ))
end

function aggregate_suff_stats(suff_l::niw_sufficient_statistics, suff_r::niw_sufficient_statistics)
    return niw_sufficient_statistics(suff_l.N+suff_r.N, suff_l.points_sum + suff_r.points_sum, suff_l.S+suff_r.S)
end


using LinearAlgebra
using JuMP
using Cbc

model = Model(Cbc.Optimizer)
set_optimizer_attribute(model, "cuts","off")
# set_optimizer_attribute(model, "msg_lev", GLPK.GLP_MSG_ON)
# set_optimizer_attribute(model, "maxNodes", 5)
month_num = 8
month = 1:month_num
# paramerters
q = [5000] # product fix cost
p = [100] # cost per product
d = diagm([400, 400, 800, 800, 1200, 1200, 1200, 1200]) # demand of product by month
s_init = [200,0,0,0, 0,0,0,0] # initial stock
h = [5] # stock cost per product month

@variable(model, x[i in month, t in month] >= 0)
@variable(model, s[i in month, t in month] >= 0)
@variable(model, y[i in month], Bin)

# 库存平衡
@constraint(model ,dem_sat1[t in month], s_init[t] + x[1,t] == d[1, t] + s[1, t])
@constraint(model, dem_sat[i in month[2:end], t in month], s[i - 1, t] + x[i, t] == d[i, t] + s[i, t] )

# 最大产量约束
@constraint(model, vub1[i in month], x[i, 1] <= (d[1,1] - s_init[1]) * y[i])
@constraint(model, vub[i in month, t in month[2:end]], x[i, t] <= d[t,t] * y[i])
# 目标函数
@expression(model ,cost, sum(p[1] * x[i,t] + h[1] * s[i, t] for i in month, t in month[i:end]) + sum(q[1] * y[i] for i in month))
@objective(model, Min, cost)
for t in month
    fix(s[t,t], 0;force=true)
end
optimize!(model)

# 显示结果
cost_val = value.(cost)
x_val = value.(x)
s_val = value.(s)
@show cost_val;
@show x_val;
@show s_val;


# K. Roosa and G. Chowell.
# Assessing parameter identifiability in compartmental dynamic models using a computational approach: application to infectious disease transmission models
# https://doi.org/10.1186/s12976-018-0097-6
# Model 2
using Logging

using StructuralIdentifiability

logger = Logging.SimpleLogger(stdout, Logging.Debug)
global_logger(logger)

ode = @ODEmodel(
    S'(t) = -b * S(t) * (I(t) + J(t) + q * A(t)) * Ninv,
    E'(t) = b * S(t) * (I(t) + J(t) + q * A(t)) * Ninv - k * E(t),
    A'(t) = k * (1 - r) * E(t) - g1 * A(t),
    I'(t) = k * r * E(t) - (alpha + g1) * I(t),
    J'(t) = alpha * I(t) - g2 * J(t),
    C'(t) = alpha * I(t),
    y(t) = C(t),
    y2(t) = Ninv
)

@time println(assess_global_identifiability(ode))


# Required packages
using DataFrames, StatsBase, Printf, LinearAlgebra

# Reading in source files
include("dat.jl");
include("branchAndBound.jl");
include("utilities.jl");
include("multiComponents.jl");

# Pitprops example
k=5; # desired sparsity
prob = pitprops; #problem selection
obj, xVal, timetoBound, timetoConverge, timeOut,
	explored, toPrint = branchAndBound(prob, k) ;

print(@sprintf("objective value: %.4e\n", obj));
print(@sprintf("time to best solution: %.4e\n", timetoBound));
print(@sprintf("time to converge: %.4e\n", timetoConverge));
print(@sprintf("nodes explored: %d\n", explored));
println("solution vector:");
print(xVal);
println("\n");

# A synthetic example
dataset = randn(20,100);
Sigma = dataset'*dataset/20;
syntheticprob = problem(dataset, Sigma);

k=5; # desired sparsity
prob = syntheticprob; #problem selection
obj, xVal, timetoBound, timetoConverge, timeOut,
	explored, toPrint = branchAndBound(prob, k, outputFlag=0) ;

print(@sprintf("objective value: %.4e\n", obj));
print(@sprintf("time to best solution: %.4e\n", timetoBound));
print(@sprintf("time to converge: %.4e\n", timetoConverge));
print(@sprintf("nodes explored: %d\n", explored));
println("solution vector:");
print(xVal);
println("\n");

# Multiple Components
k=5; # desired sparsity
prob = pitprops; #problem selection
ncomp = 4;
all_x, timeOut = multiOptimalSPCA(prob, k, ncomp);

total_variance = adjVarExplained(prob.Sigma, all_x)

println("adjusted variance explained:");
print(total_variance);
print(@sprintf("total time: %.4e\n", timeOut));
println("solution vectors:");
print(all_x);


"""
```julia
    NegBinomialLikelihood(r::Int=10)
```

[Negative Binomial likelihood](https://en.wikipedia.org/wiki/Negative_binomial_distribution) with number of failures `r`
```math
    p(y|r,f) = binomial(y+r-1,y) (1-σ(f))ʳσ(f)ʸ
```
Where `σ` is the logistic function

"""
struct NegBinomialLikelihood{T<:Real} <: EventLikelihood{T}
    r::Int
    c::Vector{T}
    θ::Vector{T}
    function NegBinomialLikelihood{T}(r::Int) where {T<:Real}
        new{T}(r)
    end
    function NegBinomialLikelihood{T}(r,c,θ) where {T<:Real}
        new{T}(r,c,θ)
    end
end

function NegBinomialLikelihood(r::Int=10) where {T<:Real}
    NegBinomialLikelihood{Float64}(r)
end

function init_likelihood(likelihood::NegBinomialLikelihood{T},inference::Inference{T},nLatent::Integer,nSamplesUsed::Int,nFeatures::Int) where T
    NegBinomialLikelihood{T}(
    likelihood.r,
    rand(T,nSamplesUsed),
    zeros(T,nSamplesUsed))
end

function pdf(l::NegBinomialLikelihood,y::Real,f::Real)
    pdf(NegativeBinomial(lr,get_p(l,f)),y)
end

function expec_count(l::NegBinomialLikelihood,f)
    broadcast((p,r)->p*r./(1.0.-p) ,get_p.(l,f),l.r)
end

function get_p(::NegBinomialLikelihood,f)
    logistic.(f)
end

function Base.show(io::IO,model::NegBinomialLikelihood{T}) where T
    print(io,"Negative Binomial Likelihood")
end

function compute_proba(l::NegBinomialLikelihood{T},μ::Vector{T},σ²::Vector{T}) where {T<:Real}
    N = length(μ)
    pred = zeros(T,N)
    for i in 1:N
        x = pred_nodes.*sqrt(max(σ²[i],zero(T))).+μ[i]
        pred[i] = dot(pred_weights,get_p.(l,x))
    end
    return pred
end

## Local Updates ##

function local_updates!(l::NegBinomialLikelihood{T},y::AbstractVector,μ::AbstractVector,Σ::AbstractVector) where {T}
    l.c .= sqrt.(abs2.(μ) + Σ)
    l.θ .= (l.r.+y)./l.c.*tanh.(0.5*l.c)
end

## Global Updates ##

@inline ∇E_μ(l::NegBinomialLikelihood{T},::AOptimizer,y::AbstractVector) where {T} = (0.5*(y.-l.r),)
@inline ∇E_Σ(l::NegBinomialLikelihood{T},::AOptimizer,y::AbstractVector) where {T} = (0.5.*l.θ,)

## ELBO Section ##

AugmentedKL(l::NegBinomialLikelihood{T},y::AbstractVector) where {T} = PolyaGammaKL(l,y)

function logabsbinomial(n,k)
    log(binomial(n,k))
end

function expec_log_likelihood(l::NegBinomialLikelihood{T},i::AnalyticVI,y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    tot = sum(logabsbinomial.(y.+(l.r-1),y))-log(2.0)*sum((y.+l.r))
    tot += 0.5*dot(μ,(y.-l.r))-0.5*dot(l.θ,μ)-0.5*dot(l.θ,diag_cov)
    return tot
end

function PolyaGammaKL(l::NegBinomialLikelihood,y::AbstractVector) where {T}
    PolyaGammaKL(y.+l.r,l.c,l.θ)
end


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    cubes, permutations = Dict{String, Int}(), Dict{String, Int}()
    for i ∈ 1:10 ^ n
        permutation = join(sort(collect(string(i ^ 3))))
        if haskey(cubes, permutation)
            permutations[permutation] += 1
            if permutations[permutation] == n
                return cubes[permutation]
            end
        else
            cubes[permutation] = i ^ 3
            permutations[permutation] = 1
        end
    end
end

compute(3)

compute(5)

@benchmark compute(5)


using LookupTables
using Test

@testset "LookupTables.jl" begin
    nx = 100
    ny = 50

    y = range(13.7, step = 0.5, length = 50)
    x = range(0.0, step = 0.1, length = 100)

    vals = zeros(nx, ny)

    for i in eachindex(x)
        for j in eachindex(y)
            vals[i, j] = x[i]^2 - 3 * x[i] * y[j]
        end
    end

    ltable = LookupTable2D(x, y, vals)

    using BenchmarkTools

    @btime lookup(ltable, 4.4, 22.2)

    @btime [[lookup(ltable, xval, yval) for xval in x] for yval in y]

end


export ultraconversion!,ultraint!

## Start of support for UFun

# diff from T -> U
function ultradiff{T<:Number}(v::Vector{T})
    #polynomial is p(x) = sum ( v[i] * x^(i-1) )
    if length(v)≤1
        w = zeros(T,1)
    else
        w = Array{T}(length(v)-1)
        for k=1:length(v)-1
            @inbounds w[k] = k*v[k+1]
        end
    end

    w
end

#int from U ->T

#TODO: what about missing truncation?
function ultraint!{T<:Number}(v::Array{T,2})
    for k=size(v,1):-1:2
        for j=1:size(v,2)
            @inbounds v[k,j] = v[k-1,j]/(k-1)
        end
    end

    @simd for j=1:size(v)[2]
        @inbounds v[1,j] = zero(T)
    end

    v
end

function ultraint!{T<:Number}(v::Vector{T})
    resize!(v,length(v)+1)
    @simd for k=length(v):-1:2
        @inbounds v[k] = v[k-1]/(k-1)
    end

    @inbounds v[1] = zero(T)

    v
end

# Convert from U -> T
function ultraiconversion{T<:Number}(v::Vector{T})
    n = length(v)
    w = Array(T,n)

    if n == 1
        w[1] = v[1]
    elseif n == 2
        w[1] = v[1]
        w[2] = 2v[2]
    elseif n ≥ 3
        @inbounds w[end] = 2v[end]
        @inbounds w[end-1] = 2v[end-1]

        for k = n-2:-1:2
            @inbounds w[k] = 2*(v[k] + .5w[k+2])
        end

        @inbounds w[1] = v[1] + .5w[3]
    end

    w
end

# Convert T -> U
function ultraconversion{T<:Number}(v::Vector{T})
    n = length(v)
    w = Array(T,n)

    if n == 1
        w[1] = v[1]
    elseif n == 2
        w[1] = v[1]
        w[2] = .5v[2]
    elseif n ≥ 3
        w[1] = v[1] - .5v[3]

        @simd for j=2:n-2
            @inbounds w[j] = .5*(v[j] - v[j+2])
        end

        w[n-1] = .5v[n-1]
        w[n] = .5v[n]
    end

    w
end

function ultraconversion!{T<:Number}(v::Vector{T})
    n = length(v) #number of coefficients

    if n ≤ 1
        #do nothing
    elseif n == 2
        @inbounds v[2] /= 2
    else
        @inbounds v[1] -= v[3]/2

        for j=2:n-2
            @inbounds v[j] = (v[j] - v[j+2])/2
        end
        @inbounds v[n-1] /= 2
        @inbounds v[n] /= 2
    end

    v
end

function ultraconversion!{T<:Number}(v::Array{T,2})
    n = size(v)[1] #number of coefficients
    m = size(v)[2] #number of funs

    if n ≤ 1
        #do nothing
    elseif n == 2
        @simd for k=1:m
            @inbounds v[2,k] /= 2
        end
    else
        for k=1:m
            @inbounds v[1,k] -= v[3,k]/2

            for j=2:n-2
                @inbounds v[j,k] = (v[j,k] - v[j+2,k])/2
            end
            @inbounds v[n-1,k] /= 2
            v[n,k] /= 2
        end
    end

    v
end

#ultraiconversion and ultraconversion are linear, so it is possible to define them on Complex numbers as so
#ultraiconversion(v::Vector{Complex{Float64}})=ultraiconversion(real(v)) + ultraiconversion(imag(v))*1.0im
#ultraconversion(v::Vector{Complex{Float64}})=ultraconversion(real(v)) + ultraconversion(imag(v))*1.0im

#using DualNumbers
#ultraconversion{T<:Number}(v::Vector{Dual{T}})=dual(ultraconversion(real(v)), ultraconversion(epsilon(v)) )


module Example2018

export Point
struct Point{T}
    x::T
    y::T
end

function Base. +(p::Point, q::Point)
    Point(p.x + q.x, p.y + q.y)
end

export MPoint
mutable struct MPoint{T}
    x::T
    y::T
end

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

export Polynomial
struct Polynomial{T}
    coeffs::Vector{T}
end

export evaluate
function evaluate(p::Polynomial, x)
    sum(c * x^(i-1) for (i,c) in enumerate(p.coeffs))
end

function Base. ==(p::Polynomial, q::Polynomial)
    d = p - q
    # all(c -> c==0, d.coeffs)
    all(d.coeffs .==0)
end

function Base. +(p::Polynomial, q::Polynomial)
    l = max(length(p.coeffs), length(q.coeffs))
    rcoeffs = zeros(eltype(p.coeffs),l)
    for (i,c) in enumerate(p.coeffs)
        rcoeffs[i] += c
    end
    for (i,c) in enumerate(q.coeffs)
        rcoeffs[i] += c
    end
    Polynomial(rcoeffs)
end

function Base. *(a, p::Polynomial)
    Polynomial(map(x->a*x, p.coeffs))
end
Base. *(p::Polynomial, a) = a*p

function Base. -(p::Polynomial)
    (-1) * p
end

function Base. -(p::Polynomial, q::Polynomial)
    p + (-q)
end

export deriv
function deriv(p::Polynomial)
    Polynomial([(i-1) * c for (i,c) in enumerate(p.coeffs)][2:end])
end

end


using Test

ProjDir = joinpath(dirname(@__FILE__), "..", "examples", "BernoulliInitTheta")
cd(ProjDir) do

  isdir("tmp") &&
    rm("tmp", recursive=true);

  bernoullimodel = "
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    theta ~ beta(1,1);
    y ~ bernoulli(theta);
  }
  "

  bernoullidata = [
    Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1]),
    Dict("N" => 10, "y" => [0, 1, 0, 0, 0, 0, 1, 0, 0, 1]),
    Dict("N" => 10, "y" => [0, 0, 0, 0, 0, 0, 1, 0, 1, 1]),
    Dict("N" => 10, "y" => [0, 0, 0, 1, 0, 0, 0, 1, 0, 1])
  ]

  theta_init = [0.1, 0.4, 0.5, 0.9]
  bernoulliinit = Dict[
    Dict("theta" => theta_init[1]),
    Dict("theta" => theta_init[2]),
    Dict("theta" => theta_init[3]),
    Dict("theta" => theta_init[4]),
  ]

  monitor = ["theta", "lp__", "accept_stat__"]

  stanmodel = Stanmodel(name="bernoulli",
    model=bernoullimodel,
    init=Stan.Init(init=bernoulliinit),
    adapt=1);

  rc, a3d, cnames = stan(stanmodel, bernoullidata, CmdStanDir=CMDSTAN_HOME)

  if rc == 0
    sdf  = read_summary(stanmodel)
    @test sdf[sdf.parameters .== :theta, :mean][1] ≈ 0.34 rtol=0.1
  end

  isdir("tmp") &&
    rm("tmp", recursive=true);

end # cd


module NtToolBox

# package code goes here

using LinearAlgebra

export load_image, imageplot, rescale, clamP, snr, perform_wavelet_transf, plot_wavelet, compute_wavelet_filter, perform_wavortho_transf, plot_levelset, gaussian_blur, bilinear_interpolate, Grad, Div, perform_haar_transf, load_sound, perform_stft, plot_spectrogram, compute_max, perform_blurring, plot_vf, meshgrid, read_mesh, compute_boundary, plot_mesh, compute_normal, perform_linprog, plot_hufftree, perform_conjugate_gradient, perform_dijkstra_fm, cconvol, subsampling, upsampling, trim_dim, rand_discr, plot_multiclasses

## m must be of type Array{Float32, 3} or Array{Float64, 3}, it makes products between the vectors of the 3rd dimension and v

function Mdot(m, v::Array{Float64, 1})
  w = ones(size(m, 1), size(m, 2))
  for i in 1:size(m, 1)
    for j in 1:size(m, 2)
      w[i, j] = (m[i, j, :]'*v)[1]
    end
  end
  return w
end

##

## Rescale linearly the dynamic of a vector to fit within a range [a,b]

function rescale(f, a = 0, b = 1)
  v = maximum(f) - minimum(f)
  g = copy(f - minimum(f))
  if v > 0
    g = g / v
  end
  return a + g*(b - a)
end

##

## Function clamp

function clamP(x, a = [], b = [])

  """
   clamp - clamp a value

     y = clamp(x,a,b);

   Default is [a,b]=[0,1].

     Copyright (c) 2004 Gabriel Peyre
  """

  if a == []
    a = 0.0
  end
  if b == []
    b = 1.0
  end
  return min(max(x, a), b)
end

#function plot_levelset(Z, level = 0, f = [])
#    """
#        f is supposed to be of the same shape as Z
#    """
#     if length(f) == 0
#         f = copy(Z)
#     end
#
#     (n, p) = size(Z)
#     (X, Y) = meshgrid(collect(0 : n - 1), collect(0 : p - 1))
#     contour(X, Y, Z, [level], linewidths = 2, colors = "red")
#     imageplot(f)
# end

function gaussian_blur(f, sigma)

      """ gaussian_blur - gaussian blurs an image
      %
      %   M = perform_blurring(M, sigma, options);
      %
      %   M is the original data
      %   sigma is the std of the Gaussian blur (in pixels)
      %
      %   Copyright (c) 2007 Gabriel Peyre
      """
      if sigma <= 0
          return
      end
      n = maximum(size(f))
      t = [collect(0:n/2); collect(-n/2:-2)]
      (Y, X) = meshgrid(t, t)
      h = exp( -(X.^2 + Y.^2)./(2.0*float(sigma)^2) )
      h = h./sum(h)
      return real(plan_ifft((plan_fft(f)*f).*(plan_fft(h)*h))*((plan_fft(f)*f).*(plan_fft(h)*h)))
end

function plot_levelset(Z, level = 0, f = [])
    """
        f is supposed to be of the same shape as Z
    """
    if length(f) == 0
        f = copy(Z)
    end

    (n, p) = size(Z)
    (X, Y) = meshgrid(collect(0 : n - 1), collect(0 : p - 1))
    contour(X, Y, Z, [level], linewidths = 2, colors="red")
    imageplot(f)
end

function compute_max(X,d)
    """ compute_max - compute maximum along dimension d

       Y,I = compute_max(X,d);

       Copyright (c) 2008 Gabriel Peyre
    """
    if ndims(X)>=2
        Y = maximum(X,d)
        I = mapslices(argmax,X,d)[:]
    elseif ndims(X)==1
        Y = maximum(X)
        I = indmaw(X)
    end
    return Y,I
end

function compute_max(X)
    if ndims(X)>=2
        Y = maximum(X,1)
        I = mapslices(argmax,X,1)[:]
    elseif ndims(X)==1
        Y = maximum(X)
        I = indmaw(X)
    end
    return Y,I
end

# function circshift(x, p)
#     """
#         Circular shift of an array.
#     """
#     y = copy(x)
#     y = cat(1, y[p[1]+1:end, :], y[1:p[1]+1, :])
#     if (size(x)[2] > 0) & (length(p) > 1)
#         y = cat(2, y[:, p[1]+1:end], y[:, 1:p[1]+1])
#     end
#     return y
# end

include("ndgrid.jl")
include("signal.jl")
include("perform_wavelet_transf.jl")
include("compute_wavelet_filter.jl")
include("Grad.jl")
include("Div.jl")
include("perform_blurring.jl")
include("read_bin.jl")
#include("isosurface.jl")
include("perform_thresholding.jl")
include("load_sound.jl")
include("perform_stft.jl")
include("plot_spectrogram.jl")
include("plot_vf.jl")
include("read_mesh.jl")
include("compute_boundary.jl")
include("plot_mesh.jl")
include("compute_normal.jl")
include("perform_linprog.jl")
include("plot_hufftree.jl")
include("perform_conjugate_gradient.jl")
include("graph.jl")
include("perform_haar_transf.jl")
#include("perform_fast_marching.jl")
include("rand_discr.jl")
include("plot_multiclasses.jl")

end # module


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 73112674-61f9-4915-9b09-216c70df0e74
begin
	using MortalityTables

	vbt = MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB") #or any other table
end

# ╔═╡ ce999823-acab-47bb-bb15-babb8018a06a
using DataFrames

# ╔═╡ a5611f56-cca3-4a06-842e-715a89d401e4
begin
	# add a table of contents to the page
	using PlutoUI
	TableOfContents()
end

# ╔═╡ 059c84bc-5d87-4781-b0f2-15563bfae09d
md"# Using MortaltiyTables.jl with DataFrames

## MortalityTables.jl standard representation

MortalityTables.jl stores the rates in a very efficient manner as a collection of vectors indexed by attained age.
"

# ╔═╡ 6e38cf1f-0c8e-47b9-baa1-e9525f0a1faa
md"First, we include the package, and then we'll pick a table, where all of the `mort.soa.org` tables are mirrored into your MortalityTables.jl installation."

# ╔═╡ c9770fc4-e460-46d9-ba6b-2ab0fc17830e
md" To see how the data is represented, we can look at the the select data for a 55 year old and see the attained age and mortality rates:"

# ╔═╡ 6bce5873-d064-46ea-9969-7809051300db
vbt.select[55]

# ╔═╡ 5f0e15fc-1107-11ec-2a58-3351c91877fc
md" This is very efficient and convienent for modeling, but a lot of times you want the data matched up with policy data in a DataFrame."

# ╔═╡ f0b7e725-b433-458a-a5bf-efd5dbc990ce
md"## Getting data into a dataframe"

# ╔═╡ b5a00124-c994-4b79-99dc-6b966bec88cd
md"### Generate sample data"

# ╔═╡ 2d497d32-b088-4027-ad5b-0c4597db3922
sample_size = 10_000

# ╔═╡ dbb95d40-6517-491b-b6a9-abccbd2a3a78
sample_data = let
	# generate fake data
	df = DataFrame(
		"sex" => rand(["Male","Female"],sample_size),
		"smoke" => rand(["Smoker","Nonsmoker"],sample_size),
		"issue_age" => rand(25:65,sample_size),
		)

	# a random offset of issue age is the current attained age
	df.attained_age = df.issue_age .+ rand(1:10,sample_size)
	df
end

# ╔═╡ e0acb2d4-668e-40aa-ac70-de4527c7a891
md"### Define the table set you want to use"

# ╔═╡ 423f35bd-6b79-464d-a8f6-a2405e5ba3ce
md"
There are a lot of different possible combinations of parameters that you might want to use, such as rates that vary by sex, risk class, table set (VBT/CSO/etc), smoking status, relative risk, ALB/ANB, etc.

It's easy to define the parameters applicable to your assumption set. Here, we'll use a dictionary to define the relationship:
"

# ╔═╡ c351704b-4bf8-4fc0-bda2-c2b1970c2a11
rate_map = Dict(
	"Male" => Dict(
		"Smoker" => MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Smoker, ANB"),
		"Nonsmoker" => MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB"),
		),

	"Female" => Dict(
		"Smoker" => MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Female Smoker, ANB"),
		"Nonsmoker" => MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Female Nonsmoker, ANB"),
		)
	)

# ╔═╡ 4d88a451-f350-4a72-ae5b-8698b0527511
md" and then we'll define a function to look up the relevant rate. Note how the function matches the levels we defined for the assumption set dictionary above."

# ╔═╡ 45f4378e-9299-429d-b361-8f99b2be6efd
function rate_lookup(assumption_map,sex,smoke,issue_age,attained_age)
	# pick the relevant table
	table = assumption_map[sex][smoke]

	# check if the select rate exists, otherwise look to the ulitmate table
	if issue_age in eachindex(table.select)
		table.select[issue_age][attained_age]
	else
		table.ultimate[attained_age]
	end
end

# ╔═╡ 6d5b5548-7cbc-435e-8058-ebd529f839f1
md"### Lining up with dataframe

By mapping each row's data to the lookup function, we get a vector of rates for our data:"

# ╔═╡ 372694fa-60a8-4305-a922-f020a4deb896
rates = map(eachrow(sample_data)) do row
	rate_lookup(rate_map,row.sex,row.smoke,row.issue_age,row.attained_age)
end

# ╔═╡ a8cfd44f-5e40-4624-b98e-00de4a623af2
md"And finally, we can just add this to the dataframe:"

# ╔═╡ 7fe85397-51a5-4d82-86ee-f6881dcf8808
sample_data.expectation = rates;

# ╔═╡ 47da227d-72ab-4eae-bb27-3a61d842bb03
sample_data

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
MortalityTables = "4780e19d-04b9-53dc-86c2-9e9aa59b5a12"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
DataFrames = "~1.2.2"
MortalityTables = "~2.1.4"
PlutoUI = "~0.7.9"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "84918055d15b3114ede17ac6a7182f68870c16f7"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "3.3.1"

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "6071cb87be6a444ac75fdbf51b8e7273808ce62f"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.35.0"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "bec2532f8adb82005476c141ec23e921fc20971b"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.8.0"

[[DataFrames]]
deps = ["Compat", "DataAPI", "Future", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrettyTables", "Printf", "REPL", "Reexport", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"]
git-tree-sha1 = "d785f42445b63fc86caa08bb9a9351008be9b765"
uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
version = "1.2.2"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[EzXML]]
deps = ["Printf", "XML2_jll"]
git-tree-sha1 = "0fa3b52a04a4e210aeb1626def9c90df3ae65268"
uuid = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
version = "1.1.0"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Future]]
deps = ["Random"]
uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[InvertedIndices]]
git-tree-sha1 = "bee5f1ef5bf65df56bdd2e40447590b272a5471f"
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
version = "1.1.0"

[[IterTools]]
git-tree-sha1 = "05110a2ab1fc5f932622ffea2a003221f4782c18"
uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
version = "1.3.0"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[Libiconv_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "42b62845d70a619f063a7da093d995ec8e15e778"
uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531"
version = "1.16.1+1"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "2ca267b08821e86c5ef4376cffed98a46c2cb205"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.1"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MortalityTables]]
deps = ["OffsetArrays", "Parsers", "Pkg", "QuadGK", "Requires", "UnPack", "XMLDict"]
git-tree-sha1 = "cae9281196e40e24cb5885b8463a6f955ae0f3eb"
uuid = "4780e19d-04b9-53dc-86c2-9e9aa59b5a12"
version = "2.1.4"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[OffsetArrays]]
deps = ["Adapt"]
git-tree-sha1 = "c870a0d713b51e4b49be6432eff0e26a4325afee"
uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
version = "1.10.6"

[[OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "438d35d2d95ae2c5e8780b330592b6de8494e779"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.3"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoUI]]
deps = ["Base64", "Dates", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "Suppressor"]
git-tree-sha1 = "44e225d5837e2a2345e69a1d1e01ac2443ff9fcb"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.9"

[[PooledArrays]]
deps = ["DataAPI", "Future"]
git-tree-sha1 = "a193d6ad9c45ada72c14b731a318bedd3c2f00cf"
uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
version = "1.3.0"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "0d1245a357cc61c8cd61934c07447aa569ff22e6"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.1.0"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "12fbe86da16df6679be7521dfb39fbc861e1dc7b"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.4.1"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[Suppressor]]
git-tree-sha1 = "a819d77f31f83e5792a76081eee1ea6342ab8787"
uuid = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
version = "0.2.0"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "368d04a820fe069f9080ff1b432147a6203c3c89"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.5.1"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[UnPack]]
git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
version = "1.0.2"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[XML2_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Pkg", "Zlib_jll"]
git-tree-sha1 = "1acf5bdf07aa0907e0a37d3718bb88d4b687b74a"
uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a"
version = "2.9.12+0"

[[XMLDict]]
deps = ["EzXML", "IterTools", "OrderedCollections"]
git-tree-sha1 = "d9a3faf078210e477b291c79117676fca54da9dd"
uuid = "228000da-037f-5747-90a9-8195ccbf91a5"
version = "0.4.1"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─059c84bc-5d87-4781-b0f2-15563bfae09d
# ╟─6e38cf1f-0c8e-47b9-baa1-e9525f0a1faa
# ╠═73112674-61f9-4915-9b09-216c70df0e74
# ╟─c9770fc4-e460-46d9-ba6b-2ab0fc17830e
# ╠═6bce5873-d064-46ea-9969-7809051300db
# ╟─5f0e15fc-1107-11ec-2a58-3351c91877fc
# ╠═f0b7e725-b433-458a-a5bf-efd5dbc990ce
# ╟─b5a00124-c994-4b79-99dc-6b966bec88cd
# ╠═ce999823-acab-47bb-bb15-babb8018a06a
# ╠═2d497d32-b088-4027-ad5b-0c4597db3922
# ╠═dbb95d40-6517-491b-b6a9-abccbd2a3a78
# ╟─e0acb2d4-668e-40aa-ac70-de4527c7a891
# ╟─423f35bd-6b79-464d-a8f6-a2405e5ba3ce
# ╠═c351704b-4bf8-4fc0-bda2-c2b1970c2a11
# ╟─4d88a451-f350-4a72-ae5b-8698b0527511
# ╠═45f4378e-9299-429d-b361-8f99b2be6efd
# ╟─6d5b5548-7cbc-435e-8058-ebd529f839f1
# ╠═372694fa-60a8-4305-a922-f020a4deb896
# ╟─a8cfd44f-5e40-4624-b98e-00de4a623af2
# ╠═7fe85397-51a5-4d82-86ee-f6881dcf8808
# ╠═47da227d-72ab-4eae-bb27-3a61d842bb03
# ╟─a5611f56-cca3-4a06-842e-715a89d401e4
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


#------------------------------------------------------------------------------

# rational number reconstruction implementation borrowed from CLUE
# and modified a bit to suit the 'Modern Computer Algebra' definitions
# returns a rational r // h of QQ field in a canonical form such that
#   r // h ≡ a (mod m)
#
# let n = max( λ(a), λ(m) ) , where λ(x) is a number of bits for x
# O(n^2)
"""
    rational_reconstruction(a, m)

Rational number reconstruction implementation borrowed from CLUE
and modified a bit to suit the 'Modern Computer Algebra' definitions.
Returns a rational of QQ field in a canonical form that
is congruent a modulo m

a, m are integers
"""
function rational_reconstruction(a::I, m::I) where {I<:Union{Int, BigInt}}
    a = mod(a, m)
    if a == 0 || m == 0
        return Nemo.QQ(0, 1)
    end
    if m < 0
        m = -m
    end
    if a < 0
        a = m - a
    end
    if a == 1
        return Nemo.QQ(1, 1)
    end
    bnd = sqrt(float(m) / 2)

    U = I[1, 0, m]
    V = I[0, 1, a]
    while abs(V[3]) >= bnd
        q = div(U[3], V[3])
        T = U .- q .* V
        U = V
        V = T
    end

    t = abs(V[2])
    r = V[3] * sign(V[2])
    # changed from `<= bnd` to `<= m / bnd`
    if t <= m / bnd && gcd(r, t) == 1
        return Nemo.QQ(r, t)
    end

    throw(DomainError(
        :($a//$m), "rational reconstruction of $a (mod $m) does not exist"
    ))
end

#------------------------------------------------------------------------------

function modular_reduction(x::FracElem, field)
    n, d = field(numerator(x)), field(denominator(x))
    if iszero(d)
        throw(DomainError(
            :($x), "modular reduction of $x (to $field) does not exist"
        ))
    end
    n // d
end

function modular_reduction(x::gfp_fmpz_elem, field)
    return field(convert(BigInt, x))
end

function modular_reduction(x::gfp_elem, field)
    return field(x)
end

#------------------------------------------------------------------------------

# Returns the minimal polynomial of the sequence `h` modulo `m`,
# the sequence `h` is formed as the given polynomial coefficients
# If n is max(deg(h), deg(m))
# O(n^2)
function minimal_polynomial(h::PolyElem, m::PolyElem)
    # The function implements Algorithm 12.9 from
    #   "Modern Computer Algebra", second edition,
    #   Joachim von zur Gathen, Jürgen Gerhard

    n = div(degree(m), 2)
    FF = parent(h)

    U = (FF(1), FF(0), m)
    V = (FF(0), FF(1), h)
    while degree(V[3]) >= n
        q = div(U[3], V[3])
        T = U .- q .* V
        U = V
        V = T
    end

    s, t = V[3], V[2]

    if iszero(coeff(t, 0))
        throw(SingularException(0))
    end

    t *= inv(coeff(t, 0))
    d = max(degree(s) + 1, degree(t))

    if gcd(s, t) == 1 && degree(t) <= n
        return reverse(t, d + 1)
    end

    throw(DomainError(
        h, "minimal polynomial of $h (mod $m) does not exist"
    ))
end

#------------------------------------------------------------------------------

# Returns such vector, that it is orthogonal to the given vector set
# The given `vectors` set must be linearly independent
# and for any vect from `vectors` the following must hold
#   length(vectors) = dim(vect) - 1
#
# vectors :: dict of elems (row_idx, row)
#
# Note that vectors may be modified, CAREFUL ABOUT CONCURRENCY!
# It is guaranteed that after the function yields,
# `vectors` will be in the original state
#
# NOT GUARANTEED FOR RETURNED VECTOR TO BE ORTHOGONAL
function find_orthogonal!(vectors::AbstractDict)
    first_vector = first(values(vectors))
    field = base_ring(first_vector)
    n = dim(first_vector)

    i = 1
    found = false
    nnz_rows = Array(1 : n)
    f = zero(first_vector)

    # some complexity estimations should be done
    while ! found
        # taking the i-th unit vector
        u = unit_sparsik(n, i, field)

        # the only change of `vectors`
        vectors[n] = u

        b = unit_sparsik(n, n, field)

        B = from_rows(n, n, field, nnz_rows, vectors)

        try
            # if the selected `u` is not independent with `vectors`
            # then the wiedemann must fail and we should choose a new vector
            # Gleb: we could use probabilistic here
            # Alex: we could, so, the answer could be "almost orthogonal"
            # (only orthogonal to some of the vectors)
            f = square_nonsingular_randomized_wiedemann(B, b)
            found = true
        catch e
            if isa(e, SingularException)
                @info e
            else
                rethrow()
            end
        end

        i += 1
        i % 10 == 0 && error("something is wrong")
    end

    # undo the change
    delete!(vectors, n)

    return f
end

# A slower analogue of the upper function
#
# vectors :: array of sparsiks
function find_orthogonal!(vectors::AbstractVector)
    return find_orthogonal!(
        Dict(i => vect for (i, vect) in enumerate(vectors))
    )
end

#------------------------------------------------------------------------------

#=
function construct_jacobians(system)
    nnz = [
        (i, var_index(v), derivative(f, v))
        for (i, f) in enumerate(system)
            for v in vars(f)
    ]
    m, n = length(system), length(degrees(first(system)))
    ring = parent(first(system))
    return from_COO(m, n, nnz, ring)
end
=#

# for the given system of polynomials in variables xi
# consturucts a set of matrices Ai over number field
# such that the Jacobian J of the provided system can represented as the sum
# J = Aᵢxⁱ
"""
    construct_jacobians(system)

For the given system of polynomials in variables xi
consturucts a set of matrices Aᵢ over number field
such that the Jacobian J of the provided system can represented as the sum
J = Aᵢxⁱ
"""
function construct_jacobians(system::AbstractArray{T}) where {T<:MPolyElem}
    domain = base_ring(first(system))
    poly_ring = parent(first(system))
    gen2idx = Dict(x=>i for (i, x) in enumerate(gens(poly_ring)))

    jacobians = Dict()

    for (p_idx, poly) in enumerate(system)
        for var in vars(poly)
            v_idx = gen2idx[var]
            # term is of form α*monomial
            for term in terms(derivative(poly, var))
                monom = monomial(term, 1)
                cf = coeff(term, 1)

                # how can we do it in a normal way..
                !haskey(jacobians, monom) && (jacobians[monom] = Dict())
                idx = (v_idx, p_idx)
                if !haskey(jacobians[monom], idx)
                    jacobians[monom][idx] = zero(domain)
                end
                jacobians[monom][idx] += cf
            end
        end
    end

    m, n = length(system), length(gens(poly_ring))
    factors = [
        from_COO(m, n, jac, domain)
        for jac in values(jacobians)
    ]

    @info "constructed a set of $(length(factors)) matrices $m×$n from the system Jacobian"

    return factors
end

function construct_jacobians(system)
    ring = parent(first(keys(system)))
    construct_jacobians([system[x] for x in gens(ring)])
end

#------------------------------------------------------------------------------

# vectors - an array of 1D iterables
# domain - an MPoly domain
#
# converts the given vectors into the polynomial form
# with respect to the generators of the given domain
function polynormalize(vectors, domain) where {T}
    polynômes = []
    vars = Nemo.gens(domain)

    for v in vectors
        push!(
            polynômes,
            sum(map(prod, zip(vars, to_dense(v))))
        )
    end

    polynômes
end

#------------------------------------------------------------------------------

# for As[i] in array As compute dok_sparsik Tr: Tr_ij = trace(A_i, A_j)
function gram_matrix(As::Array{AbstractSparseObject{FlintRationalField},1})
    F = base_ring(As[1])
    n = length(As)

    traces = Dict{Tuple{Int, Int}, elem_type(F)}(
        (i, j) => tr(As[i] * As[j])
        for i in 1 : n
            for j in i : n
    )

    nnz_coords = [
        (i, j, traces[min(i, j), max(i, j)])
        for i in 1 : n
            for j in 1 : n
                if ! iszero(traces[min(i, j), max(i, j)])
    ]

    return from_COO(n, n, nnz_coords, F)
end


# clark_west.jl
# Clark-West
#
# Copyright (C) 2020   Guilherme Bodin, Paul Söderlind
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export ClarkWestTest

struct ClarkWestTest <: ZTest
    n::Int       # number of observations
    xbar::Real   # estimated mean
    stderr::Real # population standard error
    z::Real      # t-statistic
    μ0::Real     # mean under h_0
end

"""
    ClarkWestTest(e1::AbstractVector{<:Real}, e2::AbstractVector{<:Real};lookahead::Integer=1)

Perform the Clark-West test of equal performance of two nested predition models, in terms of the
out-of-sample mean squared prediction errors.

`e1` is a vector of forecasts from the smaller (nested) model, `e2` is a vector of forecast
errors from the larger model, and `lookahead` is the number of steps ahead of the forecast.
Typically, the null hypothesis is that the two models perform equally well (a two-sided test),
but sometimes we test whether the larger model performs better, which is indicated by a
positive test statistic, for instance, above 1.645 for the 5% significance level (right tail test).

Implements: [`pvalue`](@ref)
# References
 * Clark, T. E., West, K. D. 2006, Using out-of-sample mean squared prediction errors to test
   the martingale difference hypothesis. Journal of Econometrics, 135(1): 155–186.
 * Clark, T. E., West, K. D. 2007, Approximately normal tests for equal predictive accuracy
   in nested models. Journal of Econometrics, 138(1): 291–311.

"""
function ClarkWestTest(e1::AbstractVector{<:Real}, e2::AbstractVector{<:Real};
                       lookahead::Integer=1)
    length(e1) == length(e2) || throw(DimensionMismatch("inputs must have the same length"))
    n            = length(e1)
    d            = 2*e1.*(e1 - e2)
    cw_cov       = HypothesisTests.autocov(d, collect(0:lookahead-1))
    cw_var       = (cw_cov[1] + 2 * sum(cw_cov[2:end]))/n
    statistic_cw = mean(d)/sqrt(cw_var)
    return ClarkWestTest(n,mean(d),sqrt(cw_var),statistic_cw,0.0)
end

testname(::ClarkWestTest) = "Clark West test"
population_param_of_interest(x::ClarkWestTest) = ("Mean",x.μ0, x.xbar)
default_tail(test::ClarkWestTest) = :both

function show_params(io::IO, x::ClarkWestTest, ident)
    println(io, ident, "number of observations:    $(x.n)")
    println(io, ident, "CW statistic:              $(x.z)")
    println(io, ident, "population standard error: $(x.stderr)")
end


function test_standardize()
    mat1 = [0. 0.; 0. 0.; 1. 1.; 1. 1.]
    mat2 = [-1. -1.; -1. -1.; 1. 1.; 1. 1.]
    @test isapprox(MVApp.standardize(mat1), mat2)
    @test isapprox(MVApp.reverse_standardization(mat2,mat1),mat1)
    mat3 = ones(3,2)
    mat4 = zeros(3,2)
    @test isapprox(MVApp.standardize(mat3), mat4)
    @test isapprox(MVApp.reverse_standardization(mat4,mat3),mat3)
end

function test_gradient()
    vec = [1., 2., 4., 7., 11., 16.]
    res = [1. , 1.5, 2.5, 3.5, 4.5, 5. ]
    @test isapprox(MVApp.gradient(vec), res)
    mat = [1. 2. 6.; 3. 4. 5.]
    res2 = [2.0  2.0  -1.0;2.0  2.0  -1.0]
    @test isapprox(MVApp.gradient(mat), res2)
end


"""
    OrthogonalPolynomialRatio(P,x)

is a is equivalent to the vector `P[x,:] ./ P[x,2:end]`` but built from the
recurrence coefficients of `P`.
"""
mutable struct OrthogonalPolynomialRatio{T, PP<:AbstractQuasiMatrix{T}} <: AbstractCachedVector{T}
    P::PP # OPs
    x::T
    data::Vector{T}
    datasize::Tuple{Int}

    function OrthogonalPolynomialRatio{T, PP}(P::PP, x::T) where {T,PP<:AbstractQuasiMatrix{T}}
        μ = inv(sqrt(sum(orthogonalityweight(P))))
        new{T, PP}(P, x, [Base.unsafe_getindex(P,x,1)/Base.unsafe_getindex(P,x,2)], (1,))
    end
end

OrthogonalPolynomialRatio(P::AbstractQuasiMatrix{T}, x) where T = OrthogonalPolynomialRatio{T,typeof(P)}(P, convert(T, x))

size(K::OrthogonalPolynomialRatio) = (ℵ₀,)

function cache_filldata!(R::OrthogonalPolynomialRatio, inds)
    A,B,C = recurrencecoefficients(R.P)
    x = R.x
    data = R.data
    @inbounds for n in inds
        data[n] = inv(A[n]*x + B[n] - C[n] * data[n-1])
    end

end


using DifferentialEquations
using Plots
using LinearAlgebra

function DeQ(du,u,p,t)
du[1]=-u[2]*(1+u[1]^2) / ( (u[1]^2+u[2]^2)*(1-u[2]^2));
du[2]=u[1]*(1-u[2]^2)/ ( (u[1]^2+u[2]^2)*(1+u[1]^2));
du[3]= 2*u[1]*u[2]/(u[1]^2+u[2]^2)^2 *u[3]-1/(u[2]^2+u[1]^2)^2* (1-u[2]^2)*(3*u[1]^4+u[2]^2*u[1]^2-u[2]^2+u[1]^2) /( (1+u[1]^2)^2 )*u[4];
du[4]=-(1/(u[1]^2+u[2]^2)^2* (1+u[1]^2)*(3*u[2]^4+u[1]^2*u[2]^2+u[1]^2-u[2]^2)/( (1-u[2]^2)^2) )*u[3]-(2*u[1]*u[2]/(u[1]^2+u[2]^2)^2)*u[4];
du[5]= 2*u[1]*u[2]/(u[1]^2+u[2]^2)^2 *u[5]-1/(u[2]^2+u[1]^2)^2* (1-u[2]^2)*(3*u[1]^4+u[2]^2*u[1]^2-u[2]^2+u[1]^2) /( (1+u[1]^2)^2 )*u[6];
du[6]=-(1/(u[1]^2+u[2]^2)^2* (1+u[1]^2)*(3*u[2]^4+u[1]^2*u[2]^2+u[1]^2-u[2]^2)/( (1-u[2]^2)^2) )*u[5]-(2*u[1]*u[2]/(u[1]^2+u[2]^2)^2)*u[6];
end
condition(u,t,integrator)= u[1]
affect!(integrator) = terminate!(integrator)
cb = ContinuousCallback(condition,affect!)

#
H=range(.05,stop=.2,length=10)
# for k=1:length(H)
H=.125
    h=1/(2*H)
    Y0=1.0/sqrt(h+1)
    tspan0=(0.,1e4)
    u0=[0.,Y0,0.,0.,0.,0.]
    prob = ODEProblem(DeQ,u0,tspan0)
    t,A1=solve(prob, Vern9(),reltol=1e-12,abstol=1e-20,maxiters=1e15,callback=cb)
    FinalT=t.t[end]

    u0=[0.,Y0,1.,0.,0.,1.]
    tspan=(0.,FinalT)
    prob = ODEProblem(DeQ,u0,tspan)
    t,A=solve(prob, Vern9(),reltol=1e-12,abstol=1e-12,maxiters=1e15)
    N=length(t.t)
    r=zeros(N)
    z=zeros(N)
    theta=zeros(N)
    for j=1:N
        #Create Monodromy matrix
        a=A[3,j]
        b=A[4,j]
        c=A[5,j]
        d=A[6,j]
        M=[a c; b d]
        W,S,V = svd(M)
        P=W*Diagonal(S)*W'
        U=W*V'
        r[j]=P[1,1]
        z[j]=P[1,2]
        theta[j]=asin(U[2,1])
    end
    x=r.*cos.(theta)
    y=r.*sin.(theta)

plot3d(x,y,z)


"""
An implementation of simple method to calculate Ewald energy
"""
function calc_ewald( pw::PWGrid, Sf, Xpos, Nspecies::Int, atm2species,
                     Zv::Array{Float64}; sigma=nothing )
    #
    Ω  = pw.Ω
    r  = pw.r
    Ns = pw.Ns
    G2 = pw.gvec.G2
    Npoints = prod(Ns)
    LatVecs = pw.LatVecs
    #
    # Generate array of distances
    dr = gen_dr_center( r, LatVecs )
    #
    # Generate charge density
    #
    if sigma==nothing
      sigma = 0.25*ones(Nspecies)
    end
    #
    g1  = zeros( Float64, Npoints )
    rho_is = zeros( Float64, Npoints, Nspecies )
    Rho = zeros(Float64, Npoints)
    #
    for isp = 1:Nspecies
        c1 = 2*sigma[isp]^2
        cc1 = sqrt(2*pi*sigma[isp]^2)^3
        for ip=1:Npoints
            g1[ip] = exp(-dr[ip]^2/c1)/cc1
        end
        #
        g1 = Zv[isp] * g1[:]
        #
        ctmp = R_to_G( Ns, g1 )
        for ip = 1:Npoints
            ctmp[ip] = ctmp[ip]*Sf[ip,isp]
        end
        rho_is[:,isp] = real( G_to_R(Ns, ctmp) )
        intrho = sum(rho_is[:,isp])*Ω/Npoints
        @printf("Species %d, intrho: %18.10f\n", isp, intrho)
        Rho[:] = Rho[:] + rho_is[:,isp]
    end
    #
    intrho = sum(Rho)*Ω/Npoints
    @printf("Integrated total Gaussian rho: %18.10f\n", intrho)
    #
    # Solve Poisson equation and calculate Hartree energy
    ctmp = 4.0*pi*R_to_G( Ns, Rho )
    ctmp[1] = 0.0
    for ip = 2:Npoints
        ctmp[ip] = ctmp[ip] / G2[ip]
    end
    phi = real( G_to_R( Ns, ctmp ) )
    Ehartree = 0.5*dot( phi, Rho ) * Ω/Npoints
    #
    Eself = 0.0
    Natoms = size(Xpos)[2]
    for ia = 1:Natoms
        isp = atm2species[ia]
        Eself = Eself + Zv[isp]^2/(2*sqrt(pi))*(1.0/sigma[isp])
    end
    E_nn = Ehartree - Eself
    @printf("Ehartree, Eself, E_nn = %18.10f %18.10f %18.10f\n", Ehartree, Eself, E_nn)
    return E_nn
end


isarmstrong(n) = sum(k^length(digits(n)) for k in digits(n)) == n


equations = open("math.txt") do f
    readlines(f)
end

function eval(eq::AbstractString)
    # println(eq)
    paren = r"\((?:[^)(]+|(?R))*+\)"
    math = r"\d+|\+|\*"
    for m in reverse(collect(eachmatch(paren, eq)))
        neq = m.match[2:end-1]
        eq = replace(eq, m.match => eval(neq))
    end
    op = nothing
    acc = 0
    for x in eachmatch(math, eq)
        x = x.match
        if x in ("+", "*")
            op = x
        else
            x = parse(Int, x)
            if op == "+"
                acc += x
            elseif op == "*"
                acc *= x
            elseif op === nothing
                acc = x
            end
        end
    end
    return acc
end

function eval2(eq::AbstractString)
    # println(eq)
    paren = r"\((?:[^)(]+|(?R))*+\)"
    for m in reverse(collect(eachmatch(paren, eq)))
        neq = m.match[2:end-1]
        eq = replace(eq, m.match => eval2(neq))
    end
    while occursin("+", eq)
        m = last(collect(eachmatch(r"(\d+) \+ (\d+)", eq)))
        o1, o2 = parse.(Int, m.captures)
        eq = replace(eq, m.match => o1 + o2)
        # println("$(o1) + $(o2) => $(o1 + o2)")
    end
    # println(eq)
    while occursin("*", eq)
        m = last(collect(eachmatch(r"(\d+) \* (\d+)", eq)))
        o1, o2 = parse.(Int, m.captures)
        eq = replace(eq, m.match => o1 * o2)
        # println("$(o1) * $(o2) => $(o1 * o2)")
    end
    return parse(Int, eq)
end

println(sum(map(eq -> eval(eq), equations)))

println(sum(map(eq -> eval2(eq), equations)))


check() = a & 0xffff == b & 0xffff
# Part 1
a, b = 116, 299
n = 0
update(a, b) = (a * 16807, b * 48271) .% 2147483647
for _ in 1:40000000
    a, b = update(a, b)
    n += check()
end
println(n)

# Part 2
a, b = 116, 299
n = 0
for _ in 1:5000000
    while true
        a = (a * 16807) % 2147483647
        a % 4 == 0 && break
    end
    while true
        b = (b * 48271) % 2147483647
        b % 8 == 0 && break
    end
    n += check()
end
println(n)


# 2520 is the smallest number that can be divided by each of the numbers from 1 to 10 without any remainder.

# What is the smallest positive number that is evenly divisible by all of the numbers from 1 to 20?

primes = [19, 17, 13, 11, 7, 5, 3, 2]
added_factors = [3, 2, 2, 2]

number = prod(primes)*prod(added_factors)

for i in 1:20
    if number % i != 0
        println("Not divisible by: ", i)
    end
end


@symbol_func function q_DV(cur_reactor::AbstractReactor)
  cur_q = 18.31e-3

  cur_q *= 1 + cur_reactor.Q / 5.0

  cur_q *= 2 ^ 0.6

  cur_q /= ( 1.0 + cur_reactor.kappa_95 ^ 2 ) ^ 0.6

  cur_q *= P_H(cur_reactor)

  cur_q *= cur_reactor.I_P ^ 1.2

  cur_q /= cur_reactor.epsilon ^ 1.2

  cur_q /= cur_reactor.R_0 ^ 2.2

  cur_q
end


using HomotopyContinuation

@time begin
    @var x y
    affine_sqr =
        System([2.3 * x^2 + 1.2 * y^2 + 3x - 2y + 3, 2.3 * x^2 + 1.2 * y^2 + 5x + 2y - 5])
    res = solve(affine_sqr; compile = false)
end
res = solve(affine_sqr; compile = false, start_system = :total_degree)
@var x y z
proj_square = System([
    2.3 * x^2 + 1.2 * y^2 + 3x * z - 2y * z + (3 + 0im) * z^2,
    2.3 * x^2 + 1.2 * y^2 + 5x * z + 2y * z - (5 + 0im) * z^2,
])
solve(proj_square; compile = false)
solve(affine_sqr; compile = false, start_system = :total_degree)

@var x y
affine_ov = System([
    (x^2 + y^2 + x * y - 3) * (x + 3),
    (x^2 + y^2 + x * y - 3) * (y - x + 2.1),
    2x + 5y - 3,
])
solve(affine_ov; compile = false)

function toric_ed(A)
    d, n = size(A)
    @var t[1:d] y[1:n] u[1:n]

    φ = map(j -> prod(i -> t[i]^A[i, j], 1:d), 1:n)
    Dφ = [differentiate(φ[i], t[j]) for i = 1:n, j = 1:d]

    System([φ + y - u; Dφ' * y], parameters = u)
end

F = toric_ed([3 2 1 0; 0 1 2 3])
result = monodromy_solve(
    F,
    target_solutions_count = 21,
    max_loops_no_progress = 20,
    compile = false,
)

solve(
    F,
    solutions(result),
    start_parameters = parameters(result),
    target_parameters = randn(nparameters(F)),
    compile = false,
)


using GAIO

# logistic map
μ = 4.0
f(x) = μ.*x.*(1.0.-x)

center, radius = 0.5, 0.5
P = BoxPartition(Box(center, radius), 256)
F = BoxMap(f, P; no_of_points=400)

T = TransferOperator(F, P[:])
(λ, ev) = eigs(T)

plot(abs ∘ ev[1])


using JuMP
using GLPK
model = Model(GLPK.Optimizer)
n = 5
JuMP.@variables(model, begin
    y[1:n] >= 0
end)
rtg1_idx = [1,2,4]
rtg2_idx = [3]
rtg3_idx = [5]

a = [@constraint(model,sum(y[rtg1_idx])==12),
     @constraint(model,sum(y[rtg2_idx])==12),
     @constraint(model,sum(y[rtg3_idx])==12)
]

set_name(a[1], "efefef")
set_name(a[2], "efef")
set_name(a[3], "efefe")

optimize!(model)

model.obj_dict
value.(y)