text:
using Test
using SumOfSquares
using DynamicPolynomials

function term_test(optimizer,
                   config::MOIT.TestConfig,
                   cone::SumOfSquares.PolyJuMP.PolynomialSet)
    atol = config.atol
    rtol = config.rtol

    model = _model(optimizer)

    @variable(model, α)

    @polyvar x
    cref = @constraint(model, α * x^2 in cone)

    # See https://github.com/JuliaOpt/MathOptInterface.jl/issues/676
    @objective(model, Min, α + 1)
    optimize!(model)

    @test termination_status(model) == MOI.OPTIMAL
    @test objective_value(model) ≈ 1.0 atol=atol rtol=rtol

    @test primal_status(model) == MOI.FEASIBLE_POINT
    @test value(α) ≈ 0.0 atol=atol rtol=rtol

    p = gram_matrix(cref)
    @test getmat(p) ≈ zeros(1, 1) atol=atol rtol=rtol
    @test p.x == [x]

    @test dual_status(model) == MOI.FEASIBLE_POINT
    μ = dual(cref)
    @test μ isa AbstractMeasure{Float64}
    @test length(moments(μ)) == 1
    @test moment_value(moments(μ)[1]) ≈ 1.0 atol=atol rtol=rtol
    @test monomial(moments(μ)[1]) == x^2

    ν = moment_matrix(cref)
    @test getmat(ν) ≈ ones(1, 1) atol=atol rtol=rtol
    @test ν.x == [x]

    S = SumOfSquares.SOSPolynomialSet{
        SumOfSquares.FullSpace, typeof(cone), SumOfSquares.MonomialBasis,
        Monomial{true},MonomialVector{true},Tuple{}
    }
    @test list_of_constraint_types(model) == [(Vector{VariableRef}, S)]
    test_delete_bridge(
        model, cref, 1,
        ((MOI.VectorOfVariables, MOI.Nonnegatives, 0),
         (MOI.VectorAffineFunction{Float64},
          SumOfSquares.PolyJuMP.ZeroPolynomialSet{
              SumOfSquares.FullSpace, SumOfSquares.MonomialBasis,
              Monomial{true}, MonomialVector{true}},
          0)))
end
sos_term_test(optimizer, config)   = term_test(optimizer, config, SOSCone())
sdsos_term_test(optimizer, config) = term_test(optimizer, config, SDSOSCone())
dsos_term_test(optimizer, config)  = term_test(optimizer, config, DSOSCone())


text:
function include_parameters(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    param_reference::UpdateRef,
    expression_name::Symbol,
    multiplier::Float64 = 1.0,
)
    @assert model_has_parameters(psi_container)
    time_steps = model_time_steps(psi_container)
    ## TODO, replace multiplier
    names = [r.name for r in ts_data]
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    param = get_parameter_array(container)
    mult = get_multiplier_array(container)
    expr = get_expression(psi_container, expression_name)
    for t in time_steps, r in ts_data
        param[r.name, t] = PJ.add_parameter(psi_container.JuMPmodel, r.timeseries[t])
        mult[r.name, t] = r.multiplier * multiplier
        add_to_expression!(
            expr,
            r.bus_number,
            t,
            param[r.name, t],
            r.multiplier * multiplier,
        )
    end
    return param
end

ext: jl


text:
"""
    loglikelihood(d::UnivariateDistribution, y::AbstractVector, μ::AbstractVector)

Calculates the loglikelihood of observing `y` given mean `μ` and some distribution
`d`.

Note that loglikelihood is the sum of the logpdfs for each observation.
For each logpdf from Normal, Gamma, and InverseGaussian, we scale by dispersion.
"""
function loglikelihood(d::UnivariateDistribution, y::AbstractVector{T},
                       μ::AbstractVector{T}) where {T <: Float}
    logl = 0.0
    ϕ = MendelIHT.deviance(d, y, μ) / length(y)
    @inbounds for i in eachindex(y)
        logl += loglik_obs(d, y[i], μ[i], 1, ϕ) #wt = 1 because only have 1 sample to estimate a given mean
    end
    return T(logl)
end

"""
This function is taken from GLM.jl from:
https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_JuliaStats_GLM.jl_blob_956a64e7df79e80405867238781f24567bd40c78_src_glmtools.jl-23L445&d=DwIGaQ&c=sJ6xIWYx-zLMB3EPkvcnVg&r=7sbWVWEGF5cmtB61wl7FFg&m=t0UYMxl1l6T9gQwevDjzKZl1EUq7cxc1N1Q251BQAUU&s=T5Tp_cvAeqiX2CbgYRm0yhX-Uzmeg5rRbOXvyDiwq_M&e=

Putting it here because it was not exported.

`wt`: the number of observations (y) used to estimate μ
"""
function loglik_obs end

loglik_obs(::Bernoulli, y, μ, wt, ϕ) = wt*logpdf(Bernoulli(μ), y)
loglik_obs(::Binomial, y, μ, wt, ϕ) = logpdf(Binomial(Int(wt), μ), Int(y*wt))
loglik_obs(::Gamma, y, μ, wt, ϕ) = wt*logpdf(Gamma(inv(ϕ), μ*ϕ), y)
loglik_obs(::InverseGaussian, y, μ, wt, ϕ) = wt*logpdf(InverseGaussian(μ, inv(ϕ)), y)
loglik_obs(::Normal, y, μ, wt, ϕ) = wt*logpdf(Normal(μ, sqrt(ϕ)), y)
loglik_obs(::Poisson, y, μ, wt, ϕ) = wt*logpdf(Poisson(μ), y)
# We use the following parameterization for the Negative Binomial distribution:
#    (Γ(θ+y) / (Γ(θ) * y!)) * μ^y * θ^θ / (μ+θ)^{θ+y}
# The parameterization of NegativeBinomial(r=θ, p) in Distributions.jl is
#    Γ(θ+y) / (y! * Γ(θ)) * p^θ(1-p)^y
# Hence, p = θ/(μ+θ)
loglik_obs(d::NegativeBinomial, y, μ, wt, ϕ) = wt*logpdf(NegativeBinomial(d.r, d.r/(μ+d.r)), y)

"""
The deviance of a GLM can be evaluated as the sum of the squared deviance residuals. Calculation
of sqared deviance residuals is accomplished by `devresid` which is implemented in GLM.jl
"""
function deviance(d::UnivariateDistribution, y::AbstractVector{T}, μ::AbstractVector{T}) where {T <: Float}
    dev = 0.0
    @inbounds for i in eachindex(y)
        dev += devresid(d, y[i], μ[i])
    end
    return dev
end

function update_μ!(μ::AbstractVector{T}, xb::AbstractVector{T}, l::Link) where {T <: Float}
    @inbounds for i in eachindex(μ)
        μ[i] = linkinv(l, xb[i])
    end
end

"""
This function update the linear predictors `xb` with the new proposed b. We clamp the max
value of each entry to (-20, 20) because certain distributions (e.g. Poisson) have exponential
link functions, which causes overflow.
"""
function update_xb!(v::IHTVariable{T}, x::Union{SnpArray, AbstractMatrix},
                    z::AbstractMatrix{T}) where {T <: Float}
    typeof(x) == SnpArray ? copyto!(v.xk, @view(x[:, v.idx]), center=true, scale=true) : copyto!(v.xk, @view(x[:, v.idx]))
    A_mul_B!(v.xb, v.zc, v.xk, z, view(v.b, v.idx), v.c)
    clamp!(v.xb, -20, 20)
    clamp!(v.zc, -20, 20)
end

# For NegativeBinomial
function update_r!(d::NegativeBinomial, y::AbstractVector{T}, μ::AbstractVector{T}) where {T <: Float}
    num = zero(T)
    den = zero(T)
    for i in eachindex(y)
        for j = 0:y[i] - 1
            num = num + (d.r /(d.r + j))  # numerator for r
        end
        p = d.r / (d.r + μ[i])
        den = den + log(p)  # denominator for r
    end
    new_r = -(num)/den
    d = NegativeBinomial(new_r, d.p)
    return d
end

"""
    score = X^T * W * (y - g(x^T b))

Calculates the score (gradient) for different glm models.

W is a diagonal matrix where w[i, i] = g'(x^T b) / var(μ).
"""
function score!(d::UnivariateDistribution, l::Link, v::IHTVariable{T},
                x::Union{SnpBitMatrix, AbstractMatrix}, z::AbstractMatrix{T},
                y::AbstractVector{T}) where {T <: Float}
    @inbounds for i in eachindex(y)
        # η = clamp(v.xb[i] + v.zc[i], -20, 20)
        η = v.xb[i] + v.zc[i]
        w = mueta(l, η) / glmvar(d, v.μ[i])
        v.r[i] = w * (y[i] - v.μ[i])
    end
    At_mul_B!(v.df, v.df2, x, z, v.r, v.r)
end

"""
This function computes the gradient step v.b = P_k(β + η∇f(β)) and updates idx and idc.
"""
function _iht_gradstep(v::IHTVariable{T}, η::T, J::Int, k::Int,
                       full_grad::Vector{T}) where {T <: Float}
    lb = length(v.b)
    lw = length(v.weight)
    lg = length(v.group)
    lf = length(full_grad)

    # take gradient step: b = b + ηv, v = score
    BLAS.axpy!(η, v.df, v.b)
    BLAS.axpy!(η, v.df2, v.c)

    # scale model by weight vector, if supplied
    if lw == 0
        copyto!(@view(full_grad[1:lb]), v.b)
        copyto!(@view(full_grad[lb+1:lf]), v.c)
    else
        copyto!(@view(full_grad[1:lb]), v.b .* @view(v.weight[1:lb]))
        copyto!(@view(full_grad[lb+1:lf]), v.c .* @view(v.weight[lb+1:lf]))
    end

    # project to sparsity
    lg == 0 ? project_k!(full_grad, k) : project_group_sparse!(full_grad, v.group, J, k)

    # unweight the model after projection
    if lw == 0
        copyto!(v.b, @view(full_grad[1:lb]))
        copyto!(v.c, @view(full_grad[lb+1:lf]))
    else
        copyto!(v.b, @view(full_grad[1:lb]) ./ @view(v.weight[1:lb]))
        copyto!(v.c, @view(full_grad[lb+1:lf]) ./ @view(v.weight[lb+1:lf]))
    end

    #recombute support
    v.idx .= v.b .!= 0
    v.idc .= v.c .!= 0

    # if more than J*k entries are selected, randomly choose J*k of them
    _choose!(v, J, k)

    # make necessary resizing since grad step might include/exclude non-genetic covariates
    check_covariate_supp!(v)
end

"""
When initializing the IHT algorithm, take largest elements in magnitude of each group of
the score as nonzero components of b. This function set v.idx = 1 for those indices.

`J` is the maximum number of active groups, and `k` is the maximum number of predictors per
group.
"""
function init_iht_indices!(v::IHTVariable{T}, x::Union{SnpBitMatrix, AbstractMatrix},
                           z::AbstractMatrix{T},y::Vector{T}, d::UnivariateDistribution,
                           l::Link, J::Int, k::Int) where {T <: Float}
    # find the intercept by Newton's method
    ybar = mean(y)
    for iteration = 1:20
        g1 = linkinv(l, v.c[1])
        g2 = mueta(l, v.c[1])
        v.c[1] = v.c[1] - clamp((g1 - ybar) / g2, -1.0, 1.0)
        abs(g1 - ybar) < 1e-10 && break
    end
    v.zc .= z * v.c

    # update mean vector and use them to compute score (gradient)
    update_μ!(v.μ, v.xb + v.zc, l)
    score!(d, l, v, x, z, y)

    # find J*k largest entries in magnitude and set everything else to 0.
    a = partialsort([v.df; v.df2], k * J, by=abs, rev=true)
    v.idx .= abs.(v.df) .>= abs(a)
    v.idc .= abs.(v.df2) .>= abs(a)

    # Choose randomly if more are selected
    _choose!(v, J, k)

    # make necessary resizing when necessary
    check_covariate_supp!(v)
end

"""
if more than J*k entries are selected after projection, randomly select top J*k entries.
This can happen if entries of b are equal to each other.
"""
function _choose!(v::IHTVariable{T}, J::Int, k::Int) where {T <: Float}
    while sum(v.idx) + sum(v.idc) > J * k
        idx_length = length(v.idx)
        all_idx = [v.idx; v.idc]

        nonzero_idx = findall(x -> x == true, all_idx) #find non-0 location
        pos = nonzero_idx[rand(1:length(nonzero_idx))] #randomly set a non-0 entry to 0
        pos > idx_length ? v.idc[pos - idx_length] = 0 : v.idx[pos] = 0
    end
end

"""
In `_init_iht_indices` and `_iht_gradstep`, if non-genetic cov got
included/excluded, we must resize xk and gk
"""
function check_covariate_supp!(v::IHTVariable{T}) where {T <: Float}
    if sum(v.idx) != size(v.xk, 2)
        v.xk = zeros(T, size(v.xk, 1), sum(v.idx))
        v.gk = zeros(T, sum(v.idx))
    end
end

"""
This function returns true if backtracking condition is met. Currently, backtracking condition
includes either one of the following:
    1. New loglikelihood is smaller than the old one
    2. Current backtrack (`η_step`) exceeds maximum allowed backtracking (`nstep`, default = 3)

Note for Posison, NegativeBinomial, and Gamma, we require model coefficients to be
"small" to prevent loglikelihood blowing up in first few iteration. This is accomplished
by clamping η = xb values to be in (-20, 20)
"""
function _iht_backtrack_(logl::T, prev_logl::T, η_step::Int64, nstep::Int64) where {T <: Float}
    (prev_logl > logl) && (η_step < nstep)
end

"""
    std_reciprocal(x::SnpBitMatrix, mean_vec::Vector{T})

Compute the standard error of each columns of a SnpArray in place.

`mean_vec` stores the mean for each SNP. Note this function assumes all SNPs
are not missing. Otherwise, the inner loop should only add if data not missing.
"""
function std_reciprocal(x::SnpBitMatrix, mean_vec::Vector{T}) where {T <: Float}
    m, n = size(x)
    @assert n == length(mean_vec) "number of columns of snpmatrix doesn't agree with length of mean vector"
    std_vector = zeros(T, n)

    @inbounds for j in 1:n
        @simd for i in 1:m
            a1 = x.B1[i, j]
            a2 = x.B2[i, j]
            std_vector[j] += (convert(T, a1 + a2) - mean_vec[j])^2
        end
        std_vector[j] = 1.0 / sqrt(std_vector[j] / (m - 1))
    end
    return std_vector
end

"""
    standardize!(z::Matrix{Float64})

Standardizes each column of `z` to mean 0 and variance 1. Make sure you
do not standardize the intercept.
"""
@inline function standardize!(z::AbstractMatrix{Float64})
    n, q = size(z)
    μ = _mean(z)
    σ = _std(z, μ)

    @inbounds for j in 1:q
        @simd for i in 1:n
            z[i, j] = (z[i, j] - μ[j]) * σ[j]
        end
    end
end

@inline function _mean(z::AbstractMatrix{Float64})
    n, q = size(z)
    μ = zeros(q)
    @inbounds for j in 1:q
        tmp = 0.0
        @simd for i in 1:n
            tmp += z[i, j]
        end
        μ[j] = tmp / n
    end
    return μ
end

function _std(z::AbstractMatrix{Float64}, μ::Vector{Float64})
    n, q = size(z)
    σ = zeros(q)

    @inbounds for j in 1:q
        @simd for i in 1:n
            σ[j] += (z[i, j] - μ[j])^2
        end
        σ[j] = 1.0 / sqrt(σ[j] / (n - 1))
    end
    return σ
end

"""
    project_k!(x::AbstractVector, k::Integer)

Sets all but the largest `k` entries of `x` to 0.

# Examples:
```julia-repl
using MendelIHT
x = [1.0; 2.0; 3.0]
project_k!(x, 2) # keep 2 largest entry
julia> x
3-element Array{Float64,1}:
 0.0
 2.0
 3.0
```

# Arguments:
- `x`: the vector to project.
- `k`: the number of components of `x` to preserve.
"""
function project_k!(x::AbstractVector{T}, k::Int64) where {T <: Float}
    a = abs(partialsort(x, k, by=abs, rev=true))
    @inbounds for i in eachindex(x)
        abs(x[i]) < a && (x[i] = zero(T))
    end
end

"""
    project_group_sparse!(y::AbstractVector, group::AbstractVector, J::Integer, k::Integer)

Projects the vector y onto the set with at most J active groups and at most
k active predictors per group. Currently assumes there are no unknown or overlaping
group membership.

# Examples
```julia-repl
using MendelIHT
J, k, n = 2, 3, 20
y = collect(1.0:20.0)
y_copy = copy(y)
group = rand(1:5, n)
project_group_sparse!(y, group, J, k)
for i = 1:length(y)
    println(i,"  ",group[i],"  ",y[i],"  ",y_copy[i])
end
```

# Arguments
- `y`: The vector to project
- `group`: Vector encoding group membership
- `J`: Max number of non-zero group
- `k`: Max number of non-zero predictor per group`
"""
function project_group_sparse!(y::AbstractVector{T}, group::AbstractVector{Int64},
    J::Int64, k::Int64) where {T <: Float}
    groups = maximum(group)          # number of groups
    group_count = zeros(Int, groups) # counts number of predictors in each group
    group_norm = zeros(groups)       # l2 norm of each group
    perm = zeros(Int64, length(y))   # vector holding the permuation vector after sorting
    sortperm!(perm, y, by = abs, rev = true)

    #calculate the magnitude of each group, where only top predictors contribute
    for i in eachindex(y)
        j = perm[i]
        n = group[j]
        if group_count[n] < k
            group_norm[n] = group_norm[n] + y[j]^2
            group_count[n] = group_count[n] + 1
        end
    end

    #go through the top predictors in order. Set predictor to 0 if criteria not met
    group_rank = zeros(Int64, length(group_norm))
    sortperm!(group_rank, group_norm, rev = true)
    group_rank = invperm(group_rank)
    fill!(group_count, 1)
    for i in eachindex(y)
        j = perm[i]
        n = group[j]
        if (group_rank[n] > J) || (group_count[n] > k)
            y[j] = 0.0
        else
            group_count[n] = group_count[n] + 1
        end
    end
end

"""
    maf_weights(x::SnpArray; max_weight::T = Inf)

Calculates the prior weight based on minor allele frequencies.

Returns an array of weights where `w[i] = 1 / (2 * sqrt(p[i] (1 - p[i]))) ∈ (1, ∞).`
Here `p` is the minor allele frequency computed by `maf()` in SnpArrays.

- `x`: A SnpArray
- `max_weight`: Maximum weight for any predictor. Defaults to `Inf`.
"""
function maf_weights(x::SnpArray; max_weight::T = Inf) where {T <: Float}
    p = maf(x)
    p .= 1 ./ (2 .* sqrt.(p .* (1 .- p)))
    clamp!(p, 1.0, max_weight)
    return p
end

"""
Function that saves `b`, `xb`, `idx`, `idc`, `c`, and `zc` after each iteration.
"""
function save_prev!(v::IHTVariable{T}) where {T <: Float}
    copyto!(v.b0, v.b)     # b0 = b
    copyto!(v.xb0, v.xb)   # Xb0 = Xb
    copyto!(v.idx0, v.idx) # idx0 = idx
    copyto!(v.idc0, v.idc) # idc0 = idc
    copyto!(v.c0, v.c)     # c0 = c
    copyto!(v.zc0, v.zc)   # Zc0 = Zc
end

"""
Computes the best step size η = v'v / v'Jv

Here v is the score and J is the expected information matrix, which is
computed by J = g'(xb) / var(μ), assuming dispersion is 1
"""
function iht_stepsize(v::IHTVariable{T}, z::AbstractMatrix{T},
                      d::UnivariateDistribution, l::Link) where {T <: Float}

    # first store relevant components of gradient
    copyto!(v.gk, view(v.df, v.idx))
    A_mul_B!(v.xgk, v.zdf2, v.xk, view(z, :, v.idc), v.gk, view(v.df2, v.idc))

    #use zdf2 as temporary storage
    v.xgk .+= v.zdf2
    v.zdf2 .= mueta.(l, v.xb + v.zc).^2 ./ glmvar.(d, v.μ)

    # now compute and return step size. Note non-genetic covariates are separated from x
    numer = sum(abs2, v.gk) + sum(abs2, @view(v.df2[v.idc]))
    denom = Transpose(v.xgk) * Diagonal(v.zdf2) * v.xgk
    return (numer / denom) :: T
end

"""
This is a wrapper linear algebra function that computes [C1 ; C2] = [A1 ; A2] * [B1 ; B2]
where A1 is a snpmatrix and A2 is a dense Matrix{Float}. Used for cleaner code.

Here we are separating the computation because A1 is stored in compressed form while A2 is
uncompressed (float64) matrix. This means that they cannot be stored in the same data
structure.
"""
function A_mul_B!(C1::AbstractVector{T}, C2::AbstractVector{T}, A1::SnpBitMatrix{T},
        A2::AbstractMatrix{T}, B1::AbstractVector{T}, B2::AbstractVector{T}) where {T <: Float}
    SnpArrays.mul!(C1, A1, B1)
    LinearAlgebra.mul!(C2, A2, B2)
end

function A_mul_B!(C1::AbstractVector{T}, C2::AbstractVector{T}, A1::AbstractMatrix{T},
        A2::AbstractMatrix{T}, B1::AbstractVector{T}, B2::AbstractVector{T}) where {T <: Float}
    LinearAlgebra.mul!(C1, A1, B1)
    LinearAlgebra.mul!(C2, A2, B2)
end

"""
    initialize_beta!(v::IHTVariable, y::AbstractVector, x::SnpArray, d::UnivariateDistribution, l::Link)

Fits a univariate regression (+ intercept) with each β_i corresponding to `x`'s predictor.

Used to find a good starting β. Fitting is done using scoring (newton) algorithm
implemented in `GLM.jl`. The intial intercept is separately fitted using in init_iht_indices().

Note: this function is quite slow and not memory efficient.
"""
function initialize_beta!(v::IHTVariable{T}, y::AbstractVector{T}, x::Union{SnpArray, AbstractMatrix},
                          d::UnivariateDistribution, l::Link) where {T <: Float}
    n, p = size(x)
    temp_matrix = ones(n, 2)           # n by 2 matrix of the intercept and 1 single covariate
    temp_glm = initialize_glm_object() # preallocating in a dumb ways

    intercept = 0.0
    if typeof(x) == SnpArray
        for i in 1:p
            copyto!(@view(temp_matrix[:, 2]), @view(x[:, i]), center=true, scale=true)
            temp_glm = fit(GeneralizedLinearModel, temp_matrix, y, d, l)
            v.b[i] = temp_glm.pp.beta0[2]
        end
    else
        for i in 1:p
            temp_matrix[:, 2] .= x[:, i]
            temp_glm = fit(GeneralizedLinearModel, temp_matrix, y, d, l)
            v.b[i] = temp_glm.pp.beta0[2]
        end
    end
end

"""
This is a wrapper linear algebra function that computes [C1 ; C2] = [A1 ; A2]^T * [B1 ; B2]
where A1 is a snpmatrix and A2 is a dense Matrix{Float}. Used for cleaner code.

Here we are separating the computation because A1 is stored in compressed form while A2 is
uncompressed (float64) matrix. This means that they cannot be stored in the same data
structure.
"""
function At_mul_B!(C1::AbstractVector{T}, C2::AbstractVector{T}, A1::SnpBitMatrix{T},
        A2::AbstractMatrix{T}, B1::AbstractVector{T}, B2::AbstractVector{T}) where {T <: Float}
    SnpArrays.mul!(C1, Transpose(A1), B1)
    LinearAlgebra.mul!(C2, Transpose(A2), B2)
end

function At_mul_B!(C1::AbstractVector{T}, C2::AbstractVector{T}, A1::AbstractMatrix{T},
        A2::AbstractMatrix{T}, B1::AbstractVector{T}, B2::AbstractVector{T}) where {T <: Float}
    LinearAlgebra.mul!(C1, Transpose(A1), B1)
    LinearAlgebra.mul!(C2, Transpose(A2), B2)
end

"""
This function initializes 1 instance of a GeneralizedLinearModel(G<:GlmResp, L<:LinPred, Bool).
"""
function initialize_glm_object()
    d = Bernoulli
    l = canonicallink(d())
    x = rand(100, 2)
    y = rand(0:1, 100)
    return fit(GeneralizedLinearModel, x, y, d(), l)
end

function mle_for_θ(y::AbstractVector, μ::AbstractVector;
                    θ::AbstractFloat = 1.0, maxIter=100, convTol=1.e-6)

    function first_derivative(θ::Real)
        tmp(yi, μi) = -(yi+θ)/(μi+θ) - log(μi+θ) + 1 + log(θ) + digamma(θ+yi) - digamma(θ)
        return sum(tmp(yi, μi) for (yi, μi) in zip(y, μ))
    end

    function second_derivative(θ::Real)
        tmp(yi, μi) = (yi+θ)/(μi+θ)^2 - 2/(μi+θ) + 1/θ + trigamma(θ+yi) - trigamma(θ)
        return sum(tmp(yi, μi) for (yi, μi) in zip(y, μ))
    end

    function negbin_loglikelihood(θ::Real)
        d_newton = NegativeBinomial(θ, 0.5)
        return MendelIHT.loglikelihood(d_newton, y, μ)
    end

    function newton_increment(θ::Real)
        # use gradient descent if hessian not positive definite
        dx  = first_derivative(θ)
        dx2 = second_derivative(θ)
        if dx2 < 0
            increment = first_derivative(θ) / second_derivative(θ)
        else
            increment = first_derivative(θ)
        end
        # println(dx2)
        return increment
    end

    new_θ    = 1.0
    stepsize = 1.0
    for i in 1:maxIter

        # run 1 iteration of Newton's algorithm
        increment = newton_increment(θ)
        new_θ = θ - stepsize * increment

        # linesearch
        old_logl = negbin_loglikelihood(θ)
        for j in 1:20
            if new_θ <= 0
                stepsize = stepsize / 2
                new_θ = θ - stepsize * increment
            else
                new_logl = negbin_loglikelihood(new_θ)
                if old_logl >= new_logl
                    stepsize = stepsize / 2
                    new_θ = θ - stepsize * increment
                else
                    break
                end
            end
        end

        #check convergence
        if abs(θ - new_θ) <= convTol
            return new_θ
        else
            θ = new_θ
        end
    end

    return θ
end

"""
    naive_impute(x, destination)

Imputes missing entries of a SnpArray using the mode of each SNP, and
saves the result in a new file called destination in current directory.
Non-missing entries are the same.
"""
function naive_impute(x::SnpArray, destination::String)
    n, p = size(x)
    y = SnpArray(destination, n, p)

    @inbounds for j in 1:p

        #identify mode
        entry0, entry1, entry2 = 0, 0, 0
        for i in 1:n
            if x[i, j] == 0x00
                y[i, j] = 0x00
                entry0 += 1
            elseif x[i, j] == 0x02
                y[i, j] = 0x02
                entry1 += 1
            elseif x[i, j] == 0x03
                y[i, j] = 0x03
                entry2 += 1
            end
        end
        most_often = max(entry0, entry1, entry2)
        missing_entry = 0x00
        if most_often == entry1
            missing_entry = 0x02
        elseif most_often == entry2
            missing_entry = 0x03
        end

        # impute
        for i in 1:n
            if x[i, j] == 0x01
                y[i, j] = missing_entry
            end
        end
    end

    return nothing
end


text:
# this file contains code for the simulated annealing TSP heuristic

"""
  simulated_annealing(distmat; ...)

Use a simulated annealing strategy to return a closed tour. The temperature
decays exponentially from `init_temp` to `final_temp`. Return a tuple `(path, cost)`.

### Optional keyword arguments:
- `steps::Int = 50n^2`: number of steps to take; defaults to 50n^2 where n is number of cities
- `num_starts::Int = 1`: number of times to run the simulated annealing algorithm, each time
    starting with a random path, or `init_path` if non-null. Defaults to 1.
- `init_temp::Real = exp(8)`: initial temperature which controls initial chance of accepting an
    inferior tour.
- `final_temp::Real = exp(-6.5)` final temperature which controls final chance of accepting an
    inferior tour; lower values roughly correspond to a longer period of 2-opt.
- `init_path::Union{Vector{Int}, Nothing} = nothing`: path to start the annealing from.
    Either a `Vector{Int}` or `nothing`. If a `Vector{Int}` is given, it will be used as the initial path;
    otherwise a random initial path is used.
"""
function simulated_annealing(distmat::Matrix{T} where {T<:Real};
                             steps = 50*length(distmat),
                             num_starts = 1,
                             init_temp = exp(8), final_temp = exp(-6.5),
                             init_path::Union{Vector{Int}, Nothing} = nothing)

    # check inputs
    n = check_square(distmat, "Must pass a square distance matrix to simulated_annealing.")

    # cooling rate: we multiply by a constant mult each step
    cool_rate = (final_temp / init_temp)^(1 / (steps - 1))

    # do SA with a single starting path
    function sahelper!(path)
        temp = init_temp / cool_rate # divide by cool_rate so when we first multiply we get init_temp
        n = size(distmat, 1)
        cost_cur = pathcost(distmat, path)

        for i in 1:steps
            temp *= cool_rate

            # take a step
            # keep first and last cities fixed
            first, last = rand(2:n), rand(2:n)
            if first > last
                first, last = last, first
            end
            cost_other = pathcost_rev(distmat, path, first, last)
            @fastmath accept = cost_other < cost_cur ? true : rand() < exp((cost_cur - cost_other) / temp)
            # should we accept?
            if accept
                reverse!(path, first, last)
                cost_cur = cost_other
            end
        end

        return path, cost_cur
    end

    # unpack the initial path
    if init_path == nothing
        randstart = true
        path = randpath(n)
    else
        if !legal_circuit(init_path)
            error("The init_path passed to simulated_annealing must be a legal circuit.")
        end
        randstart = false
        path = init_path
    end
    cost = pathcost(distmat, path)
    for _ in 1:num_starts
        path_this_start = randstart ? randpath(n) : deepcopy(init_path)
        otherpath, othercost = sahelper!(path_this_start)
        if othercost < cost
            cost = othercost
            path = otherpath
        end
    end

    return path, cost
end

function randpath(n)
    path = 1:n |> collect |> shuffle
    push!(path, path[1]) # loop
    return path
end


text:
# # Instructions for Assignment01
#
# ## Introduction
#
# In this assignment,
# you will learn how to use github classroom
# to download, complete, and test your assignments.
# Make sure you have followed [the instructions](https://wellesley-bisc195.github.io/assignments/Assignment01)
# to set up your account and join the classroom.

# To view additional instructions for this assignment, #jl
# see https://wellesley-bisc195.github.io/assignments/Assignment01 #jl

# To view a rendered version of this document, #jl
# see https://wellesley-bisc195.github.io/assignments/Assignment01/#assignment01_code #jl

# ## Question 1
#
# Did you complete the instructions for modifying
# `some_variables.jl` and `some_functions.jl`?
# You should have changed the variables `n` and `a`
# to `my_name` and `my_age` respectively,
# and you should be able to execute `julia src/run.jl`
# without errors.
# Once you can, type `yes` between the quotations below
#
# Note: this syntax (`question1 = "yes"`)
# is called "variable assignment,"
# and we'll learn more about this in the next lesson.

question1 = ""

# ## Question 2
#
# What is the shell command to list the contents
# of the working directory,
# **including hidden files**?
#
# Enter the command between the quotes below.

question2 = ""

# ## Question 3
#
# What is a shell command to create a new directory
# called "q3" in your current working directory?
#
# Enter the command between the quotes below.

question3 = ""

# ## Question 4
#
# If you haven't already, take a look at the syllabus.
# Which of the following is closest to the course policy
# on collaborating / sharing code?
#
# 1. I shouldn't discuss anything about the course with other students
# 2. I can freely discuss the non-assignment parts of the course,
#    but shouldn't talk about assignments with other students.
# 3. I can freely discuss all aspects of the course,
#    but shouldn't share my assignment code or look at someone else's assignment code.
# 4. I can freely discuss all aspects of the course,
#    and freely share all of my code - open source!
#
# Replace the `0` below with the correct number

question4 = 0

# ## Question 5
#
# What fraction of the course grade
# is based on completing assignments like this one?
# Express the answer as a number between 0 and 1

question5 = 0.0


text:
using InducingPoints
using Distances
using KernelFunctions
using KernelFunctions: ColVecs
using Test
using Random: seed!
include("test_utils.jl")

@testset "InducingPoints.jl" begin
    @testset "Offline" begin
        for file in readdir("offline")
            include(joinpath("offline", file))
        end
    end
    @testset "Online" begin
        for file in readdir("online")
            include(joinpath("online", file))
        end
    end
end


text:
###############################################################################
# other stuff, trivia and non-trivia
###############################################################################

function Nemo.PolynomialRing(R::Nemo.Ring, n::Int, s::String="x";
                             cached::Bool = false, ordering::Symbol = :lex)
  return Nemo.PolynomialRing(R, ["$s$i" for i=1:n], cached = cached,
                                                    ordering = ordering)
end

function total_degree(f::fmpq_mpoly)
  return Int(maximum(degree(f, i) for i=1:nvars(parent(f))))
end

function total_degree(f::Generic.MPoly)
  return Int(maximum([sum(f.exps[:, i]) for i=1:length(f)]))
end

function add!(c::fmpq_mpoly, a::fmpq_mpoly, b::fmpq_mpoly)
  ccall((:fmpq_mpoly_add, :libflint), Void,
        (Ref{fmpq_mpoly}, Ref{fmpq_mpoly}, Ref{fmpq_mpoly}, Ref{FmpqMPolyRing}),
        c, a, b, c.parent)
  return c
end


text:
using DataKnots
using CSV

employee_csv = """
    name,department,position,salary,rate
    "ANTHONY A","POLICE","POLICE OFFICER",72510,
    "DANIEL A","FIRE","FIRE FIGHTER-EMT",95484,
    "JAMES A","FIRE","FIRE ENGINEER-EMT",103350,
    "JEFFERY A","POLICE","SERGEANT",101442,
    "NANCY A","POLICE","POLICE OFFICER",80016,
    "ROBERT K","FIRE","FIRE FIGHTER-EMT",103272,
    "ALBA M","POLICE","POLICE CADET",,9.46
    "LAKENYA A","OEMC","CROSSING GUARD",,17.68
    "DORIS A","OEMC","CROSSING GUARD",,19.38
    "BRENDA B","OEMC","TRAFFIC CONTROL AIDE",64392,
    """ |> IOBuffer |> CSV.File

overtime_csv = """
    employee,month,amount
    "DANIEL A","2018-02",108
    "JAMES A","2018-01",8776
    "JAMES A","2018-03",351
    "JAMES A","2018-04",10532
    "JAMES A","2018-05",351
    "JAMES A","2018-06",8776
    "JAMES A","2018-07",10532
    "JEFFERY A","2018-05",1319
    "NANCY A","2018-01",173
    "NANCY A","2018-02",461
    "NANCY A","2018-03",461
    "NANCY A","2018-04",1056
    "NANCY A","2018-05",1933
    "ROBERT K","2018-05",1754
    """ |> IOBuffer |> CSV.File

chicago = let
    source = DataKnot(:employee => employee_csv,
                      :overtime => overtime_csv)
    @query source {
       department => begin
         employee.keep(name)
         join(begin
           overtime
           filter(employee==name)
           collect(employee=>nothing)
         end)
         group(name => department)
         collect(begin
           employee
           collect(department=>nothing)
         end)
       end}
end


text:
################################################################################
# MIT License
#
# Copyright (c) 2021 Hajime Nakagami<nakagami@gmail.com>
#
# 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.
################################################################################

mutable struct Transaction
    wp::WireProtocol
    handle::Int32
    need_begin::Bool

    function Transaction(wp::WireProtocol)
        tpb::Vector{UInt8} = [
            isc_tpb_version3, isc_tpb_write, isc_tpb_wait, isc_tpb_autocommit
        ]
        _op_transaction(wp, tpb)
        trans_handle, _, _ = _op_response(wp)
        if trans_handle < 0
            throw(DomainError("transaction error"))
        end
        new(wp, trans_handle, true)
    end

end

function close!(trans::Transaction)
    if trans.handle != -1
        _op_rollback(trans.wp, trans.handle)
        _op_response(trans.wp)
    end
    trans.handle = -1
end


text:
function parse_rules(s)
    rules = Dict{String,Vector{Tuple{String,Int}}}()
    for line in readlines(IOBuffer(s))
        outer, inners = split(line, " bags contain ")
        if startswith(inners, "no other bags.")
            continue
        end
        inners = string(inners)
        inners = replace(inners, r"(bags?|\.)" => "")
        for inner in split(inners, ", ")
            parts = split(strip(inner), ' ', limit=2)
            n = parse(Int, parts[1])
            formatted_inner = parts[2]
            if outer in keys(rules)
                push!(rules[outer], (formatted_inner, n))
            else
                rules[outer] = [(formatted_inner, n),]
            end
        end
    end
    return rules
end

function fill_can_be_included_in(rules, bag, factor)
    acc::Int = 0
    if bag in keys(rules)
        for (next, n) in rules[bag]
            acc += n * factor
            if next in keys(rules)
                acc += fill_can_be_included_in(rules, next, n * factor)
            end
        end
    end
    return acc
end

function run(s)
    # Your code here
    rules = parse_rules(s)
    return fill_can_be_included_in(rules, "shiny gold", 1)
end

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

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


text:
cleanup() = begin
    rm("matrices"; force=true, recursive=true)
    rm("something_or_other_1=0_5_5_7.json"; force=true, recursive=true)
    rm("something_or_other_1=0_5_5_7.mytricks.h5"; force=true, recursive=true)
    rm("something_or_other_1=0_5_5_7.h5"; force=true, recursive=true)
    rm("numbers.dir"; force=true, recursive=true)
    rm("numbers.h5"; force=true, recursive=true)
    rm("stuff.h5"; force=true, recursive=true)
    rm("datafile.h5"; force=true, recursive=true)
end

cleanup()
module mt0001
using Test
using DataDrop: clean_file_name, with_extension
function test()
    name = "something or other.1=0.5_5:7"
    s =  clean_file_name(name)
    @test s == "something_or_other_1=0_5_5_7"
    a = with_extension(s, ".ext")
    @test a == "something_or_other_1=0_5_5_7.ext"
    a = with_extension(s, "ext")
    @test a == "something_or_other_1=0_5_5_7.ext"
    a = with_extension(s * ".dat", "ext")
    @test a == "something_or_other_1=0_5_5_7.ext"
    a = with_extension(s * ".", "ext")
    @test a == "something_or_other_1=0_5_5_7.ext"
    a = with_extension(s * ".x", "ext")
    @test a == "something_or_other_1=0_5_5_7.ext"
    true
end
end
using .mt0001
mt0001.test()

cleanup()
module mt0002
using Test
using DataDrop: clean_file_name, with_extension
using DataDrop: store_json, retrieve_json
function test()
    name = "something or other.1=0.5_5:7"
    ext = ".json"
    s =  clean_file_name(name)
    @test s == "something_or_other_1=0_5_5_7"
    a = with_extension(s, ext)
    d = Dict("name" => a, "b" => 1, "c" => [3, 1, 3])
    store_json(a, d)
    d1 = retrieve_json(a)
    @test d == d1
    true
end
end
using .mt0002
mt0002.test()

cleanup()
module mt0003
using Test
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    name = "something.or.other.1=0.5_5:7"
    ext = ""
    s =  clean_file_name(name)
    @test s == "something_or_other_1=0_5_5_7"
    a = with_extension(s, ext)
    d = rand(3, 2)
    store_matrix(a, d)
    d1 = retrieve_matrix(a)
    @test d == d1
    true
end
end
using .mt0003
mt0003.test()

cleanup()
module mt0004
using Test
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices"
    mkpath(p)
    name = joinpath(p, "matrix d")
    ext = ""
    s = with_extension(clean_file_name(name), ext)
    d = rand(3, 2)
    store_matrix(s, d)
    d1 = retrieve_matrix(s)
    @test d == d1
    true
end
end
using .mt0004
mt0004.test()

cleanup()
module mt0005
using Test
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices"
    mkpath(p)
    ext = ""
    for t in [Int, Float32, Float64, Complex{Float64}]
        name = joinpath(p, "matrix d " * string(t))
        s = with_extension(clean_file_name(name), ext)
        d = rand(t, 3, 2)
        store_matrix(s, d)
        @test isfile(s * ".h5") == true
        d1 = retrieve_matrix(s)
        @test d == d1
    end
    true
end
end
using .mt0005
mt0005.test()

cleanup()
module mt0006
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices"
    mkpath(p)
    name = joinpath(p, "matrix d")
    ext = ""
    s = with_extension(clean_file_name(name), ext)
    d = sprand(Float64, 5, 3, 0.75)
    store_matrix(s, d)
    d1 = retrieve_matrix(s)
    @test d == d1
    true
end
end
using .mt0006
mt0006.test()

cleanup()
module mt0007
using Test
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices"
    mkpath(p)
    ext = ""
    for t in [Int, Float32, Float64, Complex{Float64}]
        name = joinpath(p, "matrix B1 " * string(t))
        s = with_extension(clean_file_name(name), ext)
        d = rand(t, 3, 2)
        store_matrix(s, "/mydata/matrices/B1", d)
        @test isfile(s * ".h5") == true
        d1 = retrieve_matrix(s, "/mydata/matrices/B1")
        @test d == d1
    end
    true
end
end
using .mt0007
mt0007.test()

cleanup()
module mt0008
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices"
    mkpath(p)
    for t in [Int, Float32, Float64, Complex{Float64}]
        name = joinpath(p, "matrix S" * string(t))
        ext = "h5f"
        s = with_extension(clean_file_name(name), ext)
        d = sprand(t, 5, 3, 0.75)
        store_matrix(s, "/data", d)
        d1 = retrieve_matrix(s, "/data")
        @test d == d1
    end
    true
end
end
using .mt0008
mt0008.test()

cleanup()
module mt0009
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension, file_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices.dir"
    @test file_extension(p) == ".dir"
    p = "matrices.h5"
    @test file_extension(p) == ".h5"
    p = "matri.ces.dir"
    @test file_extension(p) == ".dir"
    p = "matrices."
    @test file_extension(p) == "."
    p = "matrices"
    @test file_extension(p) == ""
    true
end
end
using .mt0009
mt0009.test()

cleanup()
module mt0010
using Test
using DataDrop: clean_file_name, with_extension, file_extension
using DataDrop: store_value, retrieve_value
function test()
    p = "numbers.dir"
    a = 3.13
    store_value(p, "a", a)
    ra = retrieve_value(p, "a")
    @test a == ra
    Test.@test_throws ErrorException store_value(p, "/a", a)
    ra = retrieve_value(p, "/a")
    @test a == ra
    store_value(p, "/mynumber/a", a)
    ra = retrieve_value(p, "/mynumber/a")
    @test a == ra
    b = 61
    store_value(p, "/mynumber/b", b)
    rb = retrieve_value(p, "/mynumber/b")
    @test b == rb
    true
end
end
using .mt0010
mt0010.test()

cleanup()
module mt0011
using Test
using DataDrop: clean_file_name, with_extension, file_extension
using DataDrop: store_value, retrieve_value
function test()
    p = "numbers.h5"
    a = 3.13
    store_value(p, "/mynumber/a", a)
    ra = retrieve_value(p, "/mynumber/a")
    @test a == ra
    b = 61
    store_value(p, "/mynumber/b", b)
    rb = retrieve_value(p, "/mynumber/b")
    @test b == rb
    true
end
end
using .mt0011
mt0011.test()

cleanup()
module mt0012
using Test
using DataDrop: clean_file_name, with_extension, file_extension
using DataDrop: store_value, retrieve_value
function test()
    p = "stuff.h5"
    a = 3.13
    store_value(p, "/mynumber/a", a)
    ra = retrieve_value(p, "/mynumber/a")
    @test a == ra
    b = "61"
    store_value(p, "/mystring/b", b)
    rb = retrieve_value(p, "/mystring/b")
    @test b == rb
    true
end
end
using .mt0012
mt0012.test()

cleanup()
module mt0013
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix
function test()
    p = "matrices/afile.h5"
    mkpath(p)

    d = rand(Int64, 6, 3)
    store_matrix(joinpath(p, "afile.h5"), "/data/d", d)
    d1 = retrieve_matrix(joinpath(p, "afile.h5"), "/data/d")
    @test d == d1

    b = rand(Float64, 6, 3)
    store_matrix(joinpath(p, "afile.h5"), "/data/b", b)
    b1 = retrieve_matrix(joinpath(p, "afile.h5"), "/data/b")
    @test b == b1

    d1 = retrieve_matrix(joinpath(p, "afile.h5"), "/data/d")
    @test d == d1

    b1 = retrieve_matrix(joinpath(p, "afile.h5"), "/data/b")
    @test b == b1

    true
end
end
using .mt0013
mt0013.test()

cleanup()
module mt0014
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix, store_value, retrieve_value
function test()
    p = "matrices/afile.h5"
    mkpath(p)
    f = joinpath(p, "afile.h5")

    d = rand(Int64, 6, 3)
    store_matrix(f, "/data/d", d)
    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b = rand(Float64, 6, 3)
    store_matrix(f, "/data/b", b)
    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    store_value(f, "name", "afile.h5")
    v = retrieve_value(f, "name")
    @test v == "afile.h5"

    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    true
end
end
using .mt0014
mt0014.test()

cleanup()
module mt0015
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix, store_value, retrieve_value
function test()

    f = joinpath(".", "datafile.h5")

    d = rand(Int64, 6, 3)
    store_matrix(f, "/data/d", d)
    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b = rand(Float64, 6, 3)
    store_matrix(f, "/data/b", b)
    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    store_value(f, "name", "afile.h5")
    v = retrieve_value(f, "name")
    @test v == "afile.h5"

    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    b1 = retrieve_matrix(f, "/data/b")
    @test b == b1

    true
end
end
using .mt0015
mt0015.test()

cleanup()
cleanup()
module mt0016
using Test
using SparseArrays
using DataDrop: clean_file_name, with_extension
using DataDrop: store_matrix, retrieve_matrix, store_value, retrieve_value, empty_hdf5_file
function test()

    f = joinpath(".", "datafile.h5")

    d = rand(Int64, 6, 3)
    store_matrix(f, "/data/d", d)
    d1 = retrieve_matrix(f, "/data/d")
    @test d == d1

    empty_hdf5_file(f)

    @test_throws KeyError d1 = retrieve_matrix(f, "/data/d")

    true
end
end
using .mt0016
mt0016.test()


text:
using Test
using TestSetExtensions
using LinearAlgebra
using Random
using Qaintessent
using Random

##==----------------------------------------------------------------------------------------------------------------------

"""Checks that tailored apply gives same result as general apply"""

@testset ExtendedTestSet "apply gates to state vector" begin
    N = 6
    θ = 0.7π
    ϕ = 0.4π
    n = randn(3); n /= norm(n)
    ψ = rand(ComplexF64, 2^N)

    @testset "apply basic gates" begin
        # single qubit gates
        for g in [X, Y, Z, HadamardGate(), SGate(), TGate(), RxGate(θ), RyGate(θ), RzGate(θ), RotationGate(θ, n), PhaseShiftGate(ϕ)]
            i = rand(1:N)
            cg = circuit_gate(i, g)
            cga = CircuitGate{1,AbstractGate}(cg.iwire, cg.gate) # generate same gate with type AbstractGate{1}

            @test apply(ψ, cg) ≈ apply(ψ, cga)
            @test apply(ψ, cg) ≈ sparse_matrix(cga, N) * ψ
        end
    end

    @testset "apply multiple basic gates" begin
        g = [X, Y, Z, HadamardGate(), SGate(), TGate(), RxGate(θ), RyGate(θ), RzGate(θ), RotationGate(θ, n), PhaseShiftGate(ϕ)]
        i = randperm(N)[1:3]
        j = randperm(N)[1:3]
        cgs = circuit_gate.(j, g[i])
        cga = CircuitGate[]
        for cg in cgs
            push!(cga, CircuitGate{1,AbstractGate}(cg.iwire, cg.gate))
        end

        @test apply(ψ, cgs) ≈ apply(ψ, cga)
        @test apply(ψ, cgs) ≈ sparse_matrix(cga, N) * ψ
    end

    @testset "apply basic gates exceptions" begin
        cg = CircuitGate[]
        @test_throws ErrorException("Vector of length 0 cannot be applied") apply(ψ, cg)
    end

    @testset "apply moments" begin
        # single qubit gates
        for g in [X, Y, Z, HadamardGate(), SGate(), TGate(), RxGate(θ), RyGate(θ), RzGate(θ), RotationGate(θ, n), PhaseShiftGate(ϕ)]
            i = rand(1:N)
            cg = Moment([circuit_gate(i, g)])
            cga = Moment([CircuitGate{1,AbstractGate}(cg[1].iwire, cg[1].gate)]) # generate same gate with type AbstractGate{1}

            @test apply(ψ, cg) ≈ apply(ψ, cga)
            @test apply(ψ, cg) ≈ sparse_matrix(cga, N) * ψ
        end
    end

    @testset "apply multiple moments" begin
        g = [X, Y, Z, HadamardGate(), SGate(), TGate(), RxGate(θ), RyGate(θ), RzGate(θ), RotationGate(θ, n), PhaseShiftGate(ϕ)]
        i = randperm(N)[1:3]
        j = randperm(N)[1:3]
        cgs = circuit_gate.(j, g[i])
        cga = CircuitGate[]
        for cg in cgs
            push!(cga, CircuitGate{1,AbstractGate}(cg.iwire, cg.gate))
        end

        @test apply(ψ, Moment(cgs)) ≈ apply(ψ, Moment(cga))
        @test apply(ψ, Moment(cgs)) ≈ sparse_matrix(cga, N) * ψ
    end

    @testset "apply moments exceptions" begin
        m = Moment[]
        @test_throws ErrorException("Vector of length 0 cannot be applied") apply(ψ, m)
    end

    @testset "apply basic controlled gates" begin
        # control gate
        for g in [X, Y, Z, RotationGate(θ,n), PhaseShiftGate(ϕ)]
            i = rand(1:N)
            j = rand([1:i-1; i+1:N])
            cg = circuit_gate(i, g, j)
            cga = CircuitGate{2,ControlledGate{typeof(g)}}(cg.iwire, cg.gate)

            @test apply(ψ, cg) ≈ apply(ψ, cga)
            @test apply(ψ, cg) ≈ sparse_matrix(cga, N) * ψ
        end
    end

    @testset "apply swap gate" begin
        i = rand(1:N)
        j = rand([1:i-1; i+1:N])
        cg = circuit_gate((i, j), SwapGate())
        cga = CircuitGate{2,SwapGate}(cg.iwire, cg.gate)

        @test apply(ψ, cg) ≈ apply(ψ, cga)
        @test apply(ψ, cg) ≈ sparse_matrix(cga, N) * ψ
    end

    @testset "apply 1-qubit MatrixGate" begin
        # MatrixGate: one qubit
        d = 2
        A = rand(ComplexF64, d, d)
        U, R = qr(A)
        U = Array(U);
        g = MatrixGate(U)
        i = rand(1:N)
        cg = circuit_gate(i, g)
        cga = CircuitGate{1,MatrixGate}(cg.iwire, cg.gate) # generate same gate with type AbstractGate{1}

        @test apply(ψ, cg) ≈ apply(ψ, cga)
        @test apply(ψ, cg) ≈ sparse_matrix(cga, N) * ψ
    end

    @testset "apply k-qubit MatrixGate" begin
        # MatrixGate: k qubits
        k = rand(1:N)
        A = rand(ComplexF64, 2^k, 2^k)
        U, R = qr(A)
        U = Array(U);
        g = MatrixGate(U)
        iwire = [rand(1:N)]
        for j in 1:k-1
            l = rand(1:N-j)
            i = setdiff([1:N...], iwire)[l]
            push!(iwire, i)
        end
        sort!(iwire)
        cga = CircuitGate{k,MatrixGate}((iwire...,), g)
        m = sparse_matrix(cga, N)
        @test apply(ψ, cga) ≈ m*ψ
    end
end


text:
struct OneDTMazeRoundRobin <: LearningUpdate end

update!(self::OneDTMazeRoundRobin, args...) = nothing
Base.size(learner::OneDTMazeRoundRobin) = (0,0)

function get_action_probs(self::OneDTMazeRoundRobin, state, observation, args...)

    # action_probs = zeros(4)
    # UP = 1
    # RIGHT = 2
    # DOWN = 3
    # LEFT = 4
    # LEFT_JUNCTION = 3
    # MIDDLE_JUNCTION = 8
    # RIGHT_JUNCTION = 19

    # #This relies on being tabular. For continuous TMaze, need different round robin method
    # state = if typeof(state) == Int
    #     state
    # else
    #     #It is a sparse array encoding
    #     tups = findnz(state)
    #     tups[1][1]
    # end

    # if state == LEFT_JUNCTION
    #     action_probs[DOWN] = 0.5
    #     action_probs[UP] = 0.5
    # elseif state == MIDDLE_JUNCTION
    #     action_probs[LEFT] = 0.5
    #     action_probs[RIGHT] = 0.5
    # elseif state == RIGHT_JUNCTION
    #     action_probs[DOWN] = 0.5
    #     action_probs[UP] = 0.5
    # elseif state < 3
    #     action_probs[UP] = 1
    # elseif state < 5
    #     action_probs[DOWN] = 1
    # elseif state < 8
    #     action_probs[LEFT] = 1
    # elseif state < 15
    #     action_probs[UP] = 1
    # elseif state < 17
    #     action_probs[RIGHT] = 1
    # elseif state < RIGHT_JUNCTION
    #     action_probs[UP] = 1
    # else
    #     action_probs[DOWN] = 1
    # end

    # return action_probs
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
using NESSie.Format

testfiles = ((mktemp()..., 0), # empty file
             (mktemp()..., 2)) # dummy file
write(testfiles[2][2], """
REMARK   1 Lorem ipsum dolor sit amet.
REMARK   1
ATOM      1  N   THR     1     -17.108  25.866  23.850  0.1812 1.8240
HETATM    2  O   HOH     2     -21.160  40.444  40.509 -0.8340 1.6612
ATOM      3  CA  THR     3     -16.775  27.193  23.310  0.0034 1.9080
ATOM      4  CA  VAL     4      23.595  65.525  13.234 -0.0000 2.0000
TER
END""")

@testset "readpqr" begin
    try
        for (fname, fh, len) in testfiles, T in testtypes
            seekstart(fh)
            charges = readpqr(fh, T)

            @test typeof(charges) == Vector{Charge{T}}
            @test length(charges) == len
            if fname == testfiles[2][1]
                @test charges[1].pos == T[-17.108, 25.866, 23.850]
                @test charges[1].val == T(0.1812)
                @test charges[2].pos == T[-16.775, 27.193, 23.310]
                @test charges[2].val == T(0.0034)
            end
        end
    finally
        [begin close(fh); rm(fn); end for (fn, fh, _) in testfiles]
    end
end


text:
using ApproxFunRational, AbstractIterativeSolvers, Plots

### Scalar case
tol = 1e-14#1.e-4
r = 1.
s0 = 1.0/sqrt(2)
pn = 1.
d = 4.0
L = 1.

Pss = z -> 2*r*s0^2.0/(r^2+z.^2)
#Pss = z -> exp(-z^2/2)
G = SumFun(Fun(zai(Pss),OscLaurent(0.0,L),100) + Fun(OscLaurent(0.0,L),[pn]))
H = SumFun(Fun(zai(Pss),OscLaurent(-d,L),100))
𝓒 = Cauchy(-1)
𝓢 = Cauchy(1)
One = SumFun(pad(Fun(1.,OscLaurent(0.0,L)),100))

sop = x -> x - (G-One)*(𝓒*x)

out = GMRES(sop,H,⋅,10*tol,30, x -> x)
u = sum([out[2][i]*out[1][i] for i=1:length(out[2])])
jump = x -> (𝓢*u)(x) - (𝓒*u)(x)*G(x) - H(x)
jump(.11)

𝓕 = FourierTransform(-1.0)
U = -(𝓕*(𝓒*u))

x = -10:.01:10
y1 = real(map(U,x))
plot(x,y1)

hatH = FourierTransform(-1.0)*H
x = -10:.01:10
y1 = real(map(hatH,x))
y2 = imag(map(hatH,x))
plot(x,y1)
plot!(x,y2)

k = .3
CFun = Fun(x -> u(x)*exp(1im*k*x), -25..25,400)
sum(CFun)

sum(map(x -> x(k),components(U)))
U.space

U(.3)


text:
mutable struct Arf <: Real
    arf::arf_struct
    prec::Int

    function Arf(;prec::Integer=DEFAULT_PRECISION[])
        res = new(arf_struct(0,0,0,0), prec)
        init!(res)
        finalizer(clear!, res)
        return res
    end

    function Arf(si::Int64; prec::Integer=DEFAULT_PRECISION[])
        res = new(arf_struct(0,0,0,0), prec)
        ccall(@libarb(arf_init_set_si), Cvoid, (Ref{Arf}, Int), res, si)
        finalizer(clear!, res)
        return res
    end

    # ccall(@libarb(arf_init_set_ui), Cvoid, (Ref{Arf}, Int), res, ui)
end

mutable struct Arb <: Real
    arb::arb_struct
    prec::Int

    function Arb(;prec::Integer=DEFAULT_PRECISION[])
        res = new(arb_struct(0,0,0,0, 0,0), prec)
        init!(res)
        finalizer(clear!, res)
        return res
    end
end

mutable struct Acb <: Number
    acb::acb_struct
    prec::Int

    function Acb(;prec::Integer=DEFAULT_PRECISION[])
        res = new(arb_struct(0,0,0,0, 0,0,0,0), prec)
        init!(res)
        finalizer(clear!, res)
        return res
    end
end

mutable struct Mag <: Real
    mag::mag_struct

    function Mag()
        res = new(mag_struct(0,0))
        init!(res)
        finalizer(clear!, res)
        return res
    end

    function Mag(m::mag_struct)
        res = new(mag_struct(0,0))
        ccall(
            @libarb(mag_init_set),
            Cvoid,
            (Ref{Mag}, Ref{mag_struct}),
            res,
            m,
        )
        finalizer(clear!, res)
        return res
    end

    function Mag(arf::Arf)
        res = new(mag_struct(0,0))
        ccall(
            @libarb(mag_init_set_arf),
            Cvoid,
            (Ref{Mag}, Ref{Arf}),
            res,
            arf,
        )
        finalizer(clear!, res)
        return res
    end
end

for (T, prefix) in ((Arf, :arf), (Arb, :arb), (Acb, :acb), (Mag, :mag))
    arbstruct = Symbol(prefix, :_struct)
    spref = "$prefix"
    @eval begin
        init!(a::$T) = init!(a.$prefix)
        clear!(a::$T) = clear!(a.$prefix)
        cprefix(::Type{$T}) = Symbol($spref) # useful for metaprogramming
        cstruct(t::$T) = getfield(t, cprefix($T))
    end
    T == Mag && continue
    @eval begin
        Base.precision(x::$T) = x.prec
    end
end


text:
using ImageMorphology
using ImageCore
using Test
using OffsetArrays

using Documenter
doctest(ImageMorphology, manual = false)

@testset "ImageMorphology" begin
    include("convexhull.jl")
    include("connected.jl")
    include("dilation_and_erosion.jl")
    include("thinning.jl")
    include("imfill.jl")
    include("maxtree.jl")
    include("bwdist.jl")
    include("clearborder.jl")
end


text:
module CompatHelper

import Base64
import Dates
import GitHub
import HTTP
import Pkg
import Printf
import TOML
import TimeZones
import UUIDs

include("types.jl")

include("main.jl")

include("assert.jl")
include("ci_service.jl")
include("envdict.jl")
include("get_latest_version_from_registries.jl")
include("get_project_deps.jl")
include("git.jl")
include("new_versions.jl")
include("pull_requests.jl")
include("ssh_keys.jl")
include("stdlib.jl")
include("timestamps.jl")
include("utils.jl")
include("version_numbers.jl")

end # module


text:
using AdFem
using Test
using PyPlot
using PyCall
using LinearAlgebra

np = pyimport("numpy")
m = 40
n = 20
h = 0.1
bdnode = Int64[]
bdval1 = Float64[]
bdval2 = Float64[]
X = (0:m)*h
Y = (0:n)*h
X, Y = np.meshgrid(X,Y)
for i = 1:m+1
    for j = 1:n+1
        push!(bdnode, i+(j-1)*(n+1))
        x = (i-1)*h
        y = (j-1)*h
        push!(bdval1, x^2+y^2)
        push!(bdval2, x^2-y^2)
    end
end
bdnode = [bdnode;bdnode.+(m+1)*(n+1)]
bdval = [bdval1;bdval2]

Xv = []
Yv = []
for j = 1:n
    for i = 1:m
        push!(Xv, (i-1/2)*h)
        push!(Yv, (j-1/2)*h)
    end
end

include("Core.jl")


text:
# Unit tests for implementations/nbodysimulator/nbodysimulator_result.jl

@testset "nbodysimulator_result.jl" begin
    # TODO
end


text:
using DataFrames

module QuantileRegression

    import StatsModels: DataFrameRegressionModel, Formula, coef, @formula
    import Base.LinAlg.BLAS.axpy!
    export qreg, coef, vcov, stderr, quantiles, IP, IRLS, @formula

    using DataFrames, Distributions, Base.LinAlg.BLAS, StatsModels, StatsBase

    mutable struct QRegModel
        beta::Vector{Float64}
        vcov::Matrix{Float64}
        stderr::Vector{Float64}
        q
    end

    abstract type Solver end

    immutable IP <: Solver
    end

    immutable IRLS <: Solver
        tol
        maxIter
        threshold
    end
    IRLS(;tol::Real = 1e-12, maxIter::Integer = 1_000, threshold::Real = 1e-5) = IRLS(tol, maxIter, threshold)

    include("InteriorPoint.jl")
    include("IRLS.jl")
    include("Covariance.jl")

    function qreg(f::Formula, df::AbstractDataFrame, q, s::Solver = IP())
        mf = ModelFrame(f, df)
        mm = ModelMatrix(mf)
        mr = model_response(mf)
        coef = qreg_coef(mr, mm.m, q, s)
        vcov = qreg_vcov(mr, mm.m, coef, q)
        stderr = sqrt.(diag(vcov))
        return DataFrameRegressionModel(QRegModel(coef, vcov, stderr, q), mf, mm)
    end
    qreg(f::Formula, df::AbstractDataFrame, s::Solver) = qreg(f, df, 0.5, s)

    coef(x::QRegModel) = x.beta
    coef(rm::DataFrameRegressionModel) = coef(rm.model)

    vcov(x::QRegModel) = x.vcov

    stderr(x::QRegModel) = x.stderr

    quantiles(x::QRegModel) = x.q
    quantiles(rm::DataFrameRegressionModel) = quantiles(rm.model)

    function StatsBase.coeftable(mm::QRegModel)
        cc = coef(mm)
        se = stderr(mm)
        tt = cc./se
        CoefTable(hcat(kron(mm.q,ones(length(se))), cc,se,tt),
                 ["Quantile",  "Estimate","Std.Error","t value"],
                 ["x$i" for i = 1:length(cc)])
    end
end


text:
Gtk.waitforsignal(self::Widget, signal) = Gtk.waitforsignal(gwidget(self), signal)
"""
        @waitfor widget.event

Block the current task until the `event` is triggered for `widget`.

## Examples:

```julia
window = Window("Hello, world", 800, 600)

... # Some UI stuff

if !isinteractive()
    # If the code is not running on a interactive julia session
    # (e.g. the REPL), then the line below will block the current
    # task until the window is destroyed
    @waitfor window.destroy
end
```
"""
macro waitfor(expr)
    if @capture(expr, widget_.event_)
        eventname = QuoteNode(event)
        return :( Gtk.waitforsignal($( esc(widget) ), $(eventname)) )
    end
end

"""
        @on widget.event() do args... body... end

Adds a `event` callback to the corresponding `widget`

## Usage

```julia
@on window.destroy() do args...
    # do something when the window is destroyed
end

@on button.clicked() do args...
    # do something when the button is clicked
end

```
"""
macro on(expr)
    if @capture(expr, widget_.event_() do args__ body__ end)
        eventname = QuoteNode(event)
        return esc(
            quote
                onevent($(eventname), $(widget)) do $(args...)
                    $(body...)
                end
            end
        )
    elseif @capture(expr, widget_.event_)
        eventname = QuoteNode(event)
        return esc(
            quote
                begin
                    obs = Observable{Widget}($(widget))
                    onevent($(eventname), $(widget)) do args...
                        obs[] = $(widget)
                    end
                    obs
                end
            end
        )
    end
end

function gkey(key::SymString)
    _key = Symbol("GDK_KEY_$key")
    if isdefined(Gtk.GConstants, _key)
        return getfield(Gtk.GConstants, _key)
    else
        @error "Unknown key '$key'"
        return nothing
    end
end

macro key_str(str)
    Expr(:call, :gkey, esc(str))
end

keyname(event::Gtk.GdkEventKey) = keyval_name(event.keyval)

function setprop!(self::Widget, prop::SymString, callback::Function)
    prop = string(prop)
    if startswith(prop, "on")
        eventname = string(prop[3:end])
        onevent(callback, eventname, self)
    else
        set_gtk_property!(gwidget(self), prop, callback)
    end
end

function onkeypress(callback::Function, widget::Widget)
    if hasmethod(callback, Tuple{<:Widget, Gtk.GdkEventKey})
        onevent(callback, "key-press-event", widget)
    else
        onevent("key-press-event", widget) do widget, event
            callback(event)
        end
    end
end

function onkeypress(callback::Function, key::Integer, widget::Widget)
    onkeypress(widget) do widget, event
        if key == event.keyval
            callback(widget, event)
        end
    end
end

function ondraw(callback::Function, canvas::GtkCanvas)
    @guarded draw(canvas) do c
        w, h = size(c)
        drawing = Drawing(w, h, :image)
        drawing.cr = Gtk.getgc(c)
        callback(canvas)
        finish()
    end
end

ondraw(callback::Function, canvas::Canvas) = ondraw(callback, gwidget(canvas))

function GtkTickCallback(::Ptr{GObject}, ::Ptr{GObject}, ptr::Ptr{Nothing})
    canvas = unsafe_load(convert(Ptr{GtkCanvas}, ptr))
    Gtk.draw(canvas, false)
    return true
end

function onupdate(callback::Function, canvas::GtkCanvas)
    ondraw(callback, canvas)
    ptr = @cfunction(GtkTickCallback, Bool, (Ptr{GObject}, Ptr{GObject}, Ptr{Cvoid}))
    ref = Ref(canvas)
    return GC.@preserve ptr begin
        gtk_widget_add_tick_callback(canvas, ptr, ref, C_NULL)
    end
end

function onupdate(callback::Function, canvas::Canvas; fps::Ref{Float64} = Ref(60.0))
    then = Ref{Float64}(time_ns())
    canvas.tickcb = onupdate(gwidget(canvas)) do canvas
        now = time_ns()
        dt = (now - then[]) / 10e8
        interval = inv(fps[])
        if dt > interval
            callback(dt, canvas)
            then[] = now - (dt % interval)
        end
    end
end

"""
        onmousedown(callback, canvas)

The `mouse-down` event is fired when the user `press` with the mouse on the canvas.
"""
function onmousedown(callback::Function, canvas::Canvas)
    on(canvas.mouse.mousedown) do event
        @secure callback(event) "'onmousedown' event callback triggered a exception"
    end
end
"""
        onmouseup(callback, canvas)

The `mouse-up` event is fired when the user `release` the mouse on the canvas.
"""
function onmouseup(callback::Function, canvas::Canvas)
    on(canvas.mouse.mouseup) do event
        @secure callback(event) "'onmouseup' event callback triggered a exception"
    end
end
"""
        onmousemove(callback, canvas)

The `mouse-move` event is fired when the user `moves` the cursor on the canvas.
"""
function onmousemove(callback::Function, canvas::Canvas)
    on(canvas.mouse.mousemove) do event
        @secure callback(event) "'onmousemove' event callback triggered a exception"
    end
end
"""
        onmousedrag(callback, canvas)

The `mouse-drag` event is fired when the user holds `down` and `move` the cursor on the canvas.
"""
function onmousedrag(callback::Function, canvas::Canvas)
    on(canvas.mouse.mousedrag) do event
        @secure callback(event) "'onmousedrag' event callback triggered a exception"
    end
end
"""
        onscroll(callback, canvas)

The `scroll` event is fired when the user `scrolls` with the mouse-wheel on the canvas.
"""
function onscroll(callback::Function, canvas::Canvas)
    on(canvas.mouse.scroll) do event
        @secure callback(event) "'onscroll' event callback triggered a exception"
    end
end

function onresize(callback::Function, widget::GtkWidget)
    return onevent(callback, "size-allocate", widget)
end

onresize(callback::Function, widget::Widget) = onresize(callback, gwidget(widget))


text:
"Node is the abstract type that will contain every node of the expression tree."
abstract type Node end

"""
FunctionNode represents a non-terminal node of the expression tree.

# Arguments
- `_func::Function`: the actual function that will be executed for that function node.
- `_arity::Integer`: the number of parameters for the function, also, the number
    of children of the function node..
- `_returnType::Union{DataType, Union}`: type of the returning value of a function.
- `_argTypes::Array{Union{DataType, Union}}`: types of the arguments of the function.
"""
struct FunctionNode <: Node
    _func::Function
    _arity::Integer
    _returnType::Union{DataType, Union}
    _argTypes::Array{Union{DataType, Union}}

    function FunctionNode(func::Function, arity::Integer,
                          returnType::Union{DataType, Union},
                          argTypes::Array{Union{DataType, Union}})

        new(func, arity, returnType, argTypes)
    end # function

    function FunctionNode(func::Function, arity::Integer)
        new(func, arity, Any, Array{Union{DataType, Union}}(undef, 0))
    end # function
end # struct

"""
    getFunctionName(node::FunctionNode)

Gets the identifier of a given [`FunctionNode`](@ref).
"""
getFunctionName(node::FunctionNode) = string(node._func)
# function

"""
    getArity(node::FunctionNode)

Gets the arity of a given [`FunctionNode`](@ref).
"""
getFunctionArity(node::FunctionNode) = node._arity
# function

"""
    getReturnType(node::FunctionNode)

Gets the return type of a [`FunctionNode`](@ref).
"""
getReturnType(node::FunctionNode) = node._returnType
# function

"TerminalNode is an abstract type and represents a terminal node of the expression tree."
abstract type TerminalNode <: Node end

"
VariableNode represents an identifier with a value assigned by the user.
- `_name::String`: identifier of the variable.
- `_type::DataType`: type of the variable.

See also: [`TerminalNode`](@ref)
"
struct VariableNode <: TerminalNode
    _name::String
    _type::DataType
end # struct

"""
    getVariableName(node::VariableNode)

Gets the identifier of a [`VariableNode`](@ref).
"""
getVariableName(node::VariableNode) = node._name
# function

"""
    getVariableType(node::VariableNode)

Gets the type of a [`VariableNode`](@ref).
"""
getValueType(node::VariableNode) = node._type
# function

"
ConstantNode represents a single value.

# Fields
- `_value`: value of the constant.
- `_ephemeralFunction::Function`: if set, it determines that the constant is
    ephemeral, and this function, when evaluated, gives the constant its value.
- `_varArgs::Array{Any}`: arguments of the function of the ephemeral constant, if any.

See also: [`TerminalNode`](@ref)
"
mutable struct ConstantNode <: TerminalNode
    _value
    _ephemeralFunction::Function
    _varArgs::Array{Any}

    function ConstantNode(value)
        new(value, noFunc, [])
    end # function

    ConstantNode(ephemeralFunction::Function, varArgs::Array{Any}) = new(NaN, ephemeralFunction, varArgs)
    # function
end # struct

"""
    getConstantType(node::ConstantNode)

Determines wether a [`ConstantNode`](@ref) is ephemeral or not.
"""
function isEphemeralConstant(constantNode::ConstantNode)
    !(constantNode._ephemeralFunction == noFunc)
end # function

"""
    getConstantType(node::ConstantNode)

Gets the type of a [`ConstantNode`](@ref).
"""
getValueType(node::ConstantNode) = typeof(node._value)
# function

"""
    setEphemeralConstant(node::ConstantNode)

Sets the value of an ephemeral [`ConstantNode`](@ref).
"""
function setEphemeralConstant(node::ConstantNode)
    node._value = node._ephemeralFunction(node._varArgs...)
end # function

"""
    eval(node::ConstantNode)

Evaluates a [`ConstantNode`](@ref). If it is an ephemeral constant, sets the
value for that constant.
"""
function eval(node::ConstantNode)
    if isnan(node._value)
        setEphemeralConstant(node)
    end
    node._value
end
# function

"
NoArgsFunctionNode represents a function node that has 0 arity (no arguments)
and it's considered as a terminal node in the expression tree.

# Fields
- `_func::Function`: the actual function that will be executed for that terminal node.

See also: [`TerminalNode`](@ref)
"
struct NoArgsFunctionNode <: TerminalNode
    _func::Function
end # struct

"""
    getNoArgsFunctionName(node::NoArgsFunctionNode)

Gets the identifier of a [`NoArgsFunctionNode`](@ref).
"""
getNoArgsFunctionName(node::NoArgsFunctionNode) = string(node._func)
# function

"""
    getNoArgsFunctionType(node::NoArgsFunctionNode)

Obtains the return type of a [`NoArgsFunctionNode`](@ref).
"""
getReturnType(node::NoArgsFunctionNode) = typeof(node._func())
# function

"""
    eval(node::NoArgsFunctionNode)

Evaluates a [`NoArgsFunctionNode`](@ref).
"""
eval(node::NoArgsFunctionNode) = node._func()
# function

"""
    getSymbol(node::Node)

Gets the identifier of the node.
"""
function getSymbol(node::Node)
    if typeof(node) <: ConstantNode
        node._value
    else
        node._name
    end
end # function

"""
    getType(node::Node)

Gets the type of the node.
"""
function getType(node::Node)
    if typeof(node) <: FunctionNode || typeof(node) <: NoArgsFunctionNode
        getReturnType(node)
    else
        getValueType(node)
    end
end # function

"""
    getArity(node::Node)

Gets the arity of a node.
"""
function getArity(node::Node)
    return typeof(node) <: FunctionNode ? getFunctionArity(node) : 0
end # function

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
# # Global Climatology
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/global_ocean_circulation.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/global_ocean_circulation.ipynb)
#
# Advect particles with climatological monthly mean flow at selected depth level
# (e.g. `k=10` for 95 m) from a global ocean state estimate ([ECCO v4 r2](https://eccov4.readthedocs.io/en/latest/) ; see also <https://ecco-group.org>)
# which is here repeated for `ny` years. For additional documentation e.g. see :
# [1](https://JuliaClimate.github.io/MeshArrays.jl/dev/),
# [2](https://JuliaClimate.github.io/IndividualDisplacements.jl/dev/),
# [3](https://docs.juliadiffeq.org/latest/solvers/ode_solve.html),
# [4](https://en.wikipedia.org/wiki/Displacement_(vector))
#
# [![simulated particle movie (5m)](https://user-images.githubusercontent.com/20276764/84766999-b801ad80-af9f-11ea-922a-610ad8a257dc.png)](https://youtu.be/W5DNqJG9jt0)

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ## 1. Get Software & Iput Files
#
# - packages + helper functions
# - grid and velocity files

using IndividualDisplacements, MeshArrays, OrdinaryDiffEq
using Statistics, DataFrames, MITgcmTools, OceanStateEstimation

include(joinpath(dirname(pathof(IndividualDisplacements)),"../examples/helper_functions.jl"))
IndividualDisplacements.get_ecco_velocity_if_needed();

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ## 2. Set Up Parameters & Inputs
#
# - select vertical level & duration in years
# - read grid variables & velocities
# - normalize velocities

𝑃=setup_global_ocean(k=1,ny=2);

keys(𝑃)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 3. Main Computation Loop
#
# ### 3.1 Initialize Individuals & Solution
#
# - initial particle positions randomly over Global Ocean

xy = init_global_randn(1000,𝑃)
xy = permutedims([xy[:,i] for i in 1:size(xy,2)])

I=(position=xy,velocity=dxy_dt!,
   postprocessing=postprocess_lonlat,parameters=𝑃)
𝐼=Individuals(I)

fieldnames(typeof(𝐼))

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# - initial integration from time 0 to 0.5 month

𝑇=(0.0,𝐼.𝑃.𝑇[2])
∫!(𝐼,𝑇)

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### 3.2 Iteration function example
#
# - `𝐼.𝑃.🔄(𝐼.𝑃,t_ϵ)` resets the velocity input streams to bracket t_ϵ=𝐼.𝑃.𝑇[2]+eps(𝐼.𝑃.𝑇[2])
# - `reset_lonlat!(𝐼)` randomly selects a fraction (defined in `setup_global_ocean()`) of the particles and resets their positions before each integration period. This can maintain homogeneous coverage of the Global Ocean by particles.
# - `∫!(𝐼,𝑇)` then solves for the individual trajectories over one month, after updating velocity fields (𝐼.u0 etc) if needed, and adds diagnostics to the DataFrame used to record / trace variables along the trajectory (𝐼.tr).

function step!(𝐼::Individuals)
    t_ϵ=𝐼.𝑃.𝑇[2]+eps(𝐼.𝑃.𝑇[2])
    𝐼.𝑃.🔄(𝐼.𝑃,t_ϵ)
    reset_lonlat!(𝐼)
    𝑇=Tuple(𝐼.𝑃.𝑇)
    ∫!(𝐼,𝑇)
end

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 3.3 Iterate For `ny*12` Months
#

[step!(𝐼) for y=1:1, m=1:1]

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 3.4 Compute summary statistics
#
# See [DataFrames.jl](https://juliadata.github.io/DataFrames.jl/latest/) documentation for detail and additinal functionalities.

gdf = groupby(𝐼.🔴, :ID)
sgdf= combine(gdf,nrow,:lat => mean)
sgdf[rand(1:size(sgdf,1),4),:]

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 4. Plot trajectories / individual positions
#
# ```
# using Plots
# p=plot(;xlims=(-180,180),ylims=(-90,90),legend=:none)
# p!(x,y)=scatter!(p,x,y,markersize=1.1,markerstrokewidth=0)
# [p!(gdf[i].lon,gdf[i].lat) for i in rand(collect(1:length(gdf)),10)]
# display(p)
# ```

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# Or select a background map (e.g. `lon`, `lat`, and `DL=log10(bottom depth)`)
# and a recipe to superimpose initial and final locations. Try:
#
#```
#include(joinpath(dirname(pathof(IndividualDisplacements)),"../examples/recipes_plots.jl"))
#plot_end_points(𝐼,𝐼.𝑃.Γ)
#```


text:
module AllTools
# WARNING: This module is currently useless. You need to figure out how Julia's registry works, as well as version number and using tagbot to auto publish Releases.

using Reexport
@reexport using FileTools, DataFrameTools, HypertextTools, Shorthands
# FIXME: Currently add AllTools (https://github.com/CGRG-lab/AllTools.jl.git) failed:
# ERROR: Unsatisfiable requirements detected for package FileTools [2c042965]:
#  FileTools [2c042965] log:
#  ├─FileTools [2c042965] has no known versions!
#  └─restricted to versions * by AllTools [363e5ecb] — no versions left
#    └─AllTools [363e5ecb] log:
#      ├─possible versions are: 0.1.0 or uninstalled
#      └─AllTools [363e5ecb] is fixed to version 0.1.0
#
# For the error above, see this thread: https://discourse.julialang.org/t/package-manager-resolve-complaining-of-unsatisfiable-requirements-due-to-no-known-versions/23778

# TODO: learn travis CI at: https://medium.com/starbugs/travis-ci-%E7%B0%A1%E5%96%AE%E4%BA%8B%E6%83%85%E5%B0%B1%E4%BA%A4%E7%B5%A6%E9%9B%BB%E8%85%A6%E5%8E%BB%E5%81%9A%E4%B9%8Bci-cd-%E5%88%9D%E9%AB%94%E9%A9%97-%E8%AE%93-github-pages-%E8%87%AA%E5%8B%95%E6%9B%B4%E6%96%B0-7647be30eb1c and https://youtu.be/QVmU29rCjaA?t=1041
# TODO: also see this for learning Registry: https://www.juliabloggers.com/tips-and-tricks-to-register-your-first-julia-package/ and https://discourse.julialang.org/t/creating-a-custom-registry/28007/17

end


text:
"""
    isconstant(profile)

Return `true` if the time profile `profile` is intervalwise constant.
"""
isconstant(::TimeProfile) = false
isconstant(::PGSE) = true
isconstant(::DoublePGSE) = true


text:
export sparsePCE

# TODO: Overfitting: strategy for overfitting
# TODO: multiple dispatch for bases on arbitrary weight functions

# Univariate adaptive basis algo
# Arguments:
#   * op: Function hanlde for polnomial basis
#   * modelFun: function handle for model to evaluate
#   * pceFun: function handle to compute PCE coefficients
#   * pceModel: struct storing all relevant pce information
#   * maxDegree: maximum allowed degree of basis
#   * ϵ: target error
function adaptiveBasis!(pceModel, opFun::Function, pceFun::Function; maxDegree::Int = 20, ϵ = 1e-10)
    error = Inf     # error of current PC expansion
    error_opt = Inf    # error for optimal solution
    deg = 1         # current degree
    badFit = false  # overfitting flag

    # Compare to preset threshold. Stop and return lowest error PCE, if condition is met.
    while error > ϵ && deg < maxDegree
        # 1. Compute current orthogonal basis
        # FUTURE: truncation scheme
        op_current = opFun(deg)

        # 2. Compute new PCE coefficients and error for current basis
        # computePCE(PCEmodel, true)
        modelFun = pceModel.modelFun
        pce, error = pceFun(op_current, modelFun)

        # 3. Check if result increases, store best, check for overfitting
        println("Current error: $error \t minimal error: $error_opt")
        if error < error_opt
            # Update optimal degree, error and pce coefficients
            error_opt = error
            pceModel.op = op_current
            pceModel.pceCoeffs = pce
            pceModel.error = error
            badFit = false
        elseif badFit
            # No increase in 2 iterations -> stop
            println("break")
            break
        else
            badFit = true
        end

        deg = deg + 1
        println(deg)
    end
end

# Struct, containing all relevant data for OLS regression for PCE estimation
mutable struct OLSModel{T<:AbstractOrthoPoly}
    op::T   # Chosen orthogonal basis FUTURE: Sparse Ortho Poly
    modelFun::Function  # Function handle for system model
    pceCoeffs::Array{Float64,1}
    error::Float64      # Error of PCE model
    # Y::Vector{Float64}  # Model evaluaions
    # Φ::Matrix{Float64}  # Moment matrix
    function OLSModel(op::AbstractOrthoPoly, modelFun::Function)
        new{typeof(op)}(op, modelFun, [], Inf)
    end
end

# Compute the PCE coefficients using OLS regression
# TODO: sample function
# TODO: determine good number of samples (-> maybe move sampling outside)
# function computePCE(olsModel::OLSModel, compError)
function computePceOls(op::AbstractOrthoPoly, modelFun::Function; compError=true)
    deg = op.deg
    # Draw n samples from the underlying measure
    nSample = deg * 20
    X = sampleMeasure(nSample, op)
    # Build moment matrix Φ for nSaopmple samples and deg basis polynomials
    Φ = [ evaluate(j, X[i], op) for i = 1:nSample, j = 0:deg]
    # Evaluate the model
    Y = modelFun.(X)
    # Perform OLS regression
    coefficients = leastSquares(Φ, Y)

    if compError
        error = looError(Y, Φ, coefficients)
    else
        error = nothing
    end

    return coefficients, error
end

# Compute the PCE coefficients using a Galerkin projection
function computePceProjection(op::AbstractOrthoPoly, modelFun::Function; compError=true)
end

#################################
# Running the code
#################################

# Define the model
μ = 0
σ = 1
# model(x) = exp(μ + σ * x)
# model(x) = x^7 - 21 * x^5 + 105 * x^3 - 105 * x + x^4 - 6 * x^2 + 3
# model(x) = x^4 - 3* x^2
# model(x) = x^10 - 45 * x^8 + 630 * x^6 - 3150 * x^4 + 4725 * x^2 - 945
model(x) = 0

# Function handle for basis generation
op(deg) = GaussOrthoPoly(deg, Nrec = deg * 2)

# Function handle for computing pce components with ols regression
pceFun(op, model) = computePceOls(op, model)

# Initialize OLSModel struct
op1 = op(1)
olsModel = OLSModel(op1, model)

# Compute adaptive basis with a regression approach
# adaptiveBasis(olsModel, op, pceFun)

# Compute a sparse basis for the provided pce model and parameters
# Parameters:
# * sampleFun: user provided sampling method
#   * op: The full candidate orthogonal basis
function sparsePCE(op::AbstractOrthoPoly, modelFun::Function, sampleFun::Function; Q²tgt = .999, pMax = 10, jMax = 1)
    # Parameter specification
    pMax = min(op.deg, pMax)    # pMax can be at most size of given full basis
    COND = 1e4                  # Maximum allowed matrix condition number (see Blatman2010)
    α = 0.001                   # Tuning parameter (see Blatman2010)
    ϵ = α * (1 - Q²tgt)         # Error threshold of coefficients
    sampleSize = pMax * 20      # TODO: How to determine?

    # Initilaize basis and determination coefficients
    p = 0       # current degree
    Ap = [0]    # set of basis indices for current degree
    R² = 0.0    # coefficient of determination
    Q² = 0.0    # leave-one out coefficient determination
    pce = []    # pce coefficients

    #TODO: for each y compute sparse pce and collect output
    # restore properties after each loop
    # in case of not reaching Q²_tgt do ???
    restart = true
    # Outer loop: Iterate on experimental design
    while restart
        restart = false

        # 1. Build current ED, compute Y
        X = sampleFun(sampleSize)
        Y = modelFun.(X) # TODO: Reuse old ED data, this part is very expensive!

        # 2. Initialize loop
        p = 0
        Ap = [0] # Zero element
        Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap ]
        # Φ = [ evaluate(j, X, op) for j in Ap ] # TODO does this also work?
        pce = leastSquares(Φ, Y)
        R² = 1 - empError(Y, Φ, pce)
        Q² = 1 - looError(Y, Φ, pce)
        println("Intialization:")
        println("R²: $R²")
        println("Q²: $Q²")

        # Main loop: Iterate max degree p
        while Q² ≤ Q²tgt && p ≤ pMax && !restart
            p += 1  # current max degree
            j = 0   # number of allowed interactions

            # Iterate number of interactions j
            jM = min(p, jMax) # p limits #interactions
            while Q² ≤ Q²tgt && j < jM && !restart
                j += 1
                J = [] # Temporary store potential new basis elements
                candidates = [p] # FUTURE: this needs to capture all current multi indices. ALso need to change from indices to objects (?)

                # Forward step: compute R² for all candidate terms and keep the relevant ones
                for a in candidates
                    A = Ap ∪ a
                    Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in A]
                    pce = leastSquares(Φ, Y)
                    R²new = 1 - empError(Y, Φ, pce) # Only need R² error here "due to more efficiency" (Blatman 2010)
                    println("R² (p=$p, j=$j): ", R²)

                    # Compute accuracy gain. If high enough, add cadidate polynomial
                    ΔR² = R²new - R²
                    # println("ΔR²: $ΔR² \t ϵ: $ϵ")
                    if ΔR² > ϵ
                        J = J ∪ a
                    end
                end
                println("Candidates after forward step: $J")

                # FUTURE: Sort Jp and ΦJ according to ΔR² <- why?
                # J = sort(J)
                # R = []

                # Conditioning Check: If resulting enriched basis does not yield a well-conditioned moments matrix, we have to restart
                Ap_new = Ap ∪ J
                Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap_new]
                # Check Φ
                if cond(Φ) > COND
                    # Increase experimental design and restart computations
                    restart = true
                    k = 3 # rescale factor according to Blatman
                    println("Moments matrix is ill-conditioned. Restart computation with new sample size: $(k * sampleSize)) (Old size: $sampleSize)")
                    sampleSize *= k   # TODO: build properly
                    # X = sampleMeasure(sampleSize, op)
                else
                    # If conditioning is okay, update accuracy R² for backward step
                    pce = leastSquares(Φ, Y)
                    R² = 1 - empError(Y, Φ, pce)
                end

                # Backward step: Remove candidate polynomials one by one and compute the effect on Q²
                if !restart
                    println("Ap_new: ", Ap_new)
                    Del = [] # polynomials to be discarded
                    for a in Ap_new
                        A = filter(e->e≠a,Ap_new)
                        println("A = ", A)
                        # New candiadte basis -> compute new determination coefficients
                        Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in A]
                        pce = leastSquares(Φ, Y)
                        R²new = 1 - empError(Y, Φ, pce)

                        # If decrease in accuracy is too small, throw polynomial away
                        ΔR² = R² - R²new
                        if ΔR² ≤ ϵ
                            Del = Del ∪ a
                            println("throw away a = ", a)
                        end
                    end

                    # Update basis and compute errors for next iteration
                    Ap = filter(e->e∉Del, Ap_new)
                    Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap]
                    pce = leastSquares(Φ, Y)
                    println("pce: ", pce)
                    R² = 1 - empError(Y, Φ, pce)
                    Q² = 1 - looError(Y, Φ, pce)
                    println("Q² (p=$p, j=$j): ", Q²)
                end
            end
            println()
        end
    end

    if Q² < Q²tgt
        error("Computation reached max degree $pMax. However, accuracy is below target with Q² = $(Q²).")
    else
        println("Computation reached target accuracy with Q² = $(Q²) and max degree $p")
        return Dict(:coeffs=>pce,
                    :Ap=>Ap,
                    :p=>p,
                    :R²=>R²,
                    :Q²=>Q²)
    end
end

# sparsePCE with default sampling method
function sparsePCE(op::AbstractOrthoPoly, modelFun::Function; Q²tgt = .999, pMax = 10, jMax = 1)
    sampleFun(sampleSize) = sampleMeasure(sampleSize, op)
    return sparsePCE(op, modelFun, sampleFun; Q²tgt, pMax, jMax)
end

function testFun(modelFun, deg, range; sampleSize = deg *2)
    op = GaussOrthoPoly(deg)

    X = sampleMeasure(sampleSize, op)
    println("X: ", X)
    Y = modelFun.(X) # This is the most expensive part
    println("Y: ", Y)

    # Comparison to full basis with max degree 4
    # Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in 0:deg]
    Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in range]
    println("Φ:", Φ)
    pce = leastSquares(Φ, Y)
    println("pce: ", pce)
    R²0 = 1 - empError(Y, Φ, pce)
    println("R²0: ", R²0)
    Q²0 = 1 - looError(Y, Φ, pce)
    println("Q²0: ", Q²0)
    ###
end

# SPARSE ALGO FOR VECTOR-VLAUED MODEL RESPONSES
# Parameters:
# * sampleFun: user provided sampling method
#   * op: The full candidate orthogonal basis
function sparsePCE_vec(op::AbstractOrthoPoly, modelFun::Function, sampleFun::Function; Q²tgt = .999, pMax = 10, jMax = 1)
    # Parameter specification
    pMax = min(op.deg, pMax)    # pMax can be at most size of given full basis
    COND = 1e4                  # Maximum allowed matrix condition number (see Blatman2010)
    α = 0.001                   # Tuning parameter (see Blatman2010)
    ϵ = α * (1 - Q²tgt)         # Error threshold of coefficients
    sampleSize = pMax * 20      # Initial sample size for ED TODO: How to determine?

    # Main loop
    restart = true
    while restart
        restart = false

        # 1. Build initial ED, compute Y
        X = sampleFun(sampleSize)
        Y = modelFun.(X) # This is the most expensive part

        # Outer loop: Iterate over every ED vector.
        #TODO: sotre output for each y
        # restore properties after each loop
        # in case of not reaching Q²_tgt do ???
        # bad conditioning
        while !isempty(Y) && !restart
            y = popfirst!(Y)

            # Initilaize basis and determination coefficients TODO: wehere?
            p = 0       # current degree
            Ap = [0]    # set of basis indices for current degree
            R² = 0.0    # coefficient of determination
            Q² = 0.0    # leave-one out coefficient determination
            pce = []    # current pce coefficients

            # 2. Initialize loop
            p = 0
            Ap = [0] # Zero element
            Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap ]
            # Φ = [ evaluate(j, X, op) for j in Ap ] # TODO does this also work?
            pce = leastSquares(Φ, y)
            R² = 1 - empError(y, Φ, pce)
            Q² = 1 - looError(y, Φ, pce)
            println("Intialization:")
            println("R²: $R²")
            println("Q²: $Q²")

            # Main loop: Iterate max degree p
            while Q² ≤ Q²tgt && p ≤ pMax && !restart
                p += 1  # current max degree
                j = 0   # number of allowed interactions

                # Iterate number of interactions j
                jM = min(p, jMax) # p limits #interactions
                while Q² ≤ Q²tgt && j < jM && !restart
                    j += 1
                    J = [] # Temporary store potential new basis elements
                    candidates = [p] # FUTURE: this needs to capture all current multi indices. ALso need to change from indices to objects (?)

                    # Forward step: compute R² for all candidate terms and keep the relevant ones
                    for a in candidates
                        A = Ap ∪ a
                        Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in A]
                        pce = leastSquares(Φ, y)
                        R²new = 1 - empError(y, Φ, pce) # Only need R² error here "due to more efficiency" (Blatman 2010)
                        println("R² (p=$p, j=$j): ", R²)

                        # Compute accuracy gain. If high enough, add cadidate polynomial
                        ΔR² = R²new - R²
                        # println("ΔR²: $ΔR² \t ϵ: $ϵ")
                        if ΔR² > ϵ
                            J = J ∪ a
                        end
                    end
                    println("Candidates after forward step: $J")

                    # FUTURE: Sort Jp and ΦJ according to ΔR² <- why?
                    # J = sort(J)
                    # R = []

                    # Conditioning Check: If resulting enriched basis does not yield a well-conditioned moments matrix, we have to restart
                    Ap_new = Ap ∪ J
                    Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap_new]
                    # Check Φ
                    if cond(Φ) > COND
                        # Increase experimental design and restart computations
                        restart = true
                        k = 3 # rescale factor according to [Blatman2008]
                        println("Moments matrix is ill-conditioned. Restart computation with new sample size: $(k * sampleSize)) (Old size: $sampleSize)")
                        sampleSize *= k   # TODO: build properly
                    else
                        # If conditioning is okay, update accuracy R² for backward step
                        pce = leastSquares(Φ, y)
                        R² = 1 - empError(y, Φ, pce)
                    end

                    # Backward step: Remove candidate polynomials one by one and compute the effect on Q²
                    if !restart
                        println("Ap_new: ", Ap_new)
                        Del = [] # polynomials to be discarded
                        for a in Ap_new
                            A = filter(e->e≠a,Ap_new)
                            println("A = ", A)
                            # New candiadte basis -> compute new determination coefficients
                            Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in A]
                            pce = leastSquares(Φ, y)
                            R²new = 1 - empError(y, Φ, pce)

                            # If decrease in accuracy is too small, throw polynomial away
                            ΔR² = R² - R²new
                            if ΔR² ≤ ϵ
                                Del = Del ∪ a
                                println("throw away a = ", a)
                            end
                        end

                        # Update basis and compute errors for next iteration
                        Ap = filter(e->e∉Del, Ap_new)
                        Φ = [ evaluate(j, X[i], op) for i = 1:sampleSize, j in Ap]
                        pce = leastSquares(Φ, y)
                        println("pce: ", pce)
                        R² = 1 - empError(y, Φ, pce)
                        Q² = 1 - looError(y, Φ, pce)
                        println("Q² (p=$p, j=$j): ", Q²)
                    end
                end
                println()
            end
        end
    end

    if Q² < Q²tgt
        error("Computation reached max degree $pMax. However, accuracy is below target with Q² = $(Q²).")
    else
        println("Computation reached target accuracy with Q² = $(Q²) and max degree $p")
        return Dict(:coeffs=>pce,
                    :Ap=>Ap,
                    :p=>p,
                    :R²=>R²,
                    :Q²=>Q²)
    end
end

#################################
# Projection model old
#################################

# function projectionModel(model::Function, op::AbstractOrthoPoly, maxDeg::Int)
#     t2 = Tensor(2, op)

#     coefficients = Array{Float64}(undef, maxDeg+1)

#     for k = 0:maxDeg
#         # Function for evaluation of model and basis polynomial at given quadrature nodes t
#         g(t) = model(t) .* evaluate(k, t, op)
#         # Perfom numerical integration for the inner product <x, ϕ_k>
#         proj = integrate(g, op.quad)
#         # Inner product <ϕ_k,ϕ_k> for Galerkin projection
#         norm = t2.get([k,k])
#         # Coefficient is the fraction of the two inner products
#         coefficients[k+1] = proj / norm
#     end

#     return coefficients
# end

#################################
# Regression model old
#################################

# # maybe shift this part to regression file??
# struct OLSModel{O, S, C}
#     op::AbstractOrthoPoly   # Chosen orthogonal basis
#     sysFun::Function        # System equations of the model as function handle
#     coefficients::Array{Float64,1}
#     error::Float64
#     # TODO: sample function
#     # TODO: (number of samples)

#     # inner constructor
#     function OLSModel(op::AbstractOrthoPoly, sysFun::Function, errorFun::Function; compError = true)
#         deg = op.deg
#         # Draw n samples from the underlying measure
#         nSample = deg * 20
#         X = sampleMeasure(nSample, op)
#         # Build moment matrix Φ for nSaopmple samples and deg basis polynomials
#         Φ = [ evaluate(j, X[i], op) for i = 1:nSample, j = 0:deg]
#         # Evaluate the model
#         Y = sysFun.(X)
#         # Perform OLS regression
#         coefficients = leastSquares(Φ, Y)

#         # compute the leave-one-out error on the result
#         compError ? error = errorFun(Y, Φ, coefficients) : error = nothing

#         new{typeof(op), typeof(sysFun), typeof(coefficients)}(op, sysFun, coefficients, error)
#     end
# end

#################################
# Adaptive algo old
#################################

# Univariate adaptive basis algo
# Arguments:
#   * maxDegree: maximum degree searched for
#   * op: polnomial basis of PCE model
#   * model: function handle for model to evaluate
#   * method: name of method to compute PCe coefficients
# function adaptiveBasis(method::String, op::AbstractOrthoPoly, model; maxDegree::Int = 10, Nrec::Int = maxDegree+1, ϵ = 1e-10)
# function adaptiveBasis(PCEmodel, op; maxDegree::Int = 20, Nrec::Int = maxDegree+1, ϵ = 1e-10)
#     error = Inf     # error of current PC expansion
#     error_opt = Inf    # error for optimal solution
#     deg = 1         # current degree
#     deg_opt = 1     # degree of optimal solution
#     op_opt = 0      # basis of best solution
#     PCE_opt = []    # PCE coefficients of optimal solution
#     badFit = false  # overfitting flag

#     # Compare to preset threshold. Stop and return lowest error PCE, if condition is met.
#     while error > ϵ && deg < maxDegree

#         # 1. Compute current orthogonal basis
#         op_current = op(deg)

#         # 2. Compute new PCE model for current basis
#         currentModel = PCEmodel(op_current)
#         # Get conmputed PCE coefficients from model
#         PCE  = currentModel.coefficients

#         # 3. Compute generalization error for convergence
#         error = currentModel.error

#         # 4. Check if result increases, store best
#         println("Current error: $error \t minimal error: $error_opt")
#         if error < error_opt
#             # Update optimal degree, error and pce coefficients
#             deg_opt = deg
#             error_opt = error
#             op_opt = op_current
#             PCE_opt = PCE
#             badFit = false
#         elseif badFit
#             # No increase in 2 iterations -> stop
#             println("break")
#             break
#         else
#             badFit = true
#         end

#         deg = deg + 1
#         println(deg)
#     end

#     return PCE_opt, op_opt, error_opt, deg_opt
# end

max_stars_repo_name: adriangrupp/sparsePolyChaos

max_issues_repo_name: adriangrupp/sparsePolyChaos

max_forks_repo_name: adriangrupp/sparsePolyChaos