using Printf: @printf, @sprintf
using Statistics: mean, std

####################
## main
####################
const N = 50 # the number of vector
x = 100(rand(N)) #
x_mean = mean(x)
x_std = std(x)
ind = sortperm(x)[end:-1:end-4]
####################

####################
## plot
####################

# directory where figures are printed
figdir="./fig"
if !isdir(figdir); mkdir(figdir); end

using Plots
pyplot()
#gr()
#plotlyjs()

# bar
bar(x, xlims=(0.0, Float64(N+1)), ylims=(0.0,105.), legend=false, tickfont=font(12,"sans-serif"),
    title="N=$(N), mean="*@sprintf("%.2f",x_mean)*", σ="*@sprintf("%0.2f",x_std), # GRバックエンドの場合はσが正しく表示されない
    size=(800,600))

# horizontal line
xm = [0.0,Float64(N+1)]
ym = [x_mean, x_mean]
plot!(xm, ym, line=:solid, c=:green, lab="mean")

# 1st - 5th large
scatter!(ind, x[ind], ms=8., color=:red)

# save
savefig(joinpath(figdir,"rand_bar.png"))

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


"""
    RotationError{T<:Real, D<:ErrorMap} <: StaticVector{3,T}

A three-parameter rotation error, converted to/from a `QuatRotation` using the
`ErrorMap` `D`.

# Usage
A `RotationError` is typically created using one of the following methods

    rotation_error(R1::Rotation, R2::Rotation, error_map::ErrorMap)
    R1 ⊖ R2

which compute the difference between the rotations `R1` and `R2` and convert the
result to a 3D rotation error using `error_map`.

The error can be "added" back to a rotation using the inverse operation:

    add_error(R1::Rotation, e::RotationError)
    R1::Rotation ⊕ e::RotationError
"""
struct RotationError{T,D} <: StaticVector{3,T}
    err::SVector{3,T}
    map::D
    @inline function RotationError(err::SVector{3,T}, map::D) where {T,D <: ErrorMap}
        new{T,D}(err, map)
    end
end

@inline QuatRotation(e::RotationError)::Rotation = e.map(e.err)

"""
    rotation_error(R1::Rotation, R2::Rotation, error_map::ErrorMap)

Compute the `RotationError` by calculating the "difference" between `R1` and `R2`, i.e.
`R2\\R1`, then mapped to a three-parameter error using `error_map`.

Can be equivalently called using the default map with `R1 ⊖ R2`

If `error_map::IdentityMap`, then `SVector(R1\\R2)::SVector{3}` is used as the error. Note
this only works for three-parameter rotation representations such as `RodriguesParam` or `MRP`.
"""
function rotation_error(R1::Rotation, R2::Rotation, error_map::ErrorMap)
    return RotationError(inv(error_map)(R2\R1), error_map)
end

function rotation_error(R1::Rotation, R2::Rotation, error_map::IdentityMap)
    err = params(R2\R1)
    if length(err) != 3
        throw(ArgumentError("R2\\R1 must be a three-dimensional parameterization, got $(length(err))"))
    end
    return RotationError(err, error_map)
end

#   set the default error map to the CayleyMap
@inline ⊖(R1::Rotation, R2::Rotation) = rotation_error(R1, R2, CayleyMap())

#   default to the identity map for Rodrigues Params and MRPs
@inline ⊖(R1::RodriguesParam, R2::RodriguesParam) = rotation_error(R1, R2, IdentityMap())
@inline ⊖(R1::MRP,            R2::MRP           ) = rotation_error(R1, R2, IdentityMap())

# Static Arrays interface
function (::Type{E})(t::NTuple{3}) where E <: RotationError
    E(t[1], t[2], t[3])
end
Base.@propagate_inbounds Base.getindex(e::RotationError, i::Int) = e.err[i]
@inline Base.Tuple(e::RotationError) = Tuple(e.err)

# Compose a rotation with an error
"""
    add_error(R1::Rotation, e::RotationError)

"Adds" the rotation error `e` to the rotation `R1` by converting `e` to a quaternion via
its `ErrorMap` and then composing with `R1`.

Equivalent to

    R1 ⊕ e
"""
function add_error(R1::Rotation, e::RotationError)
    R1 * QuatRotation(e)
end

function add_error(R1::R, e::RotationError{<:Any, IdentityMap}) where R <: Rotation
    # must be able to construct R from a SVector{3}
    if length(params(R1)) != 3
        throw(ArgumentError("Can't construct a rotation of type $R from three parameters"))
    end
    R1 * R(e.err)
end

@inline ⊕(R1::Rotation, e::RotationError) = add_error(R1, e)


"""
    Labor demand by A firms.
"""
function labor_demand_A(k_I, par)
    @unpack_Pars_Simple par
    ((γ * α_hat * v * C^α_C * k_I^δ_hat) / w)^(1 / (1 - α_hat))
end

"""
    Labor demand by B firms.
"""
function labor_demand_B(k_I, par)
    @unpack_Pars_Simple par
    (((1 - γ) * α_hat * v * C^α_C * k_I^δ_hat) / w)^(1 / (1 - α_hat))
end

"""
    How much capital A firms have in place given their production and selling decision.
"""
total_capital_A(k_I_b, k_I_P, par) = (1 - par.ϕ) * k_I_b + k_I_P

"""
    How much capital B firms have in place given their production and selling decision.
"""
total_capital_B(k_I_b, k_I_P) = k_I_b + k_I_P

"""
    Market clearing condition for traded capital. For any unit shipped, only 1-τ units arrive. Needs to be scaled by sector size (γ for A and 1-γ for B).
"""
market_clearing(k_I_b_B, k_I_b_A, par) = (1 - par.γ) * k_I_b_B + (1 - par.τ) * par.γ * k_I_b_A

"""
    Marginal cost of capital production by A firms.
"""
mc_A(k_I, par) = par.ca_A * par.cb_A * k_I^(par.cb_A - 1)

"""
    Marginal cost of capital production by B firms.
"""
mc_B(k_I, par) = par.ca_B * par.cb_B * k_I^(par.cb_B - 1)

"""
    Cost of capital production by A firms.
"""
c_A(k_I, par) = par.ca_A * k_I^par.cb_A
"""
    Cost of capital production by B firms.
"""
c_B(k_I, par) = par.ca_B * k_I^par.cb_B

"""
    How much capital is produced by A firms if they are indifferent between
    producing one more unit or selling one unit less.

        AND

    The firm does not sell all of its capital! Otherwise producing one more unit has a different payoff (sell and keep the rest instead of keeping it fully).
"""
k_I_prod_A(par) = ((par.p_I / (1 - par.ϕ)) / (par.ca_A * par.cb_A))^(1 / (par.cb_A - 1))

"""
    How much capital is produced by B firms if they are indifferent between
    producing one more unit or selling one unit less.

        AND

    The firm does not sell all of its capital! Otherwise producing one more unit has a different payoff (sell and keep the rest instead of keeping it fully).
"""
k_I_prod_B(par) = ((par.p_I / (1 - par.τ)) / (par.ca_B * par.cb_B))^(1 / (par.cb_B - 1))

"""
    The price for the B-good. C_A is the quantity of A-goods.
"""
function price_A(C_A, par::Pars_Simple)
    @unpack_Pars_Simple par
    γ * v * C^α_C * C_A^(-1 / σ)
end

"""
    The price for the B-good. C_B is the quantity of B-goods.
"""
function price_B(C_B, par::Pars_Simple)
    @unpack_Pars_Simple par
    (1 - γ) * v * C^α_C * C_B^(-1 / σ)
end

"""
    The wage rate that leads to market clearing (l_A + l_B = 1)
"""
function wage_rate(k_I_A, k_I_B; par)
    @unpack_Pars_Simple par
    c = (α_hat * v * C^α_C)
    c * ((γ * k_I_A^δ_hat)^(1 / (1 - α_hat)) + ((1 - γ) * k_I_B^δ_hat)^(1 / (1 - α_hat)))^(1 - α_hat)
end

"""
    Profit for both type of firms.
"""
function profit(par)
    @unpack_Pars_Simple par
    k_I_b_A, k_I_b_B = solve_firm_problems(par)

    if par.ϕ > 0 && loss_A(k_I_b_A, par) > 1e-8
        println("")
        k_I_A, l_A = capital_and_labor_con_A(k_I_b_A, par)
    else
        k_I_A, l_A = capital_and_labor_unc_A(k_I_b_A, par)
    end
    k_I_B, l_B = capital_and_labor_B(k_I_b_B, par)

    k_I_P_A = k_I_A - (1 - ϕ) * k_I_b_A # Sold capital is only partially lost.
    k_I_P_B = k_I_B - k_I_b_B

    profit_A = γ * v * C^α_C * l_A^α_hat * k_I_A^δ_hat - w * l_A - p_I * k_I_b_A - c_A(k_I_P_A, par)
    profit_B = (1 - γ) * v * C^α_C * l_B^α_hat * k_I_B^δ_hat - w * l_B - p_I / (1 - τ) * k_I_b_B - c_B(k_I_P_B, par)

    return profit_A, profit_B
end

# function labor_demand(k_I_j, par::Pars_Simple)
#     @unpack_Pars_Simple par
#     (α_hat * γ * C^α_C * (k_I_j)^δ_hat / w)^(1 / (1 - α_hat))
# end


using StaticArrays

include("partialbridge!.jl")

T = 2.0
dt = 1/1000

tt = 0.:dt:T
#struct PBIntegratedDiffusion <: ContinuousTimeProcess{ℝ{2}}
#    γ::Float64
#end

Bridge.b(t::Float64, x, P::PBIntegratedDiffusion) = ℝ{2}(x[2], βu(t, x[2], P))

@test Bridge.b!(0.0, X1.yy[1], copy( X1.yy[1]),  P) == Bridge.b(0.0, X1.yy[1], P)

Bridge.σ(t, x, P::PBIntegratedDiffusion) = ℝ{2}(0.0, P.γ)

@test Bridge.σ!(0.0, X1.yy[1], 1.0, copy( X1.yy[1]),  P) == Bridge.σ(0.0, X1.yy[1], P)

Bridge.constdiff(::PBIntegratedDiffusion) = true

#struct PBIntegratedDiffusionAux <: ContinuousTimeProcess{ℝ{2}}
#    γ::Float64
#end

Bridge.b(t::Float64, x, P::PBIntegratedDiffusionAux) = ℝ{2}(x[2], βu(t, x[2], P))

@test Bridge.b!(0.0, X1.yy[1], copy( X1.yy[1]),  Pt) == Bridge.b(0.0, X1.yy[1], Pt)

Bridge.σ(t, P::PBIntegratedDiffusionAux) =  ℝ{2}(0.0, P.γ)
Bridge.σ(t, x, P::PBIntegratedDiffusionAux) = Bridge.σ(t, P)

Bridge.B(t, P::PBIntegratedDiffusionAux) = @SMatrix [0.0 1.0; 0.0 -1.0]
Bridge.β(t, P::PBIntegratedDiffusionAux) = ℝ{2}(0, 1/2)
Bridge.a(t, P::PBIntegratedDiffusionAux) = @SMatrix [0.0 0.0; 0.0 P.γ^2]

# Generate Data
Random.seed!(1)

W2 = sample(tt, Wiener())

@test W2.yy == W1.yy

X2 = solve(Euler(), ℝ{2}(x0), W2, P)

@test X2.yy == X1.yy

L = @SMatrix [1. 0.]
Σ = @SMatrix [0.01]
v = ℝ{1}(2.5)

ϵ = 0.02
# Solve Backward Recursion

S2 = typeof(L)
S = typeof(L*L')
T = typeof(diag(L*L'))

N = length(tt)
Lt = zeros(S2, N)
Ht = zeros(S, N)
νt = zeros(T, N)

Po2 = Bridge.PartialBridgeνH(tt, P, Pt, L, v, ϵ, Σ)

Xo2 = copy(X2)
solve!(Euler(), Xo2, ℝ{2}(x0), W2, Po2)

@test norm(Xo1.yy - Xo2.yy) < 500*eps()

# Likelihood

ll2 = llikelihood(Bridge.LeftRule(), Xo2, Po2)

@test ll1 ≈ ll2

lls2 = Float64[]

    # MCMC parameter
function mcmc2(x0, tt, Po2)
    iterations = 1000
    subsamples = 0:100:iterations
    ρ = 0.9

    # initalization
    W = sample(tt, Wiener())
    X = solve(Euler(), x0, W, Po2)
    ll = llikelihood(Bridge.LeftRule(), X, Po2)

    acc = 0

    Wo = copy(W)
    Xo = copy(X)
    Wrho = copy(W)

    XX = Any[]
    if 0 in subsamples
        push!(XX, copy(X))
    end
    for iter in 1:iterations
        # Proposal
        sample!(Wrho, Wiener())
        Wo.yy .= ρ*W.yy + sqrt(1-ρ^2)*Wrho.yy

        solve!(Euler(), Xo, x0, Wo, Po2)
        llo = llikelihood(Bridge.LeftRule(), Xo, Po2)
        push!(lls2, llo)

        if log(rand()) <= llo - ll
            X, Xo = Xo, X
            W, Wo = Wo, W

            ll = llo
            acc += 1
        end
        if iter in subsamples
            push!(XX, copy(X))
        end
    end

end

@time @testset "MCMC" begin
    mcmc2(ℝ{2}(x0), tt, Po2)

    @test norm(lls1 - lls2) < sqrt(eps())

end


"""
    loglikelihood(v::IHTVariable{T, M})

Calculates the loglikelihood of observing `y` given mean `μ = E(y) = g^{-1}(xβ)`
and some univariate distribution `d`.

Note that loglikelihood is the sum of the logpdfs for each observation.
"""
function loglikelihood(v::IHTVariable{T, M}) where {T <: Float, M}
    d = v.d
    y = v.y
    μ = v.μ
    wts = v.cv_wts
    logl = zero(T)
    ϕ = MendelIHT.deviance(v) / length(y) # variance in the case of normal
    @inbounds for i in eachindex(y)
        logl += loglik_obs(d, y[i], μ[i], wts[i], ϕ)
    end
    return logl
end

"""
This function is taken from GLM.jl from:
https://github.com/JuliaStats/GLM.jl/blob/956a64e7df79e80405867238781f24567bd40c78/src/glmtools.jl#L445

`wt`: in GLM.jl, this is prior frequency (a.k.a. case) weights for observations.
We use this for cross validation weighting: if sample `i` is being fitted,
wt[i] = 1. Otherwise wt[i] = 0.
"""
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:
#    (Γ(r+y) / (Γ(r) * y!)) * μ^y * r^r / (μ+r)^{r+y}
# The parameterization of NegativeBinomial(r=r, p) in Distributions.jl is
#    Γ(r+y) / (y! * Γ(r)) * p^r(1-p)^y
# Hence, p = r/(μ+r)
loglik_obs(d::NegativeBinomial, y, μ, wt, ϕ) = wt*logpdf(NegativeBinomial(d.r, d.r/(μ+d.r)), y)

"""
    deviance(d, y, μ, wts)

Calculates the sum of the squared deviance residuals (e.g. (y - μ)^2 for Gaussian case)
Each individual sqared deviance residual is evaluated using `devresid`
which is implemented in GLM.jl
"""
function deviance(d::UnivariateDistribution, y::AbstractVector{T},
    μ::AbstractVector{T}, wts::AbstractVector{T}) where {T <: Float}
    dev = zero(T)
    @inbounds for i in eachindex(y)
        dev += wts[i] * devresid(d, y[i], μ[i])
    end
    return dev
end
deviance(v::IHTVariable{T, M}) where {T <: Float, M} =
    MendelIHT.deviance(v.d, v.y, v.μ, v.cv_wts)

"""
    update_μ!(μ, xb, l)

Update the mean (μ) using the linear predictor `xb` with link `l`.
"""
function update_μ!(μ::AbstractVecOrMat{T}, xb::AbstractVecOrMat{T}, l::Link) where {T <: Float}
    @inbounds for i in eachindex(μ)
        μ[i] = linkinv(l, xb[i])
    end
end

function update_μ!(v::IHTVariable{T, M}) where {T <: Float, M}
    μ = v.μ
    xb = v.xb
    zc = v.zc
    l = v.l
    @inbounds for i in eachindex(μ)
        μ[i] = linkinv(l, xb[i] + zc[i]) #genetic + nongenetic contributions
    end
end

"""
    update_xb!(v::IHTVariable{T, M})

Updates the linear predictors `xb` and `zc` with the new proposed `b` and `c`.
`b` is sparse but `c` (beta for non-genetic covariates) is dense.

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, M}) where {T <: Float, M}
    # compute X*beta for genetic component
    if v.memory_efficient
        fill!(v.xb, 0)
        fill!(v.xk, 0)
        Threads.@threads for j in eachindex(v.idx)
            if v.idx[j]
                t = Threads.threadid()
                @inbounds @simd for i in 1:length(v.xb)
                    v.xk[i, t] += v.x[i, j] * v.b[j]
                end
            end
        end
        sum!(v.xb, v.xk) # v.xb .= sum(v.xk, dims=2)
    else
        copyto!(v.xk, @view(v.x[:, v.idx]))
        copyto!(v.gk, view(v.b, v.idx)) # use v.gk as storage
        mul!(v.xb, v.xk, v.gk)
    end
    # compute X*beta for nongenetic component
    mul!(v.zc, v.z, v.c)
    if !(typeof(v.d) <: Normal)
        clamp!(v.xb, -20, 20)
        clamp!(v.zc, -20, 20)
    end
end

"""
    score!(v::IHTVariable{T})

Calculates the score (gradient) `X^T * W * (y - g(x^T b))` for different GLMs.
W is a diagonal matrix where `w[i, i] = dμ/dη / var(μ)` (see documentation)
"""
function score!(v::IHTVariable{T, M}) where {T <: Float, M}
    d, l, x, z, y, cv_wts = v.d, v.l, v.x, v.z, v.y, v.cv_wts
    @inbounds for i in eachindex(y)
        η = v.xb[i] + v.zc[i]
        w = mueta(l, η) / glmvar(d, v.μ[i])
        v.r[i] = w * (y[i] - v.μ[i]) * cv_wts[i] # cv_wts handles sample masking for cross validation
    end
    mul!(v.df, Transpose(x), v.r)
    mul!(v.df2, Transpose(z), v.r)
end

"""
Wrapper function to decide whether to use Newton or MM algorithm for estimating
the nuisance paramter of negative binomial regression.
"""
function mle_for_r(v::IHTVariable{T, M}) where {T <: Float, M}
    method = v.est_r
    if method == :MM
        return update_r_MM(v)
    elseif method == :Newton
        return update_r_newton(v)
    else
        throw(ArgumentError("Only support method is Newton or MM, but got $method"))
    end

    return nothing
end

"""
Performs maximum loglikelihood estimation of the nuisance paramter for negative
binomial model using MM's algorithm.
"""
function update_r_MM(v::IHTVariable{T, M}) where {T <: Float, M}
    y = v.y
    μ = v.μ
    r = v.d.r # estimated r in previous iteration
    num = zero(T)
    den = zero(T)
    for i in eachindex(y)
        for j = 0:y[i] - 1
            num = num + (r /(r + j))  # numerator for r
        end
        p = r / (r + μ[i])
        den = den + log(p)  # denominator for r
    end

    return NegativeBinomial(-num / den, T(0.5))
end

"""
Performs maximum loglikelihood estimation of the nuisance paramter for negative
binomial model using Newton's algorithm. Will run a maximum of `maxIter` and
convergence is defaulted to `convTol`.
"""
function update_r_newton(v::IHTVariable{T, M};
    maxIter=100, convTol=T(1.e-6)) where {T <: Float, M}
    y = v.y
    μ = v.μ
    r = v.d.r # estimated r in previous iteration

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

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

    function negbin_loglikelihood(r::T)
        v.d = NegativeBinomial(r, T(0.5))
        return MendelIHT.loglikelihood(v)
    end

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

    new_r    = one(T)
    stepsize = one(T)
    for i in 1:maxIter

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

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

        #check convergence
        if abs(r - new_r) <= convTol
            return NegativeBinomial(new_r, T(0.5))
        else
            r = new_r
        end
    end

    return NegativeBinomial(r, T(0.5))
end

"""
This function computes the gradient step v.b = P_k(β + η∇f(β)) and updates idx and idc.
"""
function _iht_gradstep!(v::IHTVariable{T, M}, η::T) where {T <: Float, M}
    lg = length(v.group)
    J = v.J
    k = length(v.ks) > 0 ? v.ks : v.k

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

    # project to sparsity, scaling model by weight vector, if supplied
    if lg == 0
        vectorize!(v.full_b, v.b, v.c, v.weight, v.zkeep)
        project_k!(v.full_b, k + v.zkeepn)
        unvectorize!(v.full_b, v.b, v.c, v.weight, v.zkeep)
    else
        # TODO: enable model selection for non-genetic covariates on group projection
        project_group_sparse!(v.b, v.group, J, k)
    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
    typeof(k) == Int && _choose!(v)

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

"""
    vectorize!(a::AbstractVector, b::AbstractVector, c::AbstractVector, ckeep::BitVector)

Without allocations, copies `b` into `a` and the copies `c` into
remaining parts of `a`.

`ckeep` tracks entries of `c` that will not be projected. That is, entries of `a`
corresponding to entries in `c[i]` will be filled with `Inf` if `ckeep[i] = true`.
"""
function vectorize!(a::AbstractVector, b::AbstractVector, c::AbstractVector,
    weight::AbstractVector, ckeep::BitVector)
    lb = length(b)
    lw = length(weight)
    la = length(a)

    # scale model by weight vector, if supplied
    if lw == 0
        copyto!(@view(a[1:lb]), b)
        copyto!(@view(a[lb+1:la]), c)
    else
        @inbounds for i in 1:lb
            a[i] = b[i] * weight[i]
        end
        ii = 1
        @inbounds for i in lb+1:la
            a[i] = c[ii] * weight[i]
            ii += 1
        end
    end

    # don't project certain non-genetic covariates
    a_view = @view(a[lb+1:la])
    @view(a_view[ckeep]) .= typemax(eltype(a))
end

"""
    unvectorize!(a::AbstractVector, B::AbstractMatrix, C::AbstractMatrix, Ckeep::BitVector)

Without allocations, copies the first `length(b)` part of `a` into `b`
and the remaining parts of `a` into `c`.

`ckeep` tracks entries of `c` that will not be projected. That is, entries in `c[i]`
will not be touched if `ckeep[i] = true`.
"""
function unvectorize!(a::AbstractVector, b::AbstractVector, c::AbstractVector,
    weight::AbstractVector, ckeep::BitVector)
    lb = length(b)
    lw = length(weight)
    la = length(a)

    # scale model by weight vector, if supplied
    if lw == 0
        copyto!(b, @view(a[1:lb]))
        ii = 1
        @inbounds for i in lb+1:la
            if !ckeep[ii]
                c[ii] = a[i]
            end
            ii += 1
        end
    else
        @inbounds for i in 1:lb
            b[i] = a[i] / weight[i]
        end
        ii = 1
        @inbounds for i in lb+1:la
            if !ckeep[ii]
                c[ii] = a[i] / weight[i]
            end
            ii += 1
        end
    end
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. If `init_beta=true`, then beta values will be initialized to
their univariate values (see [`initialize_beta`](@ref)), in which case we will simply
choose top `k` entries

`J` is the maximum number of active groups, and `k` is the maximum number of
predictors per group.
"""
function init_iht_indices!(v::IHTVariable, init_beta::Bool, cv_idx::BitVector, verbose::Bool=false)
    fill!(v.b, 0)
    fill!(v.b0, 0)
    fill!(v.best_b, 0)
    fill!(v.xb, 0)
    fill!(v.xk, 0)
    fill!(v.gk, 0)
    fill!(v.xgk, 0)
    fill!(v.idx, false)
    fill!(v.idx0, false)
    copyto!(v.idc, v.zkeep)
    copyto!(v.idc0, v.zkeep)
    fill!(v.r, 0)
    fill!(v.df, 0)
    fill!(v.df2, 0)
    fill!(v.c, 0)
    fill!(v.best_c, 0)
    fill!(v.c0, 0)
    fill!(v.zc, 0)
    fill!(v.zdf2, 0)
    fill!(v.μ, 0)
    fill!(v.cv_wts, 0)
    fill!(v.full_b, 0)
    v.cv_wts[cv_idx] .= 1

    init_beta && !(typeof(v.d) <: Normal) &&
        throw(ArgumentError("Intializing beta values only work for Gaussian phenotypes! Sorry!"))

    # find the intercept by Newton's method
    ybar = zero(eltype(v.y))
    @inbounds @simd for i in eachindex(v.y)
        ybar += v.y[i] * v.cv_wts[i]
    end
    ybar /= count(!iszero, v.cv_wts)
    for iteration = 1:20
        g1 = linkinv(v.l, v.c[1])
        g2 = mueta(v.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
    mul!(v.zc, v.z, v.c)

    # update mean vector and use them to compute score (gradient)
    update_μ!(v)
    score!(v)

    if init_beta
        initialize_beta!(v, cv_idx, verbose)
        project_k!(v)
    else
        # first `k` non-zero entries are chosen based on largest gradient
        vectorize!(v.full_b, v.df, v.df2, v.weight, v.zkeep)
        if length(v.ks) == 0 # no group projection
            project_k!(v.full_b, v.k + v.zkeepn) # project k + number of nongentic covariates to keep
            unvectorize!(v.full_b, v.df, v.df2, v.weight, v.zkeep)
            v.idx .= v.df .!= 0
            v.idc .= v.zkeep

            # Choose randomly if more are selected
            _choose!(v)
        else
            project_group_sparse!(v.df, v.group, v.J, v.ks)
            v.idx .= v.b .!= 0
            fill!(v.idc, true)
        end
    end

    # make necessary resizing when necessary
    check_covariate_supp!(v)

    # store relevant components of x for first iteration
    v.memory_efficient || copyto!(v.xk, @view(v.x[:, v.idx]))
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}) where {T <: Float}
    sparsity = v.k + v.zkeepn
    groups = (v.J == 0 ? 1 : v.J)

    nonzero = sum(v.idx) + sum(v.idc) - v.zkeepn
    if nonzero > groups * sparsity
        z = zero(eltype(v.b))
        non_zero_idx = findall(!iszero, v.idx)
        excess = nonzero - groups * sparsity
        for pos in sample(non_zero_idx, excess, replace=false)
            v.b[pos]   = z
            v.idx[pos] = false
        end
    end
end

"""
In `_init_iht_indices` and `_iht_gradstep!`, if non-genetic cov got
included/excluded, we must resize `xk` and `gk`.

TODO: Use ElasticArrays.jl
"""
function check_covariate_supp!(v::IHTVariable{T}) where {T <: Float}
    nzidx = sum(v.idx)
    if nzidx != size(v.xk, 2) && !v.memory_efficient
        v.xk = zeros(T, size(v.xk, 1), nzidx)
        v.gk = zeros(T, nzidx)
    end
    # @inbounds for i in eachindex(v.zkeep)
    #     v.zkeep[i] && !v.idc[i] && error("A non-genetic covariate was accidentally set to 0! Shouldn't happen!")
    # end
end

"""
    _iht_backtrack_(logl::T, prev_logl::T, η_step::Int64, nstep::Int64)

Returns true if one of the following conditions is met:
1. New loglikelihood is smaller than the old one
2. Current backtrack (`η_step`) exceeds maximum allowed backtracking (`nstep`, default = 3)
"""
function _iht_backtrack_(logl::T, prev_logl::T, η_step::Int64, nstep::Int64) where {T <: Float}
    (prev_logl > logl) && (η_step < nstep)
end

"""
    standardize!(z::AbstractVecOrMat)

Standardizes each column of `z` to mean 0 and variance 1. Make sure you
do not standardize the intercept.
"""
@inline function standardize!(z::AbstractVecOrMat)
    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)
    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, μ)
    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}
    k < 0 && throw(DomainError("Attempted to project to sparsity level $k"))
    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

function project_k!(v::IHTVariable)
    v.k < 0 && throw(DomainError("Attempted to project to sparsity level $(v.k)"))

    # copy genetic and non-genetic effects to full_grad, project, and copy back
    vectorize!(v.full_b, v.b, v.c, v.weight, v.zkeep)
    project_k!(v.full_b, v.k + v.zkeepn)
    unvectorize!(v.full_b, v.b, v.c, v.weight, v.zkeep)

    # update support
    v.idx .= v.b .!= 0
    v.idc .= v.c .!= 0
    check_covariate_supp!(v)
end

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

When `k` is an integer, projects the vector `y` onto the set with at most `J` active groups
and at most `k` active predictors per group. To have variable group sparsity level, input `k`
as a vector of integers. We will preserve `k[1]` elements for group 1, `k[2]` predictors for
group 2...etc. This function assumes there are no unknown or overlaping group membership.

Note: In the `group` vector, the first group must be 1, and the second group must be 2...etc.

# 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

J, k, n = 2, 0.9, 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`: Maximum predictors per group. Can be a positive integer or a vector of integers.
"""
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

function project_group_sparse!(y::AbstractVector{T}, group::AbstractVector{Int64},
    J::Int64, k::Vector{Int}) 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[n]
            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[n])
            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 = SnpArrays.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}, cur_logl::T, best_logl::T) where {T <: Float}
    copyto!(v.b0, v.b)     # b0 = b
    copyto!(v.idx0, v.idx) # idx0 = idx
    copyto!(v.idc0, v.idc) # idc0 = idc
    copyto!(v.c0, v.c)     # c0 = c
    if cur_logl > best_logl
        copyto!(v.best_b, v.b) # best_b = b
        copyto!(v.best_c, v.c) # best_c = c
    end
    return max(cur_logl, best_logl)
end

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

Here v is the score and J is the expected information matrix, which is
computed by J = g'(xb) / var(μ), assuming dispersion is 1. Note

Note: cross validation weights are needed in W.
"""
function iht_stepsize!(v::IHTVariable{T, M}) where {T <: Float, M}
    z = v.z # non genetic covariates
    d = v.d # distribution
    l = v.l # link function

    # first compute xgk = X*v using relevant components of gradient
    if v.memory_efficient
        fill!(v.xgk, 0)
        fill!(v.xk, 0)
        Threads.@threads for j in eachindex(v.idx)
            if v.idx[j]
                t = Threads.threadid()
                @inbounds @simd for i in 1:length(v.xgk)
                    v.xk[i, t] += v.x[i, j] * v.df[j]
                end
            end
        end
        sum!(v.xgk, v.xk) # v.xgk .= sum(v.xk, dims=2)
    else
        copyto!(v.gk, view(v.df, v.idx))
        mul!(v.xgk, v.xk, v.gk)
    end
    mul!(v.zdf2, view(z, :, v.idc), view(v.df2, v.idc))
    v.xgk .+= v.zdf2 # xgk = Xv

    # Compute sqrt(W); use zdf2 as storage
    @inbounds @simd for i in eachindex(v.zdf2)
        v.zdf2[i] = sqrt(abs2(mueta(l, v.xb[i] + v.zc[i])) / glmvar(d, v.μ[i])) * v.cv_wts[i]
    end
    v.xgk .*= v.zdf2 # xgk = sqrt(W)Xv

    # now compute and return step size. Note non-genetic covariates are separated from x
    numer = v.memory_efficient ? sum(abs2, @view(v.df[v.idx])) : sum(abs2, v.gk)
    numer += sum(abs2, @view(v.df2[v.idc]))
    denom = dot(v.xgk, v.xgk)
    η = numer / denom

    # for bad boundary cases (sometimes, k = 1 in cross validation generates weird η)
    isinf(η) && (η = T(1e-8))
    isnan(η) && (η = T(1e-8))

    return η :: T
end

"""
    initialize_beta!(v, cv_wts)

Initialze beta to univariate regression values. That is, `β[i]` is set to the estimated
beta with `y` as response, and `x[:, i]` with an intercept term as covariate.

Progress will be displayed if `verbose=true`

TODO: this function assumes quantitative (Gaussian) phenotypes. Make it work for other distributions
"""
function initialize_beta!(
    v::IHTVariable,
    cv_wts::BitVector, # cross validation weights; 1 = sample is present, 0 = not present
    verbose::Bool=true
    )
    y, x, z, β, c, T = v.y, v.x, v.z, v.b, v.c, eltype(v.b)
    xtx_store = [zeros(T, 2, 2) for _ in 1:Threads.nthreads()]
    xty_store = [zeros(T, 2) for _ in 1:Threads.nthreads()]
    xstore = [zeros(T, sum(cv_wts)) for _ in 1:Threads.nthreads()]
    ystore = y[cv_wts]
    pmeter = verbose ? Progress(nsnps(v) + ncovariates(v), 5,
        "Initializing β to univariate regression values...") : nothing
    # genetic covariates
    c0 = zero(eltype(z))
    Threads.@threads for i in 1:nsnps(v)
        id = Threads.threadid()
        copyto!(xstore[id], @view(x[cv_wts, i]))
        linreg!(xstore[id], ystore, xtx_store[id], xty_store[id])
        c0 += xty_store[id][1]
        β[i] = xty_store[id][2]
        verbose && next!(pmeter) # update progress
    end
    # non-genetic covariates
    Threads.@threads for i in 2:ncovariates(v)
        id = Threads.threadid()
        copyto!(xstore[id], @view(z[cv_wts, i]))
        linreg!(xstore[id], ystore, xtx_store[id], xty_store[id])
        c0 += xty_store[id][1]
        c[i] = xty_store[id][2]
        verbose && next!(pmeter) # update progress
    end
    c[1] = c0 / (nsnps(v) + ncovariates(v) - 1)
    clamp!(v.b, -2, 2)
    clamp!(v.c, -2, 2)
    copyto!(v.b0, β)
    copyto!(v.c0, c)
end

"""
    linreg!(x::Vector, y::Vector)

Performs linear regression with `y` as response, `x` and a vector of 1 as
covariate. `β̂` will be stored in `xty_store`.

Code inspired from Doug Bates on Discourse:
https://discourse.julialang.org/t/efficient-way-of-doing-linear-regression/31232/28
"""
function linreg!(
    x::AbstractVector{T},
    y::AbstractVector{T},
    xtx_store::AbstractMatrix{T} = zeros(T, 2, 2),
    xty_store::AbstractVector{T} = zeros(T, 2)
    ) where {T<:AbstractFloat}
    N = length(x)
    N == length(y) || throw(DimensionMismatch())
    xtx_store[1, 1] = N
    xtx_store[1, 2] = sum(x)
    xtx_store[2, 2] = sum(abs2, x)
    xty_store[1] = sum(y)
    xty_store[2] = dot(x, y)
    try # rare SNPs may have 0s everywhere, causing cholesky to fail
        ldiv!(cholesky!(Symmetric(xtx_store, :U)), xty_store)
    catch
        return xty_store
    end
    return xty_store
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

"""
    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

# small function to check sparsity parameter `k` is reasonable.
function check_group(k, group)
    if typeof(k) <: Vector
        @assert length(group) > 1 "Doubly sparse projection specified (since k" *
            " is a vector) but there are no group information."
        for i in 1:length(k)
            group_member = count(x -> x == i, group)
            group_member > k[i] || throw(DomainError("Maximum predictors for group " *
                "$i was $(k[i]) but there are only $group_member predictors is this " *
                "group. Please choose a smaller number."))
        end
    else
        @assert k >= 0 "Value of k (max predictors per group) must be nonnegative!\n"
    end
end

# helper function from https://discourse.julialang.org/t/how-to-find-out-the-version-of-a-package-from-its-module/37755
pkgversion(m::Module) = Pkg.TOML.parsefile(joinpath(dirname(string(first(methods(m.eval)).file)), "..", "Project.toml"))["version"]

function print_iht_signature(io::IO)
    v = pkgversion(MendelIHT)
    println(io, "****                   MendelIHT Version $v                  ****")
    println(io, "****     Benjamin Chu, Kevin Keys, Chris German, Hua Zhou       ****")
    println(io, "****   Jin Zhou, Eric Sobel, Janet Sinsheimer, Kenneth Lange    ****")
    println(io, "****                                                            ****")
    println(io, "****                 Please cite our paper!                     ****")
    println(io, "****         https://doi.org/10.1093/gigascience/giaa044        ****")
    println(io, "")
    io != stdout && print_iht_signature(stdout)
end
print_iht_signature() = print_iht_signature(stdout)

function print_parameters(io::IO, k, d, l, use_maf, group, debias, tol, max_iter, min_iter)
    regression = typeof(d) <: Normal ? "linear" : typeof(d) <: Bernoulli ?
        "logistic" : typeof(d) <: Poisson ? "Poisson" :
        typeof(d) <: NegativeBinomial ? "NegativeBinomial" :
        typeof(d) <: MvNormal ? "Multivariate Gaussian" : "unknown"
    println(io, "Running sparse $regression regression")
    println(io, "Number of threads = ", Threads.nthreads())
    println(io, "Link functin = $l")
    typeof(k) <: Int && println(io, "Sparsity parameter (k) = $k")
    typeof(k) <: Vector{Int} && println(io, "Sparsity parameter (k) = using group membership specified in k")
    println(io, "Prior weight scaling = ", use_maf ? "on" : "off")
    println(io, "Doubly sparse projection = ", length(group) > 0 ? "on" : "off")
    println(io, "Debias = ", debias ? "on" : "off")
    println(io, "Max IHT iterations = $max_iter")
    println(io, "Converging when tol < $tol and iteration ≥ $min_iter:\n")
    io != stdout && print_parameters(stdout, k, d, l, use_maf, group, debias, tol, max_iter, min_iter)
end
print_parameters(k, d, l, use_maf, group, debias, tol, max_iter, min_iter) =
    print_parameters(stdout, k, d, l, use_maf, group, debias, tol, max_iter, min_iter)

function check_convergence(v::IHTVariable)
    the_norm = max(chebyshev(v.b, v.b0), chebyshev(v.c, v.c0)) #max(abs(x - y))
    scaled_norm = the_norm / (max(norm(v.b0, Inf), norm(v.c0, Inf)) + 1.0)
    return scaled_norm
end

function backtrack!(v::IHTVariable, η::Float)
    # recompute gradient step
    copyto!(v.b, v.b0)
    copyto!(v.c, v.c0)
    _iht_gradstep!(v, η)

    # recompute η = xb, μ = g(η), and loglikelihood to see if we're now increasing
    update_xb!(v)
    update_μ!(v)
    if v.est_r != :None
        v.d = mle_for_r(v)
    end

    return loglikelihood(v)
end

function check_data_dim(y::AbstractVecOrMat, x::AbstractMatrix, z::AbstractVecOrMat)
    if is_multivariate(y)
        r, n1 = size(y)
        p, n2 = size(x)
        q, n3 = size(z)
        n1 == n2 == n3 || throw(DimensionMismatch("Detected multivariate analysis" *
            " but size(y, 2) = $n1, size(x, 2) = $n2, size(z, 2) = $n3 which don't " *
            "match. Recall each column of `y`, `x`, `z` should be sample " *
            "phenotypes/genotypes/covariates."))
    else
        n1 = length(y)
        n2, p = size(x)
        n3, = size(z)
        n1 == n2 == n3 || throw(DimensionMismatch("Detected univariate analysis " *
        "but length(y) = $n1, size(x, 1) = $n2, size(z, 1) = $n3 which don't match. " *
        "Recall each `y` should be a vector of phenotypes, and each row of `x`" *
        " and `z` should be sample genotypes/covariates."))
    end
end

function save_best_model!(v::IHTVariable)
    # compute η = xb with the best estimated model
    copyto!(v.b, v.best_b)
    copyto!(v.c, v.best_c)
    v.idx .= v.b .!= 0
    v.idc .= v.c .!= 0
    check_covariate_supp!(v)
    update_xb!(v)

    # update estimated mean μ with genotype predictors
    update_μ!(v.μ, v.xb, v.l)
end

"""
    debias!(v::IHTVariable)

After each IHT iteration, `β` is sparse. This function solves for the exact
solution `β̂` on the non-zero indices of `β`, a process known as debiasing.
"""
function debias!(v::IHTVariable)
    v.memory_efficient && error("Currently debiasing only works with memory_efficient=false")
    if sum(v.idx) == size(v.xk, 2)
        temp_glm = fit(GeneralizedLinearModel, v.xk, v.y, v.d, v.l)
        view(v.b, v.idx) .= temp_glm.pp.beta0
    end
end


#############################################
# Van der Pol Oscillator Controlled by RFPT #
#############################################
Adaptive=1
######################
# Control Parameters #
######################
K=1e5
B=-1
A=1e-5

#########
# Time  #
#########
δt=1e-3
LONG=Int(2e4)

########################################
# Kinematic Block Parameter (2nd order)#
########################################
Λ=3

####################################
# Parameters for Nominal Trajectory#
####################################

ω=0.5
Ampl=2
##########################
# Exact Model Parameters #
##########################

μₑ=0.4
ωₑ=0.46
αₑ=1
λₑ=0.1
mₑ=1

################################
# Approximate Model Parameters #
################################

μₐ=0.5
ωₐ=0.42
αₐ=0.9
λₐ=0.09
mₐ=0.8

############################
# Functions for simulation #
############################

# function nom_traj(t)
#   #egys=[1,-1.5,2]
#   qN=Ampl*sin(ω*δt*t)
#   qN_p=Ampl*ω*cos(ω*δt*t)
#   qN_pp=-Ampl*ω^2*sin(ω*δt*t)
# 	#println(qN[t,:])
#   return [qN;qN_p;qN_pp]
# end

#######################
# The Kinematic Block #
#######################

# function Kinematic_Block(t,err_int)
# 	#Compute the Error
# 	err=(qN-q)#'
# 	#Compute the 1st time derivative of the Error
# 	err_p=(qN_p-q_p)#'
# 	#Compute te integral of the Error (with Euler formula)
#    #println(err)
# 	err_i=err_int+err*δt
#   #println(q,q_p,q_pp)
# 	#Compute the Desired acceleration (2nd time derivative)
# 	#out=qN_pp[t]'+3*Λ^2*err+3*Λ*err_p+Λ^3*err_i
#   out=qN_pp+3*Λ^2*err+3*Λ*err_p+Λ^3*err_i #Has to be colum vector
#   #println(out)
# 	#println(err_i)
#    #println(typeof(err_i))
# 	return [out err_i]
# end

###############
# Exact Model #
###############
function Exact(u,q,q_p)
  q_pp=(u+μₑ*(1-q^2)*q_p-ωₑ^2*q-αₑ*q^3-λₑ^5)/mₑ
  return q_pp
end

################
# Approx Model #
################

function Approx(q,q_p,q_pp)
  u=mₐ*q_pp-μₐ*(1-q^2)*q_p+ωₐ^2*q+αₐ*q^3+λₐ^5
  return u
end

function sigmoid(x)
	s=x/(1+abs(x))
  return s
end

function G(past_input,past_response,xDnow)#(f,x,xᵈ)
  out=(K+past_input) * (1 + B * sigmoid(A * (past_response - xDnow))) - K
  return out #q_ppReq
end

##############################
# Define arrays for Plotting #
##############################
#time
time_mem=zeros(LONG)
#nominal trajectory
#qN_mem=zeros(LONG)
#qN_p_mem=zeros(LONG)
#qN_pp_mem=zeros(LONG)

#q_ppDes_mem=zeros(LONG)
#x=0
#f=0
#q_ppDef_mem=zeros(LONG)
#Q_mem=zeros(LONG)
#err_int=0
#qN=0
#qN_p=0
#qN_pp=0
#q=0
#q_p=0
#q_pp=0
#q_mem=zeros(LONG)
##################
# The simulation #
##################
q_mem=zeros(LONG) #A megvalósult pálya
q_p_mem=zeros(LONG)

qN_mem=zeros(LONG) #A nominális pálya
qN_p_mem=zeros(LONG)

u_mem=zeros(LONG) #A szanályozó jel elmentése

qN_pp_mem=zeros(LONG) #A nominális gyorsulások értéke
q_pp_Des_mem=zeros(LONG)#A PID korrekciós adatok
q_pp_Req_mem=zeros(LONG) #Az Adaptívan torzított jelre
q_pp_mem=zeros(LONG) #Az Adaptívan torzított jelre

h_int=0 #A követési hiba integrálja

past_input=0
past_response=0

for t=1:LONG-1
  #store time for plotting
  time_mem[t]=t*δt
  #define the nominal trajectory
  qN_mem[t]=Ampl*sin(ω*δt*t)
  qN_p_mem[t]=ω*Ampl*cos(ω*δt*t)
  qN_pp=-ω^2*Ampl*sin(ω*δt*t)
  qN_pp_mem[t]=qN_pp
  #the kinematic
  q_ppDes=Λ^3*h_int+3*Λ^2*(qN_mem[t]-q_mem[t])+3*Λ*(qN_p_mem[t]-q_p_mem[t])+qN_pp
  q_pp_Des_mem[t]=q_ppDes
  if Adaptive==1 && t>10
    q_ppReq=G(past_input,past_response,q_ppDes)
  else
    q_ppReq=q_ppDes
  end
  q_pp_Req_mem[t]=q_ppReq
  u=Approx(q_mem[t],q_p_mem[t],q_ppReq)
  u_mem[t]=u
  #println(Q)
  past_input=q_ppReq

  q_pp=Exact(u,q_mem[t],q_p_mem[t])
  #println(x)
  #println(x)
  q_pp_mem[t]=q_pp

  past_response=q_pp
  q_p_mem[t+1]=q_p_mem[t]+δt*q_pp
  q_mem[t+1]=q_mem[t]+δt*q_p_mem[t]
  h_int=h_int+δt*(qN_mem[t]-q_mem[t])

end #for t=1:LONG-1

using(PyPlot)
figure(1)
grid("on")
title("Tracking Error vs Time")
xlabel("Time [s]")
ylabel("Tracking Error")
plot(time_mem[11:LONG-1],qN_mem[11:LONG-1]-q_mem[11:LONG-1])

figure(2)
grid("on")
title("Nominal and Realized Trajectory")
xlabel("Time [s]")
ylabel("Nominal and Realized Trajectory")
plot(time_mem[3:LONG-1],qN_mem[3:LONG-1],color="red")
plot(time_mem[3:LONG-1],q_mem[3:LONG-1],color="green")

examples/RFPT_SISO/van der Pol Oscillator/van_der_Pol_RFPT.jl


include("./tools.jl");
include("./kalman.jl");
include("./subsampling.jl");
#include("./varima.jl");


"""
A module that makes expressions like `-π` behave like Irrational, rather
than Float64.

Generates IrrationalExpr objects when arithmetic operations are applied to
objects of type Irrational, and these operations are carried out in the
destination type when type conversion occurs.

For example, `BigFloat(π) + BigFloat(-π)` will evaluate to BigFloat(0),
rather than to approximately 1.2e-16.
"""

module IrrationalExpressions

import Base: +, -, *, /, convert, promote_rule, show

immutable IrrationalExpr{op, N} <: Real
  args::NTuple{N,Union{IrrationalExpr, Irrational, Rational, Integer}}
end

@generated convert{T<:AbstractFloat,op,N}(::Type{T}, x::IrrationalExpr{op,N}) =
  Expr(:call, op, [Expr(:call, :convert, T, :( x.args[$i] )) for i=1:N]...)

promote_rule{T1<:AbstractFloat, T2<:IrrationalExpr}(::Type{T1}, ::Type{T2}) = T1
promote_rule{T2<:IrrationalExpr}(::Type{BigFloat}, ::Type{T2}) = BigFloat

## Unary operators
(+)(x::IrrationalExpr) = x
(-)(x::IrrationalExpr) = IrrationalExpr{:(-),1}((x,))
(-)(x::Irrational) = IrrationalExpr{:(-),1}((x,))

## Binary operators
ops = (:(+), :(-), :(*), :(/))
types = (IrrationalExpr, Irrational, Rational, Integer)
for op in ops, i in eachindex(types), j in eachindex(types)
  if i<=2 || j<=2
    @eval $op(x::$(types[i]), y::$(types[j])) = IrrationalExpr{Symbol($op),2}((x,y))
  end
end

# We define expr() as a conveninent middle step to get to strings,
# but this also allows eval(expr(x)) as a roundabout way to get x.
expr(x) = x
expr{sym}(::Irrational{sym}) = sym
expr{op,N}(x::IrrationalExpr{op,N}) = Expr(:call, op, map(expr,x.args)...)

show(io::IO, x::IrrationalExpr) = print(io, string(expr(x)), " = ", string(Float64(x))[1:end-2], "...")

end # module


# Code for Video 2.2:

fac(n) = (n <= 1)? 1 : n * fac(n - 1)

n = 20
@printf("The %d-th Factorial number is %d", n, fac(n))
println("")
print("The calculation took")
@time fac(n)

# The 20-th Factorial number is 2432902008176640000
# The calculation took  0.000006 seconds (149 allocations: 10.167 KB)

# C:\Users\CVO>julia --lisp
# ;  _
# ; |_ _ _ |_ _ |  . _ _
# ; | (-||||_(_)|__|_)|_)
# ;-------------------|------

# > (* 3 45)
# 135

# >

 n = 108 # 108
 n # 108
 :n # :n

 typeof(:n) # Symbol

 e = :(a + b) # :(a + b)

 typeof(e) # Expr

 a = 2 # 2
 b = 3 # 3
 eval(e) # 5

 fieldnames(e)
# 3-element Array{Symbol,1}:
#  :head
#  :args
#  :typ

 e.args
# 3-element Array{Any,1}:
#  :+
#  :a
#  :b

 c = quote
           a = 2
           b = 3
           a + b
       end
# quote  # none, line 2:
#     a = 2 # none, line 3:
#     b = 3 # none, line 4:
#     a + b
# end

 eval(c) # 5

 typeof(c) # Expr

 c = quote
           a = 2
           b = 3
           $a + $b
       end
# quote  # none, line 2:
#     a = 2 # none, line 3:
#     b = 3 # none, line 4:
#     2 + 3
# end

macro timeit(ex)
    quote
        local t0 = time()
        local val = $(esc(ex))
        local t1 = time()
        println("elapsed time in seconds: ")
        @printf "%.3f\n" t1 - t0
        val
    end
end

@timeit fac(20)
# elapsed time in seconds:
# 0.002
# 2432902008176640000

macro repeat(n, body)
    quote
        for i = 1:$(esc(n))
            $(esc(body))
        end
    end
end

@repeat(7, println("Hi Julia"))
# Hi Julia
# Hi Julia
# Hi Julia
# Hi Julia
# Hi Julia
# Hi Julia
# Hi Julia

macroexpand(:@repeat(7, println("Hi Julia")))
# quote  # none, line 3:
#     for #57#i = 1:7 # none, line 4:
#         println("Hi Julia")
#     end
# end

@repeat(3, for i in 1:5 println(i*i) end)
# 1
# 4
# 9
# 16
# 25
# 1
# 4
# 9
# 16
# 25
# 1
# 4
# 9
# 16
# 25

macroexpand(:@repeat(3, for i in 1:5 println(i*i) end))
# quote  # none, line 3:
#     for #29#i = 1:3 # none, line 4:
#         for i = 1:5 # none, line 1:
#             println(i * i)
#         end
#     end
# end

# @assert  108 = (2 + 105)
# ERROR: syntax: invalid assignment location "108"

# @assert  108 == (2 + 105)
# ERROR: AssertionError: 108 == 2 + 105

 using Base.Test

# @test_approx_eq 2 2.1
# ERROR: assertion failed: |2 - 2.1| <= 4.440892098500626e-12
#   2 = 2
#   2.1 = 2.1
#   difference = 0.10000000000000009 > 4.440892098500626e-12
#  in error at error.jl:22

# @test_approx_eq_eps 2 2.1 0.1
# ERROR: assertion failed: |2 - 2.1| <= 0.1
#   2 = 2
#   2.1 = 2.1
#   difference = 0.10000000000000009 > 0.1
#  in error at error.jl:22

@test_approx_eq_eps 2 2.1 0.15

@windows ? println("I am on Windows") : println("No Windows here")
# I am on Windows

@which (1 +2im) * (2 - 8im)
# *(z::Complex{T<:Real}, w::Complex{T<:Real}) at complex.jl:113


export chandra, cyclic, katsura, fourbar, rps10, ipp, ipp2, boon, heart,
    d1, bacillus_subtilis, griewank_osborne, tritangents, cyclooctane

"""
    chandra(n)

The Chandrasekhar H-Equation for `n`.

## References

Laureano Gonzalez-Vega: "Some examples on problem solving by using the
  symbolic viewpoint when dealing with polynomial systems of equations".
  in: "Computer Algebra in Science and Engineering", Editors: J. Fleischer,
  J. Grabmeier, F.W. Hehl and W. Kuechlin.  Pages 102-116
  World Scientific Publishing, 1995.

S. Chandrasekhar: "Radiative Transfer", Dover, NY, 1960.

C.T. Kelley: "Solution of the Chandrasekhar H-equation by Newton's method".
  J. Math. Phys., 21 (1980), pp. 1625-1628.

Jorge J. More': "A collection of nonlinear model problems"
  in: "Computational Solution of Nonlinear Systems of Equations",
  Editors: Eugene L. Allgower and Kurt Georg.  Pages 723-762.
  Lectures in Applied Mathematics, Volume 26, AMS, 1990.

NOTE : the parameter c equals 0.51234.
  In general c can be any number in the interval (0,1]
"""
function chandra(n)
    @polyvar H[1:n]

    return TestSystem([(2n*H[i] - 0.51234H[i]*(1+ sum(i / (j+i) * H[j] for j=1:(n-1))) - 2n) for i=1:n],
        multi_bezout_number=(2^(n-1), [H[1:n-1], [H[n]]]),
        mixed_volume=2^(n-1))
end

"""
    cyclic(n)

The cyclic n-roots problems.

## References

Göran Björck and Ralf Fröberg:
`A faster way to count the solutions of inhomogeneous systems
 of algebraic equations, with applications to cyclic n-roots',
in J. Symbolic Computation (1991) 12, pp 329--336.
"""
function cyclic(n)
    @polyvar z[1:n]

    mv = nothing
    if n == 5
        mv = 70
    elseif n == 6
        mv = 156
    elseif n == 7
        mv = 924
    elseif n == 8
        mv = 2560
    elseif n == 9
        mv = 11016
    elseif n == 10
        mv = 35940
    elseif n == 11
        mv = 184_756
    end

    TestSystem([(sum(prod(z[(k-1) % n + 1] for k in j:j+m) for j in 1:n) for m=0:(n-2))...,prod(z)-1],
        mixed_volume = mv)
end

"""
    katsura(n)

A problem of magnetism in physics.

## References

From the PoSSo test suite.

Shigetoshi Katsura: "Users posing problems to PoSSO",
in the PoSSo Newsletter, no. 2, July 1994,
edited by L. Gonzalez-Vega and T. Recio.
Available at http://janet.dm.unipi.it/

S. Katsura, W. Fukuda, S. Inawashiro, N.M. Fujiki and R. Gebauer,
Cell Biophysics, Vol 11, pages 309--319, 1987.

W. Boege, R. Gebauer, and H. Kredel:
"Some examples for solving systems of algebraic equations by
 calculating Groebner bases", J. Symbolic Computation, 2:83-98, 1986.

Shigetoshi Katsura:
 "Spin Glass Problem by the Method of Integral Equation
  of the Effective Field".  In "New Trends in Magnetism",
edited by Mauricio D. Coutinho-Filho and Sergio M. Resende,
pages 110-121, World Scientific, 1990.
"""
function katsura(n)
  @polyvar x[0:n] # This creates variables x0, x1, ...

  return TestSystem([
    (sum(x[abs(l)+1]*x[abs(m-l)+1] for l=-n:n if abs(m-l)<=n) -
    x[m+1] for m=0:n-1)...,
    x[1] + 2sum(x[i+1] for i=1:n) - 1
  ], mixed_volume=2^n)
end

# Mechanisms

"""
    fourbar()
A four-bar design problem, so-called 5-point problem.

## References

See Morgan, A.P. and Wampler, C.W. :
 `Solving a planar four-bar design problem using continuation'
in Transaction of the ASME, J. of Mechanical Design, Vol. 112
pages 544-550, 1990.

For the coefficients, see Table 2, a = (1,0).
This is the start system, with five random precision points,
that has been used to solve twenty other test systems.
"""
function fourbar()
    @polyvar X1 X2 Y1 Y2

    return TestSystem([
        0.01692601*X1^2*Y1^2 - 0.888509280014*X1^2*Y2^2 + 0.0411717692438*X2^2*Y1^2 - 0.00437457395884*X2^2*Y2^2 + 0.331480641249*X1*X2*Y1^2 - 1.38036964668*X1*X2*Y2^2 - 0.270492270191*X1^2*Y1*Y2 + 1.44135801774*X2^2*Y1*Y2 + 0.859888946812*X1*X2*Y1*Y2 + 0.0791489659197*X1^2*Y1 - 0.00336032777032*X1^2*Y2 - 0.0620826738427*X1*Y1^2 + 0.501879647495*X1*Y2^2 + 0.647156236961*X2^2*Y1 + 0.0926311741907*X2^2*Y2 - 0.255000006226*X2*Y1^2 - 0.0896892386081*X2*Y2^2 - 0.568007271041*X1*X2*Y2 + 0.095991501961*X1*X2*Y1 + 0.165310767618*X1*Y1*Y2 - 0.563962321337*X2*Y1*Y2 - 0.0784871167595*X1*Y1 - 0.0784871167595*X2*Y2 + 0.011807283256*X1*Y2  - 0.011807283256*X2*Y1 + 0.0422876985355*X1^2 + 0.0422876985355*X2^2 + 0.0372427422943*Y1^2 + 0.0372427422943*Y2^2,
        0.518178672335*X1^2*Y1^2 - 0.0414464807343*X1^2*Y2^2 + 2.63600135179*X2^2*Y1^2 - 0.799490472298*X2^2*Y2^2 + 0.29442805494*X1^2*Y1*Y2 + 1.46551534655*X2^2*Y1*Y2 - 0.631878110759*X1*X2*Y1^2 - 1.80296540237*X1*X2*Y2^2 - 2.87586667102*X1*X2*Y1*Y2 - 0.987856648177*X1^2*Y1 - 0.530579106676*X1^2*Y2 - 0.0397576281649*X1*Y1^2 + 0.317719102869*X1*Y2^2 - 1.93710490787*X2^2*Y1  + 0.00127693327315*X2^2*Y2 - 0.581380074072*X2*Y1^2 - 0.0672137066743*X2*Y2^2 + 0.531856039949*X1*X2*Y1 + 0.949248259696*X1*X2*Y2 + 0.514166367398*X1*Y1*Y2 - 0.357476731033*X2*Y1*Y2 + 0.140965913657*X1*Y1  + 0.140965913657*X2*Y2 - 0.153347218606*X1*Y2 + 0.153347218606*X2*Y1 + 0.283274882058*X1^2  + 0.283274882058*X2^2 + 0.0382903330079*Y1^2 + 0.0382903330079*Y2^2,
        0.0233560008057*X1^2*Y1^2 - 0.00428427501149*X1^2*Y1*Y2 - 0.792756311827*X1^2*Y2^2 + 0.0492185850289*X2^2*Y2^2 + 0.0759264856293*X1*X2*Y1^2 + 1.14839711492*X1*X2*Y1*Y2 - 0.283066217262*X2^2*Y1^2 + 0.460041521291*X2^2*Y1*Y2 - 0.388399310674*X1*X2*Y2^2 - 0.0561169736293*X1*Y1^2 + 0.485064247792*X1*Y2^2 + 0.0689567235492*X1^2*Y1 - 0.115620658768*X1^2*Y2 - 0.13286905328*X2*Y1^2 - 0.084375901147*X2*Y2^2 + 0.639964831612*X2^2*Y1 + 0.101386684276*X2^2*Y2 + 0.217007343044*X1*X2*Y1 - 0.571008108063*X1*X2*Y2 + 0.0484931521334*X1*Y1*Y2 - 0.541181221422*X2*Y1*Y2 - 0.00363197918253*X1*Y2 + 0.00363197918253*X2*Y1 - 0.0781302968652*X1*Y1 - 0.0781302968652*X2*Y2 + 0.0471311092612*X1^2 + 0.0471311092612*X2^2 + 0.0324495575052*Y1^2 + 0.0324495575052*Y2^2,
        0.393707415641*X1^2*Y1^2 + 0.59841456862*X1^2*Y2^2 + 0.0735854940135*X2^2*Y1^2 + 0.0548997238169*X2^2*Y2^2 + 0.0116156836985*X1^2*Y1*Y2 + 0.0699694273575*X2^2*Y1*Y2 - 0.305757340849*X1*X2*Y1^2 - 0.364111084508*X1*X2*Y2^2 - 0.223392923175*X1*X2*Y1*Y2 + 0.0996725944534*X1^2*Y1 + 0.0113936468426*X1^2*Y2 - 0.381205205249*X1*Y1^2 - 0.473402150235*X1*Y2^2 - 0.0213613191759*X2^2*Y1 - 0.0372595571271*X2^2*Y2 + 0.148904552394*X2*Y1^2 + 0.142408744984*X2*Y2^2 - 0.0486532039697*X1*X2*Y1 + 0.121033913629*X1*X2*Y2 - 0.00649580741066*X1*Y1*Y2 + 0.092196944986*X2*Y1*Y2 - 0.0483106652705*X1*Y1  - 0.0483106652705*X2*Y2 - 0.00316794272326*X1*Y2 + 0.00316794272326*X2*Y1 + 0.00634952598374*X1^2 + 0.00634952598374*X2^2 + 0.0922886309144*Y1^2  + 0.0922886309144*Y2^2
    ], mixed_volume=80)
end

"""
    ipp()

The six-revolute-joint problem of mechanics.

## References

A. Morgan and A. Sommese
`Computing all solutions to polynomial systems
 using homotopy continuation',
 Appl. Math. Comput., Vol. 24, pp 115-138, 1987.
"""
function ipp()
    @polyvar x1 x2 x3 x4 x5 x6 x7 x8
    TestSystem([x1^2+x2^2-1,
        x3^2+x4^2-1,
        x5^2+x6^2-1,
        x7^2+x8^2-1,
        (-2.4915068E-01*x1*x3+ 1.6091354E+00*x1*x4+ 2.7942343E-01 *x2*x3+ 1.4348016E+00*x2*x4) + (4.0026384E-01*x5*x8-8.0052768E-01*x6*x7+ 7.4052388E-02*x1-8.3050031E-02*x2) - (3.8615961E-01*x3-7.5526603E-01*x4+ 5.0420168E-01*x5 -1.0916287E+00*x6+ 4.0026384E-01*x8) + 4.920729E-02,
        (1.2501635E-01*x1*x3-6.8660736E-01*x1*x4-1.1922812E-01* x2*x3-7.1994047E-01*x2*x4) - (4.3241927E-01*x5*x7-8.6483855E-01*x6*x8-3.715727E-02*x1+ 3.5436896E-02*x2)+ 8.5383482E-02*x3-3.9251967E-02*x5-4.3241927E-01*x7+ 1.387301E-02,
        (-6.3555007E-01*x1*x3-1.1571992E-01*x1*x4-6.6640448E-01 *x2*x3) + (1.1036211E-01*x2*x4+ 2.9070203E-01*x5*x7+ 1.2587767E+00*x5*x8)- (6.2938836E-01*x6*x7+ 5.8140406E-01*x6*x8+ 1.9594662E-01*x1)- (1.2280342E+00*x2-7.9034221E-02*x4+ 2.6387877E-02*x5)- 5.713143E-02*x6-1.1628081E+00*x7+1.2587767E+00*x8+ 2.162575E+00,
        (1.4894773E+00*x1*x3+ 2.3062341E-01*x1*x4+ 1.3281073E+00*x2*x3)-(2.5864503E-01*x2*x4+ 1.165172E+00*x5*x7-2.6908494E-01*x5*x8)+ (5.3816987E-01*x6*x7+ 5.8258598E-01*x6*x8-2.0816985E-01*x1)+(2.686832E+00*x2-6.9910317E-01*x3+ 3.5744413E-01*x4)+ 1.2499117E+00*x5+ 1.467736E+00*x6+ 1.165172E+00*x7+ 1.10763397E+00*x8-6.9686809E-01
        ], mixed_volume=64, nreal_solutions=10, nsolutions=32)
end

"""
    ipp2()

The 6R inverse position problem.

# References

This system occurs as Example 3.3 in a paper by Charles Wampler:
`Bezout Number Calculations for Multi-Homogeneous Polynomial Systems',
Appl. Math. Comput. vol. 51 No. 2--3, pp. 143--157.

For the original formulation of the problem, see
 Charles Wampler and Alexander Morgan:
 `Solving the 6R inverse position problem using a generic-case solution
  methodology', Mech. Mach. Theory, Vol. 26, No. 1, pp. 91-106, 1991.
"""
function ipp2()
    @polyvar z21 z22 z31 z32 z33 z41 z42 z43 z51 z52 z53
    TestSystem([ z21^2 + z22^2 - 1.25830472585 + 1.05384271933im,
        z31^2 + z32^2 + z33^2 - 1,
        z41^2 + z42^2 + z43^2 - 1,
        z51^2 + z52^2 + z53^2 - 1,
        z21*z31 + z22*z32 + (0.642935654806 + 0.819555356316im)*z33- 0.266880023988 - 0.452565255666im,
        z31*z41 + z32*z42 + z33*z43 - 0.26425551745 - 0.342483846503im,
        z41*z51 + z42*z52 + z43*z53 - 0.126010863922 - 0.864590917688im,
        (0.352598136811 + 0.116888144319im)*z51 + (0.539042485525 + 0.687058436892im)*z52 + (0.391154215376 + 0.128900893182im)*z53 - 0.179560356712 - 0.8709166566im,
        (0.984138451804 + 0.414967172346im)*z21 + (0.958341609741 + 0.847442419999im)*z22*z33 + (-0.496254299764 - 0.546020011741im)*z22 + (0.353268870498 + 0.389909226888im)*z31 + (0.964759562277 + 0.71397074519im)*z32*z43 + (0.0783739840935 - 1.33026494666im)*z32 + (-0.964759562277 - 0.71397074519im)*z33*z42 + (0.204379350351 + 0.00374294529684im)*z41 + (0.706319991205 + 0.120097702053im)*z42*z53 + (-0.706319991205 - 0.120097702053im)*z43*z52 + (0.907681632783 + 0.405209293447im)*z51 + (0.148939301127 + 0.182393186752im)*z52 + (-0.0486385369088 - 0.496934768083im)*z53 - 0.437312713588 - 0.914780691357im,
        (-0.958341609741 - 0.847442419999im)*z21*z33 + (0.496254299764 + 0.546020011741im)*z21 + (0.984138451804 + 0.414967172346im)*z22 + (-0.964759562277 - 0.71397074519im)*z31*z43 + (-0.0783739840935 + 1.33026494666im)*z31 + (0.353268870498 + 0.389909226888im)*z32 + ( 0.964759562277 + 0.71397074519im)*z33*z41 + (-0.706319991205 - 0.120097702053im)*z41*z53 + (0.204379350351 + 0.00374294529684im)*z42 + (0.706319991205 + 0.120097702053im)*z43*z51 + (-0.148939301127 - 0.182393186752im)*z51 + (0.907681632783 + 0.405209293447im)*z52 + (0.134045297503 + 0.164748774862im)*z53 - 0.719086796333 - 0.691791591267im,
        (0.958341609741 + 0.847442419999im)*z21*z32 + (-0.958341609741 - 0.847442419999im)*z22*z31 + (0.964759562277 + 0.71397074519im)*z31*z42 + (-0.964759562277 - 0.71397074519im)*z32*z41 + (0.353268870498 + 0.389909226888im)*z33 + (0.706319991205 + 0.120097702053im)*z41*z52 + (-0.706319991205 - 0.120097702053im)*z42*z51 + (0.204379350351 + 0.00374294529684im)*z43 + (0.0486385369088 + 0.496934768083im)*z51  + (-0.134045297503 - 0.164748774862im)*z52 + (0.907681632783 + 0.405209293447im)*z53 - 0.64863071126 + 0.983034576618im
    ],
    multi_bezout_number=(320, [[z21, z22, z41, z42, z43], [z31, z32, z33, z51, z52, z53]]),
    mixed_volume=288)
end

include("rps10.jl")

# Mixed

"""
    boon()

Neurofysiology, posted by Sjirk Boon. Note that the system

## References

The system has been posted to the newsgroup
sci.math.num-analysis by Sjirk Boon.

P. Van Hentenryck, D. McAllester and D. Kapur:
`Solving Polynomial Systems Using a Branch and Prune Approach'
SIAM J. Numerical Analysis, Vol. 34, No. 2, pp 797-827, 1997.

"""
function boon()
    @polyvar s1 g1 s2 g2 C1 C2
    TestSystem([s1^2+g1^2 - 1,
     s2^2+g2^2 - 1,
     C1*g1^3+C2*g2^3 - 1.2,
     C1*s1^3+C2*s2^3 - 1.2,
     C1*g1^2*s1+C2*g2^2*s2 - 0.7,
     C1*g1*s1^2+C2*g2*s2^2 - 0.7], mixed_volume=20, nsolutions=8, nreal_solutions=8)
end

"""
    heart()

The heart-dipole problem.

## References

Nelsen, C.V. and Hodgkin, B.C.:
`Determination of magnitudes, directions, and locations of two independent
dipoles in a circular conducting region from boundary potential measurements'
IEEE Trans. Biomed. Engrg. Vol. BME-28, No. 12, pages 817-823, 1981.

Morgan, A.P. and Sommese, A.J.:
`Coefficient-Parameter Polynomial Continuation'
Appl. Math. Comput. Vol. 29, No. 2, pages 123-160, 1989.
Errata: Appl. Math. Comput. 51:207 (1992)

Morgan, A.P. and Sommese, A. and Watson, L.T.:
`Mathematical reduction of a heart dipole model'
J. Comput. Appl. Math. Vol. 27, pages 407-410, 1989.
"""
function heart()
    @polyvar a b c d t u v w
    TestSystem([a + b - 0.63254;
        c + d + 1.34534;
        t*a + u*b - v*c - w*d + 0.8365348;
        v*a + w*b + t*c + u*d - 1.7345334;
        a*t^2 - a*v^2 - 2*c*t*v + b*u^2 - b*w^2 - 2*d*u*w - 1.352352;
        c*t^2 - c*v^2 + 2*a*t*v + d*u^2 - d*w^2 + 2*b*u*w + 0.843453;
        a*t^3 - 3*a*t*v^2 + c*v^3 - 3*c*v*t^2 + b*u^3 - 3*b*u*w^2 + d*w^3 - 3*d*w*u^2 + 0.9563453;
        c*t^3 - 3*c*t*v^2 - a*v^3 + 3*a*v*t^2 + d*u^3 - 3*d*u*w^2 - b*w^3 + 3*b*w*u^2 - 1.2342523
        ], mixed_volume=121, nsolutions=4, nreal_solutions=2)
end

"""
    d1()

A a sparse system, known as benchmark D1.

## References

H. Hong and V. Stahl:
`Safe Starting Regions by Fixed Points and Tightening',
Computing 53(3-4): 322-335, 1995.
"""
function d1()
    @polyvar x1 x2 x3 x4 x5 x6 x7 x8 x9 x10 x11 x12
    TestSystem([x1^2  + x2^2 - 1;
         x3^2  + x4^2 - 1;
         x5^2  + x6^2 - 1;
         x7^2  + x8^2 - 1;
         x9^2  + x10^2 - 1;
         x11^2 + x12^2 - 1;
         3*x3 + 2*x5 + x7 - 3.9701;
         3*x1*x4 + 2*x1*x6 + x1*x8 - 1.7172;
         3*x2*x4 + 2*x2*x6 + x2*x8 - 4.0616;
         x3*x9 + x5*x9 + x7*x9 - 1.9791;
         x2*x4*x9 + x2*x6*x9 + x2*x8*x9 + x1*x10 - 1.9115;
         - x3*x10*x11 - x5*x10*x11 - x7*x10*x11 + x4*x12 + x6*x12 + x8*x12 - 0.4077],
         mixed_volume=192,
         nreal_solutions=16,
         nsolutions=48)
end

"""

    bacillus_subtilis()

The system comes from biochemical reactions (Published here https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1005267).
Basically it models how the bacteris B. subtilis produces a stress response protein, sigmaB, in response to some stress input. This input to the system is the concentration of a phosphatase (a kind of protein). In this system I call that variable phos
"""
function bacillus_subtilis()
    @polyvar w w2 w2v v w2v2 vP sigmaB w2sigmaB vPp phos
    poly = [(-1 * 0.7 * w + -2 * 3600.0 * (w ^ 2 / 2) + 2 * 18.0 * w2)*(0.2 + sigmaB) + 4.0 * 0.4 * (1 + 30.0sigmaB),
     -1 * 0.7 * w2 + 3600.0 * (w ^ 2 / 2) + -1 * 18.0 * w2 + -1 * 3600.0 * w2 * v + 18.0w2v + 36.0w2v + -1 * 3600.0 * w2 * sigmaB + 18.0w2sigmaB,
     -1 * 0.7 * w2v + 3600.0 * w2 * v + -1 * 18.0 * w2v + -1 * 3600.0 * w2v * v + 18.0w2v2 + -1 * 36.0 * w2v + 36.0w2v2 + 1800.0 * w2sigmaB * v + -1 * 1800.0 * w2v * sigmaB,
     (-1 * 0.7 * v + -1 * 3600.0 * w2 * v + 18.0w2v + -1 * 3600.0 * w2v * v + 18.0w2v2 + -1 * 1800.0 * w2sigmaB * v + 1800.0 * w2v * sigmaB + 180.0vPp)*(0.2 + sigmaB) + 4.5 * 0.4 * (1 + 30.0sigmaB),
     -1 * 0.7 * w2v2 + 3600.0 * w2v * v + -1 * 18.0 * w2v2 + -1 * 36.0 * w2v2,
     -1 * 0.7 * vP + 36.0w2v + 36.0w2v2 + -1 * 3600.0 * vP * phos + 18.0vPp,
     (-1 * 0.7 * sigmaB + -1 * 3600.0 * w2 * sigmaB + 18.0w2sigmaB + 1800.0 * w2sigmaB * v + -1 * 1800.0 * w2v * sigmaB)*(0.2 + sigmaB) + 0.4 * (1 + 30.0sigmaB),
     -1 * 0.7 * w2sigmaB + 3600.0 * w2 * sigmaB + -1 * 18.0 * w2sigmaB + -1 * 1800.0 * w2sigmaB * v + 1800.0 * w2v * sigmaB,
     -1 * 0.7 * vPp + 3600.0 * vP * phos + -1 * 18.0 * vPp + -1 * 180.0 * vPp,
     (phos + vPp) - 2.0]
    TestSystem(poly; nsolutions=44)
end

"""

    tritangents()

See https://www.juliahomotopycontinuation.org/examples/tritangents/.
"""
function tritangents()
    @polyvar h[1:3] # variables for the plane
    @polyvar x[1:3] y[1:3] z[1:3] #variables for the contact points
    @polyvar c[1:20] #variables for the cubic

    #the quadric
    Q = x[3] - x[1] * x[2]
    #the cubic
    C = c ⋅ unique(kron([x;1], [x;1], [x;1]))

    #generate the system P for the contact point x
    P_x = [
      h ⋅ x - 1;
      Q;
      C;
      det([h differentiate(Q, x) differentiate(C, x)])
    ]

    #generate a copy of P for the other contact points y,z
    P_y = [p([h; x; c] => [h; y; c]) for p in P_x]
    P_z = [p([h; x; c] => [h; z; c]) for p in P_x]

    #define F
    F = [P_x; P_y; P_z]

    #create random complex coefficients for C
    # c₁ is result of randn(ComplexF64, 20)
    c₁ = Complex{Float64}[-0.524114+0.798188im, 0.23203+0.203011im, 0.24588+0.265859im, -0.0293688-0.0838781im, 1.29037+1.19918im, 0.225558+0.113139im, 0.757771+0.327947im, 0.716762-1.14998im, -0.721386+0.253249im, 0.22066+0.286244im, 1.06596-1.24645im, -0.287841-0.15352im, 0.980128+0.798647im, -0.918701-1.5926im, -0.325193-0.163909im, 0.0885901-0.437345im, 0.195358-1.07946im, 0.451222-0.677876im, -0.498035+0.368011im, -1.34059+1.35779im]
    #plug in c₁ for c
    poly = [f([h; x; y; z; c] => [h; x; y; z; c₁]) for f in F]
    TestSystem(poly; nsolutions=720)
end

"""

    griewank_osborne()

This is an illustration of a system for which Newton's method fails near an isolated root.
No matter how close one starts to the multiplicity-three isolated root at the origin, (x,y) = (0,0), Newton's method diverges.

Reference: Griewank & Osborne, 1983
"""
function griewank_osborne()
    @polyvar x y
    TestSystem([(29/16)*x^3 - 2*x*y, x^2 - y])
end

"""

    cyclooctane(c::Float64=2.0)

This is a system that describes the confirmation space of cyclooctane, where the squared distance between neighboring carbon atoms is `c`.  See

    juliahomotopycontinuation.org/examples/cyclooctane/

for a detailed description.
"""
function cyclooctane(c::Float64=2.0)
    @polyvar z[1:3, 1:6]
    z_vec = vec(z)[1:17]
    Z = [zeros(3) z[:,1:5] [z[1,6]; z[2,6]; 0] [√c; 0; 0]]

    F1 = [(Z[:, i] - Z[:, i+1]) ⋅ (Z[:, i] - Z[:, i+1]) - c for i in 1:7]
    F2 = [(Z[:, i] - Z[:, i+2]) ⋅ (Z[:, i] - Z[:, i+2]) - 8c/3 for i in 1:6]
    F3 = (Z[:, 7] - Z[:, 1]) ⋅ (Z[:, 7] - Z[:, 1]) - 8c/3
    F4 = (Z[:, 8] - Z[:, 2]) ⋅ (Z[:, 8] - Z[:, 2]) - 8c/3
    f = [F1; F2; F3; F4]
    TestSystem(f)
end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ d6a054e2-c7bf-11eb-3732-c3135a35f956
# load libraries

begin
	using CSV
	using DataFrames
	using Plots
end

# ╔═╡ b2ca8e4b-fe95-43cb-a5bc-605f69d0d581
md"""
# Sorting Algorithms
*By Paige Weber*

This project covers the implementation, computational complexity, and real-world performance of several common sorting algorithms. The algorithms implemented and tested are insertion sort (the control), merge sort, heapsort, and quicksort. Mergesort is implemented both sequentially and in parallel. Quicksort is implemented naïvely (to demonstrate the worst-case performance on an already-sorted list) and then implemented with a median-of-three pivot.

C was chosen as the language each algorithm was written in.

"""

# ╔═╡ aa573021-c79a-46c1-b4e0-021e4c3434ed
md"""
# Data Structures

Lists were represented as simple arrays, contiguous in memory, allocated with `aligned_alloc`.

The heap for heapsort is represented using a 1-dimensional array (in fact, the same array the initial list is stored in).

Accessing a known index has complexity `O(1)`. This is the only operation used.

"""

# ╔═╡ e72161bc-dab5-4585-9cf8-bc32ec84e988
md"""
# Sample Code

All sorting code may be found in the file `src/sorting-algorithms.c`.
"""

# ╔═╡ 3723574c-959a-4b04-91bf-6207c4ff6fd5
md"""
## Insertion Sort
Insertion sort is quite simple, and only took ~5 lines of actual code:

```c
void insertionSort(float* arr, size_t n) {
  for (size_t i = 1; i < n; i++) {
    size_t j = i;
    while (j > 0 && arr[j] < arr[j-1]) {
      swap(&arr[j], &arr[j-1]);
      j--;
    }
  }
}
```

"""

# ╔═╡ 209e2ecb-af90-4c48-a583-e9dda280198f
md"""
## Merge Sort

Because every other algorithm is significantly more complicated to implement and the full source code is available in this repository, I will omit whatever possible without making the code snippets less clear. Removed code will be indicated by commenting `// (removed) ...`.

My implementation does not create new arrays for each split and then rejoin them on merge. Instead, it uses indices to track which part of the array should be considered for the current instance of `mergesort`. Because this approach may require many swaps, a separate "buffer" array is used to store the sorted part of the array, which is then copied back into the main array.

The actual mergesort function behaves as follows. "Splitting" the array in two just means assigning a new pointer to the beginning of the array (and buffer) and setting the size to the correct value.

```c
void mergesort(float * array, float * buffer, size_t n) {
  // base case: an array of size 1 is always sorted
  if(n == 1)
    return;

  // create new pointers and sizes for the left and right halves. This is how
  // we "split" the array
  size_t left_n = n/2;
  float * left_array = array;
  float * left_buffer = buffer;

  size_t right_n = n - n/2;
  float * right_array = array + left_n;
  float * right_buffer = buffer + left_n;

  // mergesort each half
  mergesort(left_array, left_buffer, left_n);
  mergesort(right_array, right_buffer, right_n);

  // merge the sorted portions back together
  merge(left_array, left_n, right_array, right_n, buffer);
}
```

The merge is somewhat more complicated, requiring careful tracking of indices for both halves as well as the buffer.

```c
void merge(float * a, size_t a_n, float * b, size_t b_n, float * buffer) {
  size_t a_i = 0; size_t b_i = 0; buffer_i = 0;

  while(a_i < a_n && b_i < b_n) {
    if(a[a_i] > b[b_i]) {
      // then b should come first
      buffer[buffer_i] = b[b_i];
      b_i++;
    }
    else {
      // else a is <= b so we put a next
      buffer[buffer_i] = a[a_i];
      a_i++;
    }

    buffer_i++;
  }

  // (removed) if either 'a' or 'b' has items left, copy them to 'merged'
  // (removed) then move stuff from 'buffer' back into original array
}
```

As an extra exercise, I also parallelized mergesort, because one of the main advantages of this algorithm is the ease with which it can be parallelized. Performance improvements over sequential code will be considered later in the report.

```c

void mergesort_parallel(float * array, float * buffer, size_t n) {
  // (removed) small tasks are done sequentially

  #pragma omp task
  mergesort_parallel(array, buffer, n/2);

  #pragma omp task
  mergesort_parallel(array + n/2, buffer + n/2, n - n/2);

  // taskwait ensures that both mergesorts are completed before merging
  #pragma omp taskwait
  merge(array, n/2, array + n/2, n - n/2, buffer);
}
```

"""

# ╔═╡ fd11a60a-c3dd-4bb1-aeb3-e3be1b87f4f7
md"""
## Heapsort

The heap is created in-place. No additional memory is used. Because a heap is a complete binary tree, we can use a 1 dimensional array and know that for any node at index `i`, `i * 2 + 1` and `i * 2 + 2` are the left and right children of that node. In our case, the heap used is a max heap.

Heapsort only requires heapifying the array, then repeatedly swapping the largest element to the end and re-heapifying the root.

```c
void heapSort(float* arr, size_t n) {
  // heapify all subtrees (create initial max heap)

  // we want to move from right to left, because big numbers will be towards
  // the right

  // start at i = n/2 because anything at indices greater than n/2 will be a
  // leaf node
  for (int i = n/2 - 1; i > -1; i--) {
    heapify(arr, n, i);
  }

  // move largest element to end
  for (size_t i = n-1; i > 0; i--) {
    swap(&arr[0], &arr[i]);

    // heapify with 0 as root. Only heapify to i
    heapify(arr, i, 0);
  }
}
```

Heapify is also fairly simple. First, we check if either of a node's children are greater than itself. If so, we swap root with the larger child and repeat. The process terminates when the root is larger than both its children.

```c
// turn array[root:n] into max heap
void heapify(float* arr, size_t n, size_t root) {
  int largest = root;
  int left = 2*root + 1;
  int right = 2*root + 2;

  // compare root to its children to see which is largest
  if (left < n && arr[left] > arr[largest]) largest = left;
  if (right < n && arr[right] > arr[largest]) largest = right;

  if (largest != root) {
    // then we are not done. Swap and heapify again
    swap(&arr[root], &arr[largest]);
    heapify(arr, n, largest);
  }
}
```

"""

# ╔═╡ 1ca8d8c9-fdd1-4172-aa75-9206c48e0212
md"""
## Quicksort

This naïve implementation is very simple. First, we check if the size of the array is less than 2 (in which case we know it is already sorted). This is our base case. Then, we set the pivot to be the first element of the array and compare it with every other element in the array. If the other element is smaller than the pivot, the pivot is moved and that element is placed immediately before the pivot.

```c
void quicksort(float* arr, size_t i, size_t n) {
  if ((n - i) < 2) return;

  // pivot is first item
  size_t pivot_i = i;

  for (size_t j = i+1; j < n; j++) {
    // check if other is less than pivot. If so, move pivot to the right and
    // place other immediately before pivot.
    if (arr[j] < arr[pivot_i]) {
      swap(&arr[pivot_i], &arr[pivot_i+1]);
      pivot_i++;

      // this case is important! This covers the instance where j immediately
      // follows the pivot, and so moving the pivot also swapped other before
      // the pivot.
      if(j != pivot_i) {
        swap(&arr[j], &arr[pivot_i-1]);
      }
    }
  }

  // recursively quicksort left and right halves
  quicksort(arr, i, pivot_i);
  quicksort(arr, pivot_i+1, n);
}
```

This has very poor performance when the list is already sorted or reverse-sorted (both common cases). To overcome this, we also implemented a version of quicksort with the median-of-three strategy for choosing the pivot. This version also uses insertion sort for `n <= 15`, which acts as our base-case.

As we will show later, these changes have significant performance implications.

```c
// recursive part of modified quicksort
void quicksort_modified(float* arr, size_t i, size_t n) {
  // insertion sort on small parts
  if (n < 16) {
    insertionSort(arr, n);
    return;
  }

  // the three indices for median-of-three
  size_t hi = n-1;
  size_t lo = i;
  size_t mid = lo + ((hi - lo) / 2);

  // these three if statements ensure elements at lo, mid, and hi are sorted in
  // relation to each other. This way we know that mid is our median-of-three.
  if (arr[mid] < arr[lo]) swap(&arr[mid], &arr[lo]);
  if (arr[hi] < arr[lo]) swap(&arr[hi], &arr[lo]);
  if (arr[mid] > arr[hi]) swap(&arr[mid], &arr[hi]);

  lo++;
  hi--;

  // check left side of pivot for anything that is larger than pivot
  while (lo < mid) {
    if (arr[lo] > arr[mid]) {
      swap(&arr[mid], &arr[mid-1]);
      mid--;

      if (lo != mid) swap(&arr[lo], &arr[mid+1]);
    }
    else lo++;
  }

  // check right side of pivot for anything smaller than pivot
  while (hi > mid) {
    if (arr[hi] < arr[mid]) {
      swap(&arr[mid], &arr[mid+1]);
      mid++;

      if (hi != mid) swap(&arr[hi], &arr[mid-1]);
    }
    else hi--;
  }

  // recursively quicksort left and right halves
  quicksort_modified(arr, i, mid);
  quicksort_modified(arr, mid+1, n);
}
```

"""

# ╔═╡ 72ad1767-88e1-457c-8a76-33f119c7cc15
md"""
# Complexity Analysis

Complexity for my specific implementations are covered below, with notes to explain key points. Except where indicated, all measures are time complexity.

Algorithm | Average-case | Worst-case | Best-case | Space complexity | Notes
----------|--------------|------------|-----------|------------------|-------
Insertion sort|Θ(n^2)|O(n^2)|Ω(n)|O(1)|Best-case is when list is already sorted
Merge sort|Θ(n*log(n))|O(n*log(n))|Ω(n*log(n))|O(n)|Space complexity is O(n) because of the buffer. Space complexity is the same in implementations where each recursive call creates new arrays that are merged and deleted later.
Heap sort|Θ(n*log(n))|O(n*log(n))|Ω(n*log(n))|O(1)|
Quick sort|Θ(n*log(n))|O(n^2)|Ω(n*log(n))|O(1)|Worst-case is triggered for both sorted and reverse-sorted lists.
Quick sort (modified)|Θ(n*log(n))|O(n^2)|Ω(n*log(n))|O(1)|Even though worst-case is technically still O(n^2), it's unlikely that we'll hit that case unless the range of values is very small. That would require all three items considered in the median to be equal and very small or very large compared to the rest of the items.

"""

# ╔═╡ 3a653422-4f34-42fe-a974-ff1a57d85a9b
md"""
# Results: Performance

To investigate performance, first an overview of the average-case performance of each algorithm is considered. After that, plots are generated which demonstrate special case performance and draw other interesting comparisons (like sequential vs parallel merge-sort).

Julia is used for data processing and visualization. Comments are used liberally to explain code.
"""

# ╔═╡ 3b969425-8ee1-4bfd-af35-27238bf7ae7b
# load csv data, process
begin
	CSV_FILE_NAME = "results-2021-06-07_2021.csv"

	df = CSV.File(CSV_FILE_NAME, normalizenames=true) |> DataFrame
	df.array_type

	# strip whitespace from columns with strings
	for col in [:sorting_algorithm, :array_type]
		# '@.' syntax applies the broadcast operator ('.') to all operations
		# on that line
		@. df[!, col] = strip(df[!, col])
	end

end

# ╔═╡ 48ffcfe0-f854-4fac-b071-897561e233cc
md"""
## Average-case Performance Comparison
"""

# ╔═╡ 0e259561-54b0-47cb-88f4-7663731d2191
begin
	# for now, only consider randomly-initialized arrays (the average case)
	df_randomarray = filter(:array_type => a -> strip(a) == "random", df)

	# and draw a line graph
	plot(
		df_randomarray.array_size,
		df_randomarray.avg_elements_per_second,
		group = df_randomarray.sorting_algorithm,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)", :log10),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ fb6d4070-922a-4bbe-84ef-b8524e883dbb
md"""
Insertion sort is dramatically slower than all the other algorithms (note the logarithmic scale!). It's so slow, in fact, that sorting a list of size $6.4e4$ elements took 4.1 seconds on my Ryzen 5 3600. By comparison, the second-slowest algorithm (heapsort) took less than 7 *milliseconds*.

Because each sort was repeated 10 times, insertion sort quickly became unreasonably slow, and thus it was only tested for arrays of size less than $1e5$.

Before moving on to individual algorithm performance, let's look at the same comparison with insertion sort removed.
"""

# ╔═╡ 7b64a513-938a-41ea-9722-1816d10bf9bd
begin
	df_randomarray_noinsertion = filter(:sorting_algorithm => s -> s != "insertion",
		df_randomarray)
	plot(
		df_randomarray_noinsertion.array_size,
		df_randomarray_noinsertion.avg_elements_per_second,
		group = df_randomarray_noinsertion.sorting_algorithm,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)", :log10),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ 73c15edb-3185-426b-9f7a-85619bd84df4
md"""
This gives us a much clearer picture of the performance of the sane algorithms. Naturally, performance decreases as array size increases and data must be loaded from RAM. However, heapsort drops off significantly more starting at around $50^6$. Also notable is the fact that mergesort and mergesort parallel seem to have nearly identical performance, as do quicksort and quicksort-modified. Let's hone in on those cases now.

## Mergesort vs Parallel Mergsort
"""

# ╔═╡ ad68f965-fa47-4e38-aad1-94ebe182c62f
begin
	# for now, only consider randomly-initialized arrays (the average case)
	df_randomarray_mergesort = filter(
		:sorting_algorithm => s -> s == "merge" || s == "merge_parallel",
		df_randomarray)

	plot(
		df_randomarray_mergesort.array_size,
		df_randomarray_mergesort.avg_elements_per_second,
		group = df_randomarray_mergesort.sorting_algorithm,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)"),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ 814dabd4-e6c3-42e4-8138-959c85e896ec
md"""
Even with the return to a linear scale, its impossible to tell the lines apart. Performance is effectively identical. Let's also inspect the raw data for the largest test.
"""

# ╔═╡ 9a1a3824-8a05-44c0-9c5c-8b7f1033f188
filter(:array_size => n -> n == 32768000,
	df_randomarray_mergesort)

# ╔═╡ cca83604-fbd6-4d57-80a3-c7023a2f8768
md"""
The average time taken within 0.1%, so it's safe to say the performance is effectively identical.

Some quick tests suggest that for very large arrays the parallel code does show speedup. But showing speedup is outside the scope of this project, so no further work will be done here.

For now, let's return to quicksort.

## Quicksort vs Modified Quicksort
"""

# ╔═╡ 96a438f8-a8b9-43f7-bf78-6847a5aa83b5
begin
	# for now, only consider randomly-initialized arrays (the average case)
	df_randomarray_quicksort = filter(
		:sorting_algorithm => s -> s == "quick" || s == "quick_modified",
		df_randomarray)

	plot(
		df_randomarray_quicksort.array_size,
		df_randomarray_quicksort.avg_elements_per_second,
		group = df_randomarray_quicksort.sorting_algorithm,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)"),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ dcbe0287-9253-4f8f-9682-dd4c4a967063
md"""
Here the performance does seem to be slightly better for modified quicksort, especially on smaller arrays. Let's inspect the raw data for a small and large case.
"""

# ╔═╡ e2fbdeb3-881a-4ccc-8a29-5f3bc5b4b03c
begin
	df_rq_short = filter(:array_size => n -> n == 32768000 || n == 32000,
		df_randomarray_quicksort)
	sort!(df_rq_short, [:array_size])
end

# ╔═╡ 2d15b628-660b-4cb8-9169-a1ee45eb82b6
md"""
Modified quicksort is slightly (~2% faster). This is statistically significant but not really helpful in the real world.

Our final comparison will investigate cases were the worst- or best- cases are significantly differnt from the average-case.

## Special Cases

### Sorting a Sorted List

All algorithms except naïve quicksort perform dramatically better on sorted lists than random lists. As a result, this case is usually unintersting and will only be considered for quicksort and insertion sort (where performance on a sorted list is `O(n)`, and therefore *especially* fast).

### Insertion Sort
"""

# ╔═╡ c2c99efa-5ec8-41e4-a97f-e0a100f690e5
begin
	# only looking at insertion sort
	df_insertion = filter(:sorting_algorithm => s -> s == "insertion", df)

	plot(
		df_insertion.array_size,
		df_insertion.avg_elements_per_second,
		group = df_insertion.array_type,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)", :log10),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ 8d26154e-f838-4bc2-a588-419f57ca78b0
md"""
When the list is already sorted, insertion performance is `O(n)`! However, when the list is reversed, performance is about half the average case.

### Quicksort
"""

# ╔═╡ e6fff4bd-5dad-4372-976a-55aecbe52faa
begin
	# only looking at quick sort
	df_quick = filter(:sorting_algorithm => s -> s == "quick", df)

	plot(
		df_quick.array_size,
		df_quick.avg_elements_per_second,
		group = df_quick.array_type,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)", :log10),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ a93735a6-f8b2-4ceb-84d5-bde4b89a929e
md"""
This graph clearly shows the dramatic reduction in performance on sorted (and reverse-sorted) lists. It was so slow that, like insertion sort, arrays larger than $6.4e4$ were not tested.

### Quicksort and Modified Quicksort
"""

# ╔═╡ 815e9742-c22f-404c-9ff0-b33c8b674173
begin
	# only looking at quick sort
	df_quick_m = filter(:sorting_algorithm => s -> s == "quick" || s == "quick_modified", df)

	plot(
		df_quick_m.array_size,
		df_quick_m.avg_elements_per_second,
		group = df_quick_m.array_type,
		line = :path,
		title = "Average-case sort performance, higher is better",
		xaxis = ("Number of elements", :log10),
		yaxis = ("Performance (elements/s)", :log10),
		bg = RGB(0.2, 0.2, 0.2)
	)
end

# ╔═╡ 0d687829-1c16-4cd4-9395-c4c248a72cd4
md"""
Our modified quicksort does perform much better on these sorted lists, though! Using the median-of-three strategy to choose a pivot effectively solves quicksort's biggest issue.
"""

# ╔═╡ Cell order:
# ╟─b2ca8e4b-fe95-43cb-a5bc-605f69d0d581
# ╟─aa573021-c79a-46c1-b4e0-021e4c3434ed
# ╟─e72161bc-dab5-4585-9cf8-bc32ec84e988
# ╟─3723574c-959a-4b04-91bf-6207c4ff6fd5
# ╟─209e2ecb-af90-4c48-a583-e9dda280198f
# ╟─fd11a60a-c3dd-4bb1-aeb3-e3be1b87f4f7
# ╟─1ca8d8c9-fdd1-4172-aa75-9206c48e0212
# ╟─72ad1767-88e1-457c-8a76-33f119c7cc15
# ╟─3a653422-4f34-42fe-a974-ff1a57d85a9b
# ╠═d6a054e2-c7bf-11eb-3732-c3135a35f956
# ╠═3b969425-8ee1-4bfd-af35-27238bf7ae7b
# ╟─48ffcfe0-f854-4fac-b071-897561e233cc
# ╠═0e259561-54b0-47cb-88f4-7663731d2191
# ╟─fb6d4070-922a-4bbe-84ef-b8524e883dbb
# ╠═7b64a513-938a-41ea-9722-1816d10bf9bd
# ╟─73c15edb-3185-426b-9f7a-85619bd84df4
# ╠═ad68f965-fa47-4e38-aad1-94ebe182c62f
# ╟─814dabd4-e6c3-42e4-8138-959c85e896ec
# ╠═9a1a3824-8a05-44c0-9c5c-8b7f1033f188
# ╟─cca83604-fbd6-4d57-80a3-c7023a2f8768
# ╠═96a438f8-a8b9-43f7-bf78-6847a5aa83b5
# ╟─dcbe0287-9253-4f8f-9682-dd4c4a967063
# ╠═e2fbdeb3-881a-4ccc-8a29-5f3bc5b4b03c
# ╟─2d15b628-660b-4cb8-9169-a1ee45eb82b6
# ╠═c2c99efa-5ec8-41e4-a97f-e0a100f690e5
# ╟─8d26154e-f838-4bc2-a588-419f57ca78b0
# ╠═e6fff4bd-5dad-4372-976a-55aecbe52faa
# ╟─a93735a6-f8b2-4ceb-84d5-bde4b89a929e
# ╠═815e9742-c22f-404c-9ff0-b33c8b674173
# ╟─0d687829-1c16-4cd4-9395-c4c248a72cd4


module LaxFriedrichs
export NumFlux
# Lax Friedrichs flux:
function NumFlux(F,a::Float64,b::Float64,alpha::Float64)
    0.5*(F.flux(a)+F.flux(b) - alpha*(b-a))
end
end


# Adapted from https://discourse.julialang.org/t/expression-parser/41880/7
# code by Alan R. Rogers, Professor of Anthropology, University of Utah

import Nemo: fmpq_mpoly

function myeval(e::Union{Expr,Symbol,Number}, map::Dict{Symbol,fmpq_mpoly})
    try
        return _myeval(e, map)
    catch ex
        println("Can't parse \"$e\"")
        rethrow(ex)
    end
end

function _myeval(s::Symbol, map::Dict{Symbol,fmpq_mpoly})
    if haskey(map, s)
        return map[s]
    else
        throw(UndefVarError(s))
    end
end

# Numbers are converted to fmpq.
function _myeval(x::Number, map::Dict{Symbol,fmpq_mpoly})
    return fmpq(x)
end

# To parse an expression, convert the head to a singleton
# type, so that Julia can dispatch on that type.
function _myeval(e::Expr, map::Dict{Symbol,fmpq_mpoly})
    return _myeval(Val(e.head), e.args, map)
end

# Call the function named in args[1]
function _myeval(::Val{:call}, args, map::Dict{Symbol,fmpq_mpoly})
    return _myeval(Val(args[1]), args[2:end], map)
end

# Addition
function _myeval(::Val{:+}, args, map::Dict{Symbol,fmpq_mpoly})
    x = 0
    for arg ∈ args
        x += _myeval(arg, map)
    end
    return x
end

# Subtraction and negation
function _myeval(::Val{:-}, args, map::Dict{Symbol,fmpq_mpoly})
    len = length(args)
    if len == 1
        return -_myeval(args[1], map)
    else
        return _myeval(args[1], map) - _myeval(args[2], map)
    end
end

# Multiplication
function _myeval(::Val{:*}, args, map::Dict{Symbol,fmpq_mpoly})
    x = 1
    for arg ∈ args
        x *= _myeval(arg, map)
    end
    return x
end

# Division
function _myeval(::Val{:/}, args, map::Dict{Symbol,fmpq_mpoly})
    return _myeval(args[1], map) / _myeval(args[2], map)
end


"""
    A = JopLeakyIntegration(spc; α=0.99, dim=1)

`A` is Leaky integration (i.e. a running sum) along dimension `dim`
of `spc::JetSpace{T,N}`.  The result is scaled by `α`.
"""

#=
From stanford: http://sep.stanford.edu/sep/prof/geelib/leakint.lop
    tt = 0.0

    if (adj)
        do i =  n, 1, -1 {
            tt = rho*tt + yy(i)
            xx(i) += tt
        }
    else
        do i =  1, n {
            tt = rho*tt + xx(i)
            yy(i) += tt
        }
    }
=#

function JopLeakyIntegration(spc::JetSpace{T}; α=0.99, dim=1) where {T}
    JopLn(dom = spc, rng = spc, df! = JopLeakyIntegration_df!, df′! = JopLeakyIntegration_df′!,
        s = (dim=dim, α=T(α)))
end

export JopLeakyIntegration

function JopLeakyIntegration_df!(d::AbstractArray{T}, m::AbstractArray{T}; dim, α, kwargs...) where {T}
    n1 = size(d,1)
    n2 = div(length(d),n1)
    _d = reshape(d,n1,n2)
    _m = reshape(m,n1,n2)
    for k2 = 1:n2
        tmp = T(0)
        for k1 = 1:n1
            tmp = α * tmp + _m[k1,k2]
            _d[k1,k2] = tmp
        end
    end
    d
end

function JopLeakyIntegration_df′!(m::AbstractArray{T}, d::AbstractArray{T}; dim, α, kwargs...) where {T}
    n1 = size(d,1)
    n2 = div(length(d),n1)
    _d = reshape(d,n1,n2)
    _m = reshape(m,n1,n2)
    for k2 = 1:n2
        tmp = T(0)
        for k1 = n1:-1:1
            tmp = α * tmp + _d[k1,k2]
            _m[k1,k2] = tmp
        end
    end
    m
end