text:
# This package is a work in progress for the estimation and identification of Vector Autoregressive (VAR) models.
# Credits:
# Kilian and Kim 2011, Cremfi codes, Geertler and Karadi 2015

module VectorAutoregressions
using Parameters, GrowableArrays, LinearAlgebra, Statistics
using Random
using Statistics: mean, std, quantile
using LinearAlgebra: I, cholesky, LowerTriangular, diag, eigvals, det

struct Intercept end

struct VAR
    mData::AbstractArray
    Y::AbstractArray
    X::AbstractArray
    β::AbstractArray
    ϵ::AbstractArray
    Σ::AbstractArray
    p::Int64
    inter::Intercept
    VAR(mData,Y,X,β,ϵ,Σ,p,inter) = p <= 0 ? error("Lag-length error: 'p' must be strictly positive") : new(mData,Y,X,β,ϵ,Σ,p,inter)
end

function VAR(y::AbstractArray,p::Int64,i::Bool)
    i == false ? ((mData,Y,X,β,ϵ,Σ,p) = fit(y,p)) : ((mData,Y,X,β,ϵ,Σ,p) = fit(y,p,Intercept()))
    return VAR(mData,Y,X,β,ϵ,Σ,p,Intercept())
end

abstract type  CIs end

struct IRFs
    IRF::AbstractArray
    CI::CIs
end

struct CIs_asy <: CIs
    CIl::AbstractArray
    CIh::AbstractArray
end

struct CIs_boot <: CIs
    CIl::AbstractArray
    CIh::AbstractArray
end

function IRFs_a(V::VAR,H::Int64,i::Bool)
    if i == true
        mVar1 = get_VAR1_rep(V,V.inter)
        mIRF = irf_chol(V, mVar1, H)
        mStd,mCov_Σ = irf_ci_asymptotic(V, H, V.inter)
    else
        mVar1 = get_VAR1_rep(V)
        mIRF = irf_chol(V, mVar1, H)
        mStd,mCov_Σ = irf_ci_asymptotic(V, H)
    end
    mCIl = mIRF - 1.96.*mStd
    mCIh = mIRF + 1.96.*mStd
    return IRFs(mIRF,CIs_asy(mCIl,mCIh))
end

function IRFs_b(V::VAR,H::Int64,nrep::Int64,i::Bool)
    if i == true
        mVar1 = get_VAR1_rep(V,V.inter)
        mIRF = irf_chol(V, mVar1, H)
        CI = irf_ci_bootstrap(V, H, nrep,V.inter)
    else
        mVar1 = get_VAR1_rep(V)
        mIRF = irf_chol(V, mVar1, H)
        CI = irf_ci_bootstrap(V, H, nrep)
    end
    return IRFs(mIRF,CI)
end

function IRFs_ext_instrument(V::VAR,Z::AbstractArray,H::Int64,nrep::Int64, α::AbstractArray, intercept::Bool)
    mIRF = irf_ext_instrument(V, Z, H, intercept)
    CI = irf_ci_wild_bootstrap(V, Z, H, nrep, α, intercept)
    return IRFs(mIRF,CI)
end

function IRFs_localprojection(z::AbstractArray{Float64}, p::AbstractArray{Int64}, H::Int64, A0inv::AbstractArray{Float64},cov_Σ::AbstractArray{Float64})
    T,K = size(z)
    vA0inv = vec(A0inv')
    mIRF = GrowableArray(copy(vA0inv))      # IRF for Horizon 0 --> use auxiliary model for identification
    cov_mIRF = zeros(K^2,1)
    for h = 1:H                          # IRF for Horizon 1~H
        ph = p[h]                     # lag-order for horizon h
        ys = z[ph+h:T,:]
        yt = z[ph:T-h,:]                 # RHS variable of interest: y(t)
        x = ones(T-ph-h+1,1)             # constant term
        for i = 2:ph
            x = [x z[ph+1-i:T-i-h+1,:]]  # other RHS lags: y(t-1)~y(t-p+1)
        end
        Mx = get_annhilation_matrix(x)    # annhilation matrix **
        β  = get_lp_beta(ys,yt,Mx)        # IRF by Local Projection **
        push!(mIRF, vec(A0inv'*β))      # structural IRF
        Σ_u = newey_west(ys,yt,Mx,β,h)
        invytMxyt  = inv(yt'*Mx*yt)
        Σ_β = kron(Σ_u, invytMxyt)                                 # var(vec(β))
        Σ_mIRF = kron(Matrix(1.0I,K,K) , A0inv')*Σ_β*kron(Matrix(1.0I,K,K) , A0inv')'    # var(vec(A0inv*β))
        cov_mIRF = [cov_mIRF reshape(diag(Σ_mIRF), K^2, 1)]
    end
    mIRF = Array(mIRF')
    mStd = sqrt.(cov_mIRF+cov_Σ)
    mCIl = mIRF - 1.96.*mStd
    mCIh = mIRF + 1.96.*mStd
    return IRFs(mIRF,CIs_asy(mCIl, mCIh)), mStd
end

function IRFs_localprojection(z::AbstractArray{Float64}, p::AbstractArray{Int64}, H::Int64)
    T,K = size(z)
    vA0inv = vec(eye(K))
    mIRF = GrowableArray(copy(vA0inv))      # IRF for Horizon 0 --> reduce form
    cov_mIRF = zeros(K^2,1)
    for h = 1:H      # IRF for Horizon 1~H
        ph = p[h]                     # lag-order for horizon h
        ys = z[ph+h:T,:]
        yt = z[ph:T-h,:]                 # RHS variable of interest: y(t)
        x = ones(T-ph-h+1,1)             # constant term
        for i = 2:ph
            x = [x z[ph+1-i:T-i-h+1,:]]  # other RHS lags: y(t-1)~y(t-p+1)
        end
        Mx = get_annhilation_matrix(x)    # annhilation matrix **
        β  = get_lp_beta(ys,yt,Mx)        # IRF by Local Projection **
        push!(mIRF, vec(β))                      # reduced form IRF
        Σ_u = newey_west(ys,yt,Mx,β,h)
        invytMxyt  = inv(yt'*Mx*yt)
        Σ_β = kron(Σ_u, invytMxyt)                               # var(vec(β))
        cov_mIRF = [cov_mIRF reshape(diag(Σ_β), K^2, 1)]
    end
    mIRF = mIRF'
    mStd = sqrt.(cov_mIRF)
    mCIl = mIRF - 1.96.*mStd
    mCIh = mIRF + 1.96.*mStd
    return IRFs(mIRF,CIs_asy(mCIl, mCIh)), mStd
end

function IRFs_localprojection(z::AbstractArray{Float64}, p::Int64, H::Int64, A0inv::AbstractArray{Float64},cov_Σ::AbstractArray{Float64})
    T,K = size(z)
    vA0inv = vec(A0inv')
    mIRF = GrowableArray(copy(vA0inv))      # IRF for Horizon 0 --> use auxiliary model for identification
    cov_mIRF = zeros(K^2,1)
    for h = 1:H           # IRF for Horizon 1~H
        ph = p                     # lag-order for horizon h
        ys = z[ph+h:T,:]
        yt = z[ph:T-h,:]                 # RHS variable of interest: y(t)
        x = ones(T-ph-h+1,1)             # constant term
        for i = 2:ph
            x = [x z[ph+1-i:T-i-h+1,:]]  # other RHS lags: y(t-1)~y(t-p+1)
        end
        Mx = get_annhilation_matrix(x)    # annhilation matrix **
        β  = get_lp_beta(ys,yt,Mx)        # IRF by Local Projection **
        push!(mIRF, vec(A0inv'*β))      # structural IRF
        Σ_u = newey_west(ys,yt,Mx,β,h)
        invytMxyt  = inv(yt'*Mx*yt)
        Σ_β = kron(Σ_u, invytMxyt)                               # var(vec(β))
        Σ_mIRF = kron(Matrix(1.0I,K,K) , A0inv')*Σ_β*kron(Matrix(1.0I,K,K) , A0inv')'    # var(vec(A0inv*β))
        cov_mIRF = [cov_mIRF reshape(diag(Σ_mIRF), K^2, 1)]
    end
    mIRF = mIRF'
    mStd = sqrt.(cov_mIRF+cov_Σ)
    mCIl = mIRF - 1.96.*mStd
    mCIh = mIRF + 1.96.*mStd
    return IRFs(mIRF,CIs_asy(mCIl, mCIh)), mStd
end

function IRFs_localprojection(z::AbstractArray{Float64}, p::Int64, H::Int64)
    T,K = size(z)
    vA0inv = vec(eye(K))
    mIRF = GrowableArray(copy(vA0inv))      # IRF for Horizon 0 --> reduce form
    cov_mIRF = zeros(K^2,1)
    for h = 1:H     # IRF for Horizon 1~H
        ph = p                     # lag-order for horizon h
        ys = z[ph+h:T,:]
        yt = z[ph:T-h,:]                 # RHS variable of interest: y(t)
        x = ones(T-ph-h+1,1)             # constant term
        for i = 2:ph
            x = [x z[ph+1-i:T-i-h+1,:]]  # other RHS lags: y(t-1)~y(t-p+1)
        end
        Mx = get_annhilation_matrix(x)    # annhilation matrix **
        β  = get_lp_beta(ys,yt,Mx)        # IRF by Local Projection **
        push!(mIRF, vec(β))                      # reduced form IRF
        Σ_u = newey_west(ys,yt,Mx,β,h)
        invytMxyt  = inv(yt'*Mx*yt)
        Σ_β = kron(Σ_u, invytMxyt)                               # var(vec(β))
        cov_mIRF = [cov_mIRF reshape(diag(Σ_β), K^2, 1)]
    end
    mIRF = mIRF'
    mStd = sqrt.(cov_mIRF)
    mCIl = mIRF - 1.96.*mStd
    mCIh = mIRF + 1.96.*mStd
    return IRFs(mIRF,CIs_asy(mCIl, mCIh)), mStd
end

function newey_west(ys::AbstractArray,yt::AbstractArray,Mx::AbstractArray,β::AbstractArray,h::Int64)
    T,K = size(yt)
    u = Mx*ys - Mx*yt*β    # residual from LPs
    μ_u = zeros(1,K)
    for i = 1:K
        μ_u[:,i] .= mean(u[:,i])
    end
    iu = size(u,1)
    u0 = u-kron(ones(iu,1),μ_u)
    ρ0 = (u0'*u0)./iu        # sample cov(u)
    Σ_u = ρ0
    M = h    # truncation point
    for j = 1:M
        Rj = (u0[1:end-j,:]'*u0[j+1:end,:])./(size(u,1))   # R(j)
        Σ_u = Σ_u + (1-j/(M+1))*(Rj+Rj')                 # varcov(u)
    end
    return Σ_u
end

function get_lp_component(z::AbstractArray,p::Int64,H::Int64)
    T,K = size(z)
    ys = z[p+1:T,:]
    yt = z[p:T-1,:]                  # RHS variable of interest: y(t)
    x = ones(T-p,1)                  # constant term
    for j=2:p
        x = [x z[p+1-j:T-j,:]]       # other RHS variables: y(t-1)~y(t-p+1)
    end
    return ys,yt,x
end

function lp_estimator(ys::AbstractArray,yt::AbstractArray,x::AbstractArray,t::Int64)
    Mx = get_annhilation_matrix(x)    # annhilation matrix **
    β  = get_lp_beta(ys,yt,Mx)        # IRF by Local Projection **
    u  = get_lp_residual(ys,yt,Mx,β)  # residual from LPs
    Σ  = get_variance_estimator(u,t)  # variance of errors
    return β, u, Σ
end

get_variance_estimator(u::AbstractArray,t::Int64) = u'u/t
get_annhilation_matrix(x) = size(x,1) |> λ -> eye(λ) - x/(x'*x)*x'
get_lp_beta(ys::AbstractArray,yt::AbstractArray,Mx::AbstractArray) = (yt'*Mx*yt)\(yt'*Mx*ys)
get_lp_residual(ys::AbstractArray,yt::AbstractArray,Mx::AbstractArray,β::AbstractArray) = Mx*ys - Mx*yt*β

function lp_lagorder(z::AbstractArray,pbar::Int64,H::Int64,ic::String)
    T,K = size(z)
    t     = T-pbar
    vIC  = Array{Int64}(undef, H)
    Ys,Yt,X = get_lp_component(z,pbar,H)
    for j = 1:H                                  # loop for horizon h of IRF
        IC = zeros(pbar,1)                       # the vector of AIC(p)
        ys = Ys[j:end,:]                         # dependent variable: yt+s
        yt = Yt[1:end-j+1,:]                     # regressor of interest: yt
        for m = 1:pbar                           # loop for lag order selection
            x = X[1:end-j+1,1:(m-1)*K+1]         # other independent variables
            β, u, Σ = lp_estimator(ys,yt,x,t)
            if ic == "aic"
                IC[m]  = log(det(Σ)) + 2*(K^2*m)/t                      # AIC statistic
            elseif ic == "bic"
                IC[m]  = log(det(Σ)) + (K^2*m)*log(t)/t                 # SIC statistic
            elseif ic == "aicc"
                b = t/(t-(m*K+K+1))
                IC[m] = t*(log(det(Σ))+K) + 2*b*(K^2*m+K*(K+1)/2)       # HURVICH AND TSAI
            elseif ic == "hqc"
                IC[m]  = log(det(Σ)) + 2*log(log(t))*K^2*m/t            # HQ
            elseif error("ic must be aic, bic, aicc or hqc")
            end
        end
        _, ind = findmin(vec(IC))
        vIC[j] = Int(ind)
    end
    return vIC
end

function fit(y::AbstractArray,p::Int64)
    (T,K) = size(y)
    T < K && error("error: there are more covariates than observation")
    X = y
    y = transpose(y)
    Y = y[:,p+1:T]
    X = lagmatrix(X,p)'
    β = (Y*X')/(X*X')
    ϵ = Y - β*X
    Σ = ϵ*ϵ'/(T-p-p*K-1)
    return y',Y,X,β,ϵ,Σ,p
end

function fit(y::AbstractArray,p::Int64,inter::Intercept)
    (T,K) = size(y)
    T < K && error("error: there are more covariates than observation")
    X = y
    y = transpose(y)
    Y = y[:,p+1:T]
    X = lagmatrix(X,p,inter)'
    β = (Y*X')/(X*X')
    ϵ = Y - β*X
    Σ = ϵ*ϵ'/(T-p-p*K-1)
    return y',Y,X,β,ϵ,Σ,p
end

function lagmatrix(x::AbstractArray,p::Int64,inter::Intercept)
    sk = 1
    T, K = size(x)
    k    = K*p+1
    idx  = repeat(1:K, p)
    X    = Array{eltype(x)}(undef, (T-p, k))
    # building X (t-1:t-p) allocating data from D matrix - avoid checking bounds
    for j = 1+sk:(sk+K*p)
        for i = 1:(T-p)
            lg = round(Int, ceil((j-sk)/K)) - 1 # create index [0 0 1 1 2 2 ...etc]
            @inbounds X[i, j] = x[i+p-1-lg, idx[j-sk]]
        end
    end
    for j=1:T-p
        @inbounds X[j,1] = 1.0
    end
    return X
end

function lagmatrix(x::AbstractArray,p::Int64)
    sk = 1
    T, K = size(x)
    k    = K*p+1
    idx  = repeat(1:K, p)
    X    = Array{eltype(x)}(undef, (T-p, k))
    # building X (t-1:t-p) allocating data from D matrix - avoid checking bounds
    for j = 1+sk:(sk+K*p)
        for i = 1:(T-p)
            lg = round(Int, ceil((j-sk)/K)) - 1 # create index [0 0 1 1 2 2 ...etc]
            @inbounds X[i, j] = x[i+p-1-lg, idx[j-sk]]
        end
    end
    return X[:,2:end]
end

function lagmatrix(x::AbstractVector,p::Int64)
    sk = 1
    T = length(x)
    K = 1
    k    = K*p+1
    idx  = repeat(1:K, p)
    X    = Array{eltype(x)}(undef, (T-p, k))
    # building X (t-1:t-p) allocating data from D matrix - avoid checking bounds
    for j = 1+sk:(sk+K*p)
        for i = 1:(T-p)
            lg = round(Int, ceil((j-sk)/K)) - 1 # create index [0 0 1 1 2 2 ...etc]
            @inbounds X[i, j] = x[i+p-1-lg, idx[j-sk]]
        end
    end
    return X[:,2:end]
end

function commutation(n::Int64, m::Int64)
    # returns Magnus and Neudecker's commutation matrix of dimensions n by m
    k = reshape(kron(vec(Matrix(1.0I,n,n)), Matrix(1.0I,m,m)), n*m, n*m)
    return k
end

function duplication(n::Int64)
    # Returns Magnus and Neudecker's duplication matrix of size n
    # VERY AMBIGUOUS FUNC
    a = LowerTriangular(ones(n, n))
    inds = (a .== 1) # find(a)::Vector{Int64}
    a[inds] .= 1:Int(sum(inds))
    aT = transpose(a) - I
    for i in 1:n
        aT[i, i] = 0
    end
    a = a + aT
    j = Int.(vec(a)) #convert(Vector{Int64}, vec(a))::Vector{Int64}
    m = trunc.(Int,(n*(n+1)/2))::Int64
    d = zeros(n*n,m)
    for r = 1:size(d,1)
        r, j[r]
        d[r, j[r]] = 1.0
    end
    return d
end

function elimat(m::Int64)
    # elimat(m) returns the elimination matrix Lm
    # The elimination matrix Lm is for any matrix F, Vech(F)=Lm Vec(F)
    A = Matrix(1.0I,m^2,m^2)
    L = A[1:m,:]
    for n in 2:m
        S = A[m*(n-1)+1:n*m,:]
        S = S[n:end,:]
        L = [L;S]
    end
    return L
end

function var_lagorder(z::AbstractArray,pbar::Int64,ic::String)
    T,K = size(z)::Tuple{Int64,Int64}
    t   = Int(T-pbar)
    IC  = zeros(pbar,1)
    Y = z[pbar+1:T,:]                           # dependent variable
    for p = 1:pbar
        X = ones(Int(t),1)
        for i = 1:p
            X = [X z[pbar+1-i:T-i,:]]            # construct lagged regressors
        end
        β  = (X'*X)\(X'*Y)                    # estimate by multivariate LS **
        u     = Y-X*β                         # errors
        Σ   = (u'*u/t)::AbstractArray{Float64}
        if ic == "aic"                           # variance of errors
            IC[p]     = log(det(Σ))+2*p*K^2/t                      # AIC statistic
        elseif ic == "bic"
            IC[p]   = log(det(Σ))+(K^2*p)*log(t)/t                # SIC statistic
        elseif ic == "aicc"
            b = t/(t-(p*K+K+1))                    # COEFFICIENT FOR AICC
            IC[p] = t*(log(det(Σ))+K)+2*b*(K^2*p+K*(K+1)/2)     # AICC
        elseif ic == "hqc"
            IC[p]   = log(det(Σ))+2*log(log(t))*K^2*p/t # HQC
        elseif error("ic must be aic, bic, aicc or hqc")
        end
    end
    _, length_ic = findmin(vec(IC))
    println("Using $ic the best lag-length is $length_ic")
    return Int(length_ic)
end

function get_VAR1_rep(V::VAR)
    K = size(V.Σ,1)
    B = vcat(V.β, hcat(eye(K*(V.p-1)), zeros(K*(V.p-1),K))::AbstractArray{Float64,2})::AbstractArray{Float64,2}
    B = convert(Array{Float64,2},B)
end

function get_VAR1_rep(V::VAR,inter::Intercept)
    K = size(V.Σ,1)
    # v = [V.β[:,1]; zeros(K*(V.p-1),1)]
    B = vcat(V.β[:,2:end], hcat(eye(K*(V.p-1)), zeros(K*(V.p-1),K)))::AbstractArray{Float64,2}
    B = convert(Array{Float64,2},B)
end

function irf_ci_asymptotic(V::VAR, H::Int64)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    SIGa = kron(inv(V.X*V.X'/(T-V.p)),V.Σ)
    # Calculation of stdev follows Lutkepohl(2005) p.111,93
    A = get_VAR1_rep(V)
    A0inv = cholesky(V.Σ).L
    STD   = zeros(K^2,H+1)
    COV2   = zeros(K^2,H+1)
    J = [Matrix(1.0I,K,K)  zeros(K,K*(V.p-1))]
    L = elimat(K)
    Kk = commutation(K,K)
    Hk = L'/(L*(Matrix(1.0I,K^2,K^2) +Kk)*kron(A0inv,Matrix(1.0I,K,K) )*L')
    D = duplication(K)
    Dplus = (D'*D)\D'
    SIGsig = 2*Dplus*kron(V.Σ,V.Σ)*Dplus'
    Cbar0 = kron(Matrix(1.0I,K,K) ,J*Matrix(1.0I,K*V.p,K*V.p) *J')*Hk
    STD[:,1] = vec((reshape(diag(real(sqrt.(complex(Cbar0*SIGsig*Cbar0'/(T-V.p))))),K,K))')
    COV2[:,1] = vec((reshape(diag((Cbar0*SIGsig*Cbar0'/(T-V.p))),K,K))')
    for h=1:H
        Gi = zeros(K^2,K^2*V.p)
        for m=0:(h-1)
            Gi += kron(J*(A')^(h-1-m),J*(A^m)*J')
        end
        C = kron(A0inv',Matrix(1.0I,K,K) )*Gi
        Cbar = kron(Matrix(1.0I,K,K) ,J*A^h*J')*Hk
        STD[:,h+1] = vec((reshape(diag(real(sqrt.(complex(C*SIGa*C'+Cbar*SIGsig*Cbar')/(T-V.p)))),K,K))')
        COV2[:,h+1] = vec((reshape(diag(((Cbar*SIGsig*Cbar')/(T-V.p))),K,K))')
    end
    return STD,COV2
end

function irf_ci_asymptotic(V::VAR, H::Int64, inter::Intercept)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    SIGa = kron(inv(V.X*V.X'/(T-V.p)),V.Σ)
    SIGa = SIGa[K+1:end,K+1:end]
    # Calculation of stdev follows Lutkepohl(2005) p.111,93
    A = get_VAR1_rep(V,V.inter)
    A0inv = cholesky(V.Σ).L
    STD   = zeros(K^2,H+1)
    COV2   = zeros(K^2,H+1)
    J = [Matrix(1.0I,K,K)  zeros(K,K*(V.p-1))]
    L = elimat(K)
    Kk = commutation(K,K)
    Hk = L'/(L*(Matrix(1.0I,K^2,K^2) +Kk)*kron(A0inv,Matrix(1.0I,K,K) )*L')
    D = duplication(K)
    Dplus = (D'*D)\D'
    SIGsig = 2*Dplus*kron(V.Σ,V.Σ)*Dplus'
    Cbar0 = kron(Matrix(1.0I,K,K) ,J*Matrix(1.0I,K*V.p,K*V.p)*J')*Hk
    STD[:,1] = vec((reshape(diag(real(sqrt.(complex(Cbar0*SIGsig*Cbar0'/(T-V.p))))),K,K))')
    COV2[:,1] = vec((reshape(diag((Cbar0*SIGsig*Cbar0'/(T-V.p))),K,K))')
    for h=1:H
        Gi = zeros(K^2,K^2*V.p)
        for m=0:(h-1)
            Gi += kron(J*(A')^(h-1-m),J*(A^m)*J')
        end
        C = kron(A0inv',Matrix(1.0I,K,K) )*Gi
        Cbar = kron(Matrix(1.0I,K,K) ,J*A^h*J')*Hk
        STD[:,h+1] = vec((reshape(diag(real(sqrt.(complex(C*SIGa*C'+Cbar*SIGsig*Cbar')/(T-V.p)))),K,K))')
        COV2[:,h+1] = vec((reshape(diag(((Cbar*SIGsig*Cbar')/(T-V.p))),K,K))')
    end
    return STD,COV2
end

get_boot_init_int_draw(T::Int64,p::Int64) = Int64(trunc.((T-p+1)*rand()+1))
get_boot_init_vector_draw(T::Int64,p::Int64) = Array{Int64}(trunc.((T-p)*rand(T-p) .+ 1))

function build_sample(V::VAR)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    # Draw block of initial pre-sample values
    y = zeros(K,T)                             # bootstrap data
    u = zeros(K,T)                             # bootstrap innovations
    iDraw = get_boot_init_int_draw(T,V.p)      # position of initial draw
    y[:,1:V.p] = V.Y[:,iDraw:iDraw+V.p-1]      # drawing pre-sample obs
    # Draw innovations
    vDraw = get_boot_init_vector_draw(T,V.p)   # index for innovation draws
    u[:, V.p+1:T] = V.ϵ[:,vDraw]               # drawing innovations
    @inbounds for i = V.p+1:T
        y[:,i] = u[:,i]
        for j =  1:V.p
            y[:,i] += V.β[:,(j-1)*K + 1:j*K]*y[:,i-j]
        end
    end
    return y
end

function build_sample(V::VAR,inter::Intercept)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    # Draw block of initial pre-sample values
    y = zeros(K,T)                             # bootstrap data
    u = zeros(K,T)                             # bootstrap innovations
    iDraw = get_boot_init_int_draw(T,V.p)      # position of initial draw
    y[:,1:V.p] = V.Y[:,iDraw:iDraw+V.p-1]      # drawing pre-sample obs
    # Draw innovations
    vDraw = get_boot_init_vector_draw(T,V.p)   # index for innovation draws
    u[:, V.p+1:T] = V.ϵ[:,vDraw]               # drawing innovations
    @inbounds for i = V.p+1:T
        y[:,i] = V.β[:,1] + u[:,i]
        for j =  1:V.p
            y[:,i] += V.β[:,(j-1)*K + 2:j*K+1]*y[:,i-j]
        end
    end
    return y
end

col_mean(x::AbstractArray) = mean(x, dims = 2)
test_bias_correction(x::AbstractArray) =  any(abs.(eigvals(x)).<1)

function get_boot_ci(V::VAR,H::Int64,nrep::Int64, bDo_bias_corr::Bool,inter::Intercept)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    mIRFbc = zeros(nrep, K^2*(H+1))
    @inbounds for j = 1:nrep
        # Recursively construct sample
        yr = build_sample(V,V.inter)
        yr = Array((yr .- col_mean(yr))')  # demean yr bootstrap data
        #pr = V.p # also using lag length selection
        Vr = VAR(yr,V.p,true)
        # Bias correction: if the largest root of the companion matrix
        # is less than 1, do BIAS correction
        mVar1 = get_VAR1_rep(Vr,V.inter)::AbstractArray{Float64,2}
        bBias_corr_test = test_bias_correction(mVar1)
        if all([bDo_bias_corr, bBias_corr_test])
            mVar1 = bias_correction(Vr,mVar1)
        end
        mIRF = irf_chol(Vr,mVar1,H)
        mIRFbc[j,:] = vec(mIRF')'
    end                     # end bootstrap
    return mIRFbc
end

function get_boot_ci(V::VAR,H::Int64,nrep::Int64, bDo_bias_corr::Bool)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    mIRFbc = zeros(nrep, K^2*(H+1))
    @inbounds for j = 1:nrep
        # Recursively construct sample
        yr = build_sample(V)
        yr = Array((yr .- col_mean(yr))')  # demean yr bootstrap data
        #pr = V.p # also using lag length selection
        Vr = VAR(yr,V.p,false)
        # Bias correction: if the largest root of the companion matrix
        # is less than 1, do BIAS correction
        mVar1 = get_VAR1_rep(Vr)::AbstractArray{Float64,2}
        bBias_corr_test = test_bias_correction(mVar1)
        if all([bDo_bias_corr, bBias_corr_test])
            mVar1 = bias_correction(Vr,mVar1)
        end
        mIRF = irf_chol(Vr,mVar1,H)
        mIRFbc[j,:] = vec(mIRF')'
    end                     # end bootstrap
    return mIRFbc
end

function get_companion_vcv(V::VAR)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    mSigma = vcat(hcat(V.Σ, zeros(K,K*V.p-K)), zeros(K*V.p-K,K*V.p))::AbstractArray{Float64,2}
    mSigma = convert(Array{Float64,2},mSigma)
end

function get_sigma_y(V::VAR, mVar1::AbstractArray, mSigma::AbstractArray)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    vSigma = (eye((K*V.p)^2) - kron(mVar1,mVar1)) \ vec(mSigma)    # Lutkepohl p.29 (2.1.39)
    mSigma_y = reshape(vSigma, K*V.p::Int64, K*V.p::Int64)
    return convert(Array{Float64,2},mSigma_y)
end

function get_bias(mSigma::AbstractArray,mSigma_y::AbstractArray,B::AbstractArray,I::AbstractArray,mSum_eigen::AbstractArray)
    return mBias = mSigma*(inv(I - B) + B/(I-B^2) + mSum_eigen)/(mSigma_y)
end

function bias_correction(V::VAR,mVar1::AbstractArray)
    # Bias-correction Pope (1990)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    mSigma = get_companion_vcv(V)
    mSigma_y = get_sigma_y(V,mVar1,mSigma)
    II = Matrix(1.0I,K*V.p, K*V.p)
    B = Array(mVar1')
    vEigen = eigvals(mVar1)
    global mSum_eigen = zeros(K*V.p,K*V.p)
    for h = 1:K*V.p
        mSum_eigen += vEigen[h].\(II - vEigen[h]*B)
    end
    mBias = get_bias(mSigma,mSigma_y,B,II,mSum_eigen)
    mAbias = -mBias/T
    return mBcA = get_proportional_bias_corr(mVar1,mAbias)
end

function get_proportional_bias_corr(mVar1::AbstractArray,Abias::AbstractArray;iBc_stab::Int64 = 9, iδ::Int64 = 1)
    mBcA = zeros(size(mVar1))
    while iBc_stab >= 1
        # Adjust bias-correction proportionately
        mBcA = real(mVar1-iδ*Abias)
        bcmod = abs.(eigvals(mBcA))
        any(bcmod.>= 1) ?  iBc_stab = 1 : iBc_stab = 0
        iδ += -0.01
        iδ <= 0 && (iBc_stab = 0)
    end
    return mBcA
end

function get_boot_conf_interval(mIRFbc::AbstractArray,H::Int64,K::Int64)
    N = size(mIRFbc,2)
    mCILv = zeros(1,N)
    mCIHv = zeros(1,N)
    for i = 1:N
        mCILv[:,i] .= quantile(vec(mIRFbc[:,i]),0.025)
        mCIHv[:,i] .= quantile(vec(mIRFbc[:,i]),0.975)
    end
    mCIL  = reshape(mCILv',H+1,K^2)'
    mCIH  = reshape(mCIHv',H+1,K^2)'
    return mCIL, mCIH
end

function irf_ci_bootstrap(V::VAR, H::Int64, nrep::Int64; bDo_bias_corr::Bool=true)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    iScale_ϵ = sqrt((T-V.p)/(T-V.p-K*V.p-1))
    u = V.ϵ*iScale_ϵ   # rescaling residual (Stine, JASA 1987)
    mIRFbc = get_boot_ci(V,H,nrep,bDo_bias_corr)
    # Calculate 95 perccent interval endpoints
    mCIL, mCIH = get_boot_conf_interval(mIRFbc::AbstractArray,H::Int64,K::Int64)
    return CIs_boot(mCIL,mCIH)
end

function irf_ci_bootstrap(V::VAR, H::Int64, nrep::Int64, inter::Intercept; bDo_bias_corr::Bool=true)
    K,T = size(V.Y)::Tuple{Int64,Int64}
    iScale_ϵ = sqrt((T-V.p)/(T-V.p-K*V.p-1))
    u = V.ϵ*iScale_ϵ   # rescaling residual (Stine, JASA 1987)
    mIRFbc = get_boot_ci(V,H,nrep,bDo_bias_corr,V.inter)
    # Calculate 95 perccent interval endpoints
    mCIL, mCIH = get_boot_conf_interval(mIRFbc::AbstractArray,H::Int64,K::Int64)
    return CIs_boot(mCIL,mCIH)
end

function irf_chol(V::VAR, mVar1::AbstractArray, H::Int64)
    K = size(V.Σ,1)
    mSigma = cholesky(V.Σ).L::LowerTriangular  # Cholesky or reduced form
    J = [eye(K,K) zeros(K,K*(V.p-1))]
    mIRF = zeros(K^2,H+1)
    mIRF[:,1] = reshape((J*mVar1^0*J'*mSigma)',K^2,1)
    for i = 1:H
        mIRF[:,i+1] = reshape((J*mVar1^i*J'*mSigma)',K^2,1)
    end
    return mIRF
end

function irf_reduce_form(V::VAR, mVar1::AbstractArray, H::Int64)
    K = size(V.Σ,1)
    mSigma = Matrix(1.0I,K,K)
    J = [Matrix(1.0I,K,K) zeros(K,K*(V.p-1))]
    mIRF = zeros(K^2,H+1)
    mIRF[:,1] = reshape((J*mVar1^0*J'*mSigma)',K^2,1)
    for i = 1:H
        mIRF[:,i+1] = reshape((J*mVar1^i*J'*mSigma)',K^2,1)
    end
    return mIRF
end

function irf_ext_instrument(V::VAR,Z::AbstractArray,H::Int64,intercept::Bool)
    # Version improved by G. Ragusa
    y,B,Σ,U,p = V.Y',V.β,V.Σ,V.ϵ,V.p
    (T,K) = size(y)
    (T_z, K_z) = size(Z)
    ZZ = Z[p+1:end,:]
    ΣηZ = U[:,T-T_z+p+1:end]*ZZ
    Ση₁Z = ΣηZ[1:1,:]
    H1 = ΣηZ*Ση₁Z'./(Ση₁Z*Ση₁Z')
    A0inv = [H1 zeros(K,K-1)]
    A = [B[:,2:end];[Matrix(1.0I,K*(p-1),K*(p-1)) zeros(K*(p-1),K)]]
    J = [Matrix(1.0I,K,K) zeros(K,K*(p-1))]
    IRF = GrowableArray(A0inv[:,1])
    #HD = GrowableArray(zeros(K,K))
    for h in 1:H
        C = J*A^h*J'
        push!(IRF, (C*A0inv)[:,1])
    end
    return IRF
end

function irf_ext_instrument(V::VAR,Z::Vector{Float64},H::Int64,intercept::Bool)
    # Version improved by G. Ragusa
    y,B,Σ,U,p = V.Y',V.β,V.Σ,V.ϵ,V.p
    (T,K) = size(y)
    T_z = size(Z)
    ZZ = Z[p+1:end,:]
    ΣηZ = U[:,T-T_z+p+1:end]*ZZ
    Ση₁Z = ΣηZ[1:1,:]
    H1 = ΣηZ*Ση₁Z'./(Ση₁Z*Ση₁Z')
    A0inv = [H1 zeros(K,K-1)]
    A = [B[:,2:end];[Matrix(1.0I,K*(p-1),K*(p-1)) zeros(K*(p-1),K)]]
    J = [Matrix(1.0I,K,K) zeros(K,K*(p-1))]
    IRF = GrowableArray(A0inv[:,1])
    #HD = GrowableArray(zeros(K,K))
    for h in 1:H
        C = J*A^h*J'
        push!(IRF, (C*A0inv)[:,1])
    end
    return IRF
end

function irf_ci_wild_bootstrap(V::VAR,Z::AbstractArray,H::Int64,nrep::Int64,α::AbstractArray,intercept::Bool)
    # Wild Bootstrap
    # Version improved by G. Ragusa
    y,Y,A,u,p = V.mData,V.Y',V.β,V.ϵ,V.p
    count = 1
    (T,K) = size(y)
    (T_z, K_z) = size(Z)
    IRFS = GrowableArray(Matrix{Float64}(undef, H+1, K))
    CILv = Array{Float64}(undef, (length(α), size(IRFS,2)))
    CIHv = similar(CILv)
    lower_bound = Array{Float64}(undef, (H+1, length(α)))
    upper_bound = similar(lower_bound)
    res = u'
    oneK = ones(1,K)
    oneKz = ones(1,K_z)
    varsb = zeros(T,K)
    if intercept == true
        Awc = A[:,2:end]
        Ac  = A[:,1]
    else
        Awc = A
        Ac = zeros(A[:,1])
    end
    rr  = Array{Int16}(undef, T,1)
    while count < nrep+1
        rand!(rr, [-1, 1])
        resb = res.*rr[p+1:end]
        #  Deterministic initial values
        for j in 1:K
            for i in 1:p
                @inbounds varsb[i,j] = y[i,j]
            end
        end
        @inbounds for j = p+1:T
            lvars = transpose(varsb[j-1:-1:j-p,:])
            varsb[j,:] = Awc*vec(lvars) + Ac + resb[j-p,:];
        end
        proxies = Z.*rr[T-T_z+1:end]
        Vr = VAR(varsb,V.p,true)
        IRFr = irf_ext_instrument(Vr,proxies,H,intercept)
        #IRFrmat[count, :, :] = vec(IRFr')'
        push!(IRFS, IRFr)
        count += 1
    end
    CIl = Array{Float64}(undef, (length(α), H+1, K))
    CIh = similar(CIl)
    for i in 1:K
        # FIX THIS POINT--AT THE MOMENT ONLY FIRST VARIABLE CI
        lower = mapslices(u->quantile(u, α./2), IRFS[2:end,:,i],dims=1)'
        upper = mapslices(u->quantile(u, 1 .- α./2), IRFS[2:end,:,i],dims=1)'
        CIl[:,:,i] = lower'
        CIh[:,:,i] = upper'
    end
    return CIs_boot(CIl, CIh)
end

function t_test(V::VAR)
    K = size(V.Σ,1)
    Kp = size(V.X,1)
    H = kron(inv(V.X*V.X'),V.Σ)
    SE = sqrt.(diag(H))
    T = reshape(vec(V.β)./SE,K,Kp)
    return T
end

function gen_var1_data!(y::AbstractArray,B::AbstractArray,Σ::AbstractArray,burnin::Int64)
    T,K = size(y)
    for j = 2:T
        y[j,:] = B*y[j-1,:] + Σ*randn(K,1)
    end
    y = y[burnin+1:end,:]
    return y .- mean(y, dims = 1)
end

eye(n) = float.(I(n))

function eye(n, m)
    out = zeros(n, m)
    n = min(n, m)
    out[1:n, 1:n] .= eye(n)
    out
end

export VAR, IRFs_a, IRFs_b, IRFs_ext_instrument, IRFs_localprojection, gen_var1_data!
export lp_lagorder,var_lagorder, irf_ci_asymptotic

end # end of the module

max_stars_repo_name: lucabrugnolini/VAR

max_issues_repo_name: lucabrugnolini/VAR

max_forks_repo_name: lucabrugnolini/VAR


text:
module Solve
export solve_problem

using JuMP
using AmplNLWriter
using Lazy: @>>
using Iterators
import Base: show

function getin(nested,keys,default)
  current = nested
  for key in keys
    if haskey(current,key)
      current = current[key]
    else
      return default
    end
  end
  return current
end

function setin(nested,keys,to)
  current = nested
  keys = collect(keys)
  for key in keys[1:(end-1)]
    if !haskey(current,key)
      current[key] = Dict()
    end
    current = current[key]
  end
  current[keys[end]] = to

  nested
end

const qorders = [[1 1 1];[1 2 3];[1 3 2];[2 1 3];[2 3 1];[3 1 2];[3 2 1]]

type Required
  i::Int
  j::Int
  hours::Int
end

type Problem
  # the original json representation describing the problem
  json::Dict

  # problem dimensions
  nstudents::Int
  ncourses::Int

  # the names of students and courses
  # used for updating the json representation
  students::Array{AbstractString}
  courses::Array{AbstractString}

  # number of hours a single assignment unit represents
  hunits::Float64

  # the ranking of each assignment choice (higher ranking, better choice)
  ranks::Array{Float64,2}

  # the set of required assignments
  required::Array{Required}

  # the total number of hours each student needs
  loads::Array{Int}

  # the total nubmer of hours each course needs
  course_total::Array{Int}

  # the number of students each course can have assigned
  course_range::Array{Int,2}

  # the set of courses for each quarter
  quarterc::Array{Array{Int}}

  # the prefered ordering of hour totals for each quarter
  qorder::Array{Int}

  # minimum and maximum difference between the total nubmer of hours and the
  # actual nubmer of hours each student has
  min_off::Int
  max_off::Array{Int}

  # the maximum number of units for any assignment
  maxunits::Int

  # the maximum number of units a student can have for any quarter
  quarter_maxunits::Int

  # the maximum number of courses a student can have for any quarter
  quarter_maxcourses::Int

  # the overall weight given to course rank in the objective function
  rankweight::Float64

  # the overall weight given to the number of courses assigned to a student
  countweight::Float64

  # the overall weight given to the hour totals for each quarter in
  # the objective function
  hourweight::Float64

  # the overall weight given to how similar the solution is to
  # some other similar solution in the objective function
  closetoweight::Float64
end

function Base.show(io::IO,p::Problem)
  println(io,"TA Assignment problem with")
  println(io,"$(p.nstudents) students and")
  println(io,"$(p.ncourses) courses.")
end

function setup_problem(problem,prefs)
  courses = collect(keys(problem["courses"]))
  course_indices = Dict([name => i for (i,name) in enumerate(courses)])
  students = collect(keys(problem["students"]))
  student_indices = Dict([name => i for (i,name) in enumerate(students)])

  function quarter_courses(quarter)
    @>> problem["courses"] begin
      filter((k,c) -> c["quarter"] == quarter)
      map(c -> course_indices[c.first])
    end
  end

  function assignment_rank(istudent,icourse)
    student,course = students[istudent],courses[icourse]
    extra_hours = getin(problem,["students",student,"allow_more_hours"],false)
    student_weight = (extra_hours ?
                      prefs["student_extrahour_weight"] :
                      prefs["student_weight"])

    (getin(problem["assignments"],[student,course,"studentRank"],
           prefs["default_student_rank"]) * student_weight) +
             (getin(problem["assignments"],[student,course,"instructorRank"],
                    prefs["default_instructor_rank"]) * prefs["instructor_weight"])
  end

  nstudents = length(problem["students"])
  ncourses = length(problem["courses"])

  hunits = prefs["hour_unit"]
  maxunits = round(Int,prefs["max_units"])
  quarter_maxcourses = round(Int,prefs["quarter_max_courses"])
  quarter_maxunits = round(Int,prefs["quarter_max_units"])

  ranks = [assignment_rank(i,j) for i = 1:nstudents, j = 1:ncourses]

  reqs = Array{Required,1}()
  for i in 1:nstudents
    for j in 1:ncourses
      student,course = students[i],courses[j]
      hours = getin(problem["assignments"],[student,course,"hours"],0)
      fix = getin(problem["assignments"],[student,course,"fix"],false)
      if hours > 0 && fix
        units = round(Int,hours / hunits)
        push!(reqs,Required(i,j,units))
      end
    end
  end

  course_total = map(courses) do cid
    round(Int,float(problem["courses"][cid]["hours"]["total"])/hunits)
  end

  course_range = Array{Int,2}(ncourses,2)
  for (i,cid) in enumerate(courses)
    course_range[i,:] =
      round(Int,map(float,problem["courses"][cid]["hours"]["range"]))
  end

  quarterc = Array{Array{Int}}(3)
  quarterc[1] = quarter_courses("fall")
  quarterc[2] = quarter_courses("winter")
  quarterc[3] = quarter_courses("spring")

  qorder = map(students) do name
    orderi =
      round(Int,float(getin(problem,["students",name,"quarter_load"],0)))+1
    qorders[orderi]
  end

  loads = map(students) do name
    hours = float(getin(problem,["students",name,"total_hours"],0))
    round(Int,hours/hunits)
  end

  max_off = map(students) do name
    if getin(problem,["students",name,"allow_more_hours"],false)
      prefs["max_extra_over_units"]
    else
      prefs["max_over_units"]
    end
  end

  min_off = -prefs["max_under_units"]

  mean_student_weight = (sum(students) do x
    if getin(problem,["students",x,"allow_more_hours"],false)
      prefs["student_extrahour_weight"]
    else
      prefs["student_weight"]
    end
  end)/length(students)

  maxrank = (2*mean_student_weight + 2*prefs["instructor_weight"])
  rankweight = 1 / ncourses / maxrank * prefs["rank_weight"]

  countweight = 1 / ncourses * prefs["count_weight"]

  # 1 / numstudents / max dot product
  hourweight = 1 / nstudents / sum([3,2,1].^2) * prefs["hour_weight"]

  # 1 / numcourses
  closetoweight = 1 / ncourses

  Problem(problem,nstudents,ncourses,students,courses,hunits,ranks,reqs,loads,
          course_total,course_range,quarterc,qorder,min_off,max_off,maxunits,
          quarter_maxunits,quarter_maxcourses,rankweight,countweight,hourweight,
          closetoweight)
end

type Solution
  result::Symbol
  assignments::Array{Int,2}
end

Solution(r=:None) = Solution(r,Array{Int,2}())
function Solution(p::Problem,probjson::Dict)
  assignments = zeros(Int,(p.nstudents,p.ncourses))
  for i in 1:p.nstudents
    for j in 1:p.ncourses
      student,course = p.students[i],p.courses[j]
      hours = getin(probjson["assignments"],[student,course,"hours"],0)
      if hours > 0
        units = round(Int,hours / p.hunits)
        assignments[i,j] = units
      end
    end
  end
  solved = all(1:p.ncourses) do j
    p.course_total[j] == sum(assignments[:,j])
  end

  Solution((solved ? :Optimal : :Partial),assignments)
end

function setup_constraints(m,assignment,p::Problem,old::Solution)
  # maintain the correct representation of assignment hours
  for k = 2:p.maxunits
    @constraint(m,assignment[i=1:p.nstudents,j=1:p.ncourses,k] .≤
                assignment[i=1:p.nstudents,j=1:p.ncourses,k-1])
  end

  # the number of TAs for a course must fall within the required range
  for j in 1:p.ncourses
    @constraint(m,p.course_range[j,1] ≤
                sum{assignment[i,j,1],i=1:p.nstudents} ≤
                p.course_range[j,2])
  end

  # all classes must have the right number of total hours
  @constraint(m,p.course_total[j=1:p.ncourses]' .==
              sum{assignment[i,j,k], i=1:p.nstudents,k=1:p.maxunits})

  # students cannot have more than the max hours or courses in a quarter
  if !isempty(p.quarterc[1])
    @constraint(m,sum{assignment[i=1:p.nstudents,j,k], j=p.quarterc[1],
                      k=1:p.maxunits} .≤ p.quarter_maxunits)
    @constraint(m,sum{assignment[i=1:p.nstudents,j,1], j=p.quarterc[1]} .≤
                p.quarter_maxcourses)
  end

  if !isempty(p.quarterc[2])
    @constraint(m,sum{assignment[i=1:p.nstudents,j,k], j=p.quarterc[2],
                      k=1:p.maxunits} .≤ p.quarter_maxunits)
    @constraint(m,sum{assignment[i=1:p.nstudents,j,1], j=p.quarterc[2]} .≤
                p.quarter_maxcourses)
  end

  if !isempty(p.quarterc[3])
    @constraint(m,sum{assignment[i=1:p.nstudents,j,k], j=p.quarterc[3],
                      k=1:p.maxunits} .≤ p.quarter_maxunits)
    @constraint(m,sum{assignment[i=1:p.nstudents,j,1], j=p.quarterc[3]} .≤
                p.quarter_maxcourses)
  end

  # students' overall load cannot be too far off from their target load
  @constraint(m,p.min_off .≤
              sum{assignment[i,j,k], j=1:p.ncourses,k=1:p.maxunits} -
              p.loads[i=1:p.nstudents].≤
              p.max_off[i])

  # all required assignments must be satisfied
  for req = p.required
    @constraint(m,sum{assignment[req.i,req.j,k],k=1:p.maxunits} ==
                req.hours)
  end

  # new solutions cannot have the same exact sum of assignment ranks in the
  # same assignment locations as an old, valid solution: this guarantees the
  # new solution is different.
  #
  # NOTE: Some different solutions will be missed by this approach because they
  # will not satisfy this constraint. However, they will be trivially different:
  # differing only in the distribution of hours alloted between courses that are
  # equally ranked. This approach has the advantage that it can be expressed as
  # a linear constraint
  if old.result == :Optimal
    position = old.assignments .> 0
    oldsum = sum(old.assignments .* position .* p.ranks)
    @constraint(m,sum{assignment[i,j,k]*p.ranks[i,j]*position[i,j],
                      i=1:p.nstudents,j=1:p.ncourses,
                      k=1:p.maxunits} ≤ oldsum-1)
  end

  m
end

function setup_objective(m,assignment,p::Problem)
  const divisor = 1 / p.ncourses

  @objective(m,:Max,
             # we want the highest rank possible
             p.rankweight *
             sum{assignment[i,j,k]*p.ranks[i,j],
                 i=1:p.nstudents,j=1:p.ncourses,k=1:p.maxunits} +

             # all else being equal, we want to minimize the number of courses a
             # student is assigned
             p.countweight *
             (-sum{assignment[i,j,1], i=1:p.nstudents,j=1:p.ncourses}) +

             # we prefer qorder[1] > qorder[2] > qorder[3]
             3*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i],1]],
                                k=1:p.maxunits} +

             2*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i],2]],
                                k=1:p.maxunits} +

             1*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i],3]],
                                k=1:p.maxunits})
  m
end

function setup_objective(m,assignment,p::Problem,closeto)
  @objective(m,:Max,
             # we want the highest rank possible
             p.rankweight *
             sum{assignment[i,j,k]*p.ranks[i,j],
                 i=1:p.nstudents,j=1:p.ncourses,k=1:p.maxunits} +

             # we want a solution close to the old one
             p.closetoweight *
             sum{assignment[i,j,k] * closeto[i,j],
                 i=1:p.nstudents,j=1:p.ncourses,k=1:p.maxunits} +

             # all else being equal, we want to minimize the number of courses a
             # student is assigned
             p.countweight *
             (-sum{assignment[i,j,1], i=1:p.nstudents,j=1:p.ncourses}) +

             # we prefer qorder[1] > qorder[2] > qorder[3]
             3*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i]][1]],
                                k=1:p.maxunits} +

             2*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i]][2]],
                                k=1:p.maxunits} +

             1*p.hourweight*sum{assignment[i,j,k],
                                i=1:p.nstudents,
								j=p.quarterc[qorders[p.qorder[i]][3]],
                                k=1:p.maxunits})
  m
end

function assign_hours(p::Problem,old=Solution();closeto=false)
  m = Model()

  @variable(m,0 ≤ assignment[1:p.nstudents,1:p.ncourses,1:p.maxunits] ≤ 1,Int)

  m = setup_constraints(m,assignment,p,old)

  if closeto
    m = setup_objective(m,assignment,p,old.assignments)
  else
    m = setup_objective(m,assignment,p)
  end

  result = solve(m)

  if result == :Optimal
    a = getvalue(assignment)
    Solution(result,sum(a,3)[:,:,1])
  else
    Solution(result)
  end
end

function represent_solution(p::Problem,solution::Solution)
  represent_solution(p,solution.assignments)
end

function represent_solution(p::Problem,solution::Array{Int,2})
  result = deepcopy(p.json)

  for i = 1:p.nstudents
    for j = 1:p.ncourses
      if solution[i,j] > 0
        result = setin(result,
                       ["assignments",p.students[i],p.courses[j],"hours"],
                       p.hunits * solution[i,j])
      else
        result = setin(result,
                       ["assignments",p.students[i],p.courses[j],"hours"],0)
      end
    end
  end

  result
end

function validate(p)
  course_hours = sum(p.course_total)*p.hunits
  min_hours = (sum(p.loads) + p.min_off*p.nstudents)*p.hunits
  max_hours = (sum(p.loads) + sum(p.max_off))*p.hunits
  if course_hours < min_hours

    return """The TAs have too many assigned hours: at least $min_hours must be
available, but the courses require only $course_hours. Reduce the number of
hours that each TA must be assigned, or increase the number of hours a TA can
fall below their required amount."""

  elseif course_hours > max_hours

    return """The TAs have too few assigned hours: the courses require
$course_hours hours, but TAs can be assigned at most $max_hours. Increase the
number of hours that each TA must be assigned."""

  end

  return ""
end

function solve_problem(problem,prefs)
  p = setup_problem(problem,prefs)

  message = validate(p)
  if !isempty(message)
    return Dict("result" => "error",
                "message" => message)
  end

  if prefs["closeto"] || prefs["differentfrom"]
    s = assign_hours(p,Solution(p,problem),closeto=prefs["closeto"])
  else
    s = assign_hours(p)
  end

  if s.result == :Optimal
    Dict("result" => "optimal",
         "solution" => represent_solution(p,s))
  elseif s.result == :Infeasible
    Dict("result" => "infeasible")
  else
    Dict("result" => "error",
         "message" => "Solver returned: "*string(s.result))
  end
end

end

ext: jl

max_stars_repo_name: haberdashPI/csd-ta-assign

max_issues_repo_name: haberdashPI/csd-ta-assign

max_forks_repo_name: haberdashPI/csd-ta-assign


text:
export bissect
function bissect(h       :: AbstractNLPModel,
                      nlpstop :: AbstractStopping;
                      t1 = h.meta.x0, t2 = 100.0,
                      verbose = true)

    (tₐ, tᵦ) = trouve_intervalle(h, t1, 1.0)
    gₐ = grad(h, tₐ)
    gᵦ = grad(h, tᵦ)
    iter = 0

    OK = update_and_start!(nlpstop, x = tₐ, gx = gₐ, g0 = copy(gₐ))

    verbose &&
        @printf(" iter        tₐ        tᵦ         gₐ        gᵦ        \n")
    verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ[1], tᵦ[1], gₐ[1], gᵦ[1])
    while !OK
      tₚ = (tₐ .+ tᵦ) ./ 2
      gₚ = grad(h, tₚ)
      if true in (gₚ .<= 0.0)
        tₐ = tₚ
        gₐ = gₚ
        OK = update_and_stop!(nlpstop, x = tₐ, gx = gₐ)
      else
        tᵦ = tₚ
        gᵦ = gₚ
        OK = update_and_stop!(nlpstop, x = tᵦ, gx = gᵦ)
      end

      iter += 1
      verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ[1], tᵦ[1], gₐ[1], gᵦ[1])
    end

    topt = (tₐ .+ tᵦ) ./ 2.0
    update!(nlpstop.current_state, x = topt, gx = grad(h, topt))
    optimal = nlpstop.meta.optimal
    return optimal, nlpstop
end


text:
using Base.Test
using Reactive

# Processes 1 messages from the Reactive event queue.
onestep() = Reactive.run(1)

x = Signal(2)
@test value(x) == 2

mul3 = map(a ->  3a, x)
@test value(mul3) == 6

push!(x, 3)
onestep()

@test value(x) == 3
@test value(mul3) == 9


text: # TODO: Add tests here.


text:
@testset "Node" begin
    lt_bounds = [
        Cerberus.BoundUpdate(_CVI(1), _LT(0.0)),
        Cerberus.BoundUpdate(_CVI(3), _LT(1.0)),
        Cerberus.BoundUpdate(_CVI(5), _LT(0.0)),
    ]
    gt_bounds = [
        Cerberus.BoundUpdate(_CVI(2), _GT(1.0)),
        Cerberus.BoundUpdate(_CVI(3), _GT(0.0)),
        Cerberus.BoundUpdate(_CVI(6), _GT(1.0)),
    ]
    lt_constrs =
        [Cerberus.AffineConstraint{_LT}(_CSAF([2.3], [_CVI(5)], 1.9), _LT(4.5))]
    gt_constrs = [
        Cerberus.AffineConstraint{_GT}(
            _CSAF([2.3, 4.5], [_CVI(5), _CVI(3)], 6.7),
            _GT(6.7),
        ),
    ]
    depth = 6
    dual_bound = 3.2

    n1 = @inferred Cerberus.Node()
    n2 = @inferred Cerberus.Node(
        copy(lt_bounds),
        copy(gt_bounds),
        lt_constrs,
        gt_constrs,
        depth,
    )
    n3 = @inferred Cerberus.Node(
        copy(lt_bounds),
        copy(gt_bounds),
        lt_constrs,
        gt_constrs,
        depth,
        dual_bound,
    )

    @test n1.depth == 0
    @test n2.depth == depth
    @test n3.depth == depth

    @test isempty(n1.lt_bounds)
    @test n2.lt_bounds == lt_bounds
    @test n3.lt_bounds == lt_bounds

    @test isempty(n1.gt_bounds)
    @test n2.gt_bounds == gt_bounds
    @test n3.gt_bounds == gt_bounds

    @test isempty(n1.lt_general_constrs)
    @test n2.lt_general_constrs == lt_constrs
    @test n3.lt_general_constrs == lt_constrs

    @test isempty(n1.gt_general_constrs)
    @test n2.gt_general_constrs == gt_constrs
    @test n3.gt_general_constrs == gt_constrs

    @test n1.dual_bound == -Inf
    @test n2.dual_bound == -Inf
    @test n3.dual_bound == dual_bound

    @testset "copy" begin
        _n1 = copy(n1)
        _n2 = copy(n2)
        _n3 = copy(n3)

        @test _n1.depth == 0
        @test _n2.depth == depth
        @test _n3.depth == depth

        @test isempty(_n1.lt_bounds)
        @test _n2.lt_bounds == lt_bounds
        @test _n3.lt_bounds == lt_bounds

        @test isempty(_n1.gt_bounds)
        @test _n2.gt_bounds == gt_bounds
        @test _n3.gt_bounds == gt_bounds

        @test isempty(_n1.lt_general_constrs)
        @test _n2.lt_general_constrs == lt_constrs
        @test _n3.lt_general_constrs == lt_constrs

        @test isempty(_n1.gt_general_constrs)
        @test _n2.gt_general_constrs == gt_constrs
        @test _n3.gt_general_constrs == gt_constrs

        @test _n1.dual_bound == -Inf
        @test _n2.dual_bound == -Inf
        @test _n3.dual_bound == dual_bound
    end
    @testset "apply_branching!" begin
        let bd = Cerberus.BoundUpdate(_CVI(4), _LT(7.0))
            Cerberus.apply_branching!(n1, bd)
            Cerberus.apply_branching!(n2, bd)
            Cerberus.apply_branching!(n3, bd)

            @test n1.depth == 1
            @test n2.depth == depth + 1
            @test n3.depth == depth + 1

            new_lt_bounds = vcat(lt_bounds, bd)

            @test n1.lt_bounds == [bd]
            @test n2.lt_bounds == new_lt_bounds
            @test n3.lt_bounds == new_lt_bounds

            @test isempty(n1.gt_bounds)
            @test n2.gt_bounds == gt_bounds
            @test n3.gt_bounds == gt_bounds
        end

        # Now apply a dominated branch
        let bd = Cerberus.BoundUpdate(_CVI(1), _GT(0.0))
            Cerberus.apply_branching!(n1, bd)
            Cerberus.apply_branching!(n2, bd)
            Cerberus.apply_branching!(n3, bd)

            @test n1.depth == 2
            @test n2.depth == depth + 2
            @test n3.depth == depth + 2

            new_lt_bounds =
                vcat(lt_bounds, Cerberus.BoundUpdate(_CVI(4), _LT(7.0)))
            new_gt_bounds = vcat(gt_bounds, bd)

            @test n1.lt_bounds == [Cerberus.BoundUpdate(_CVI(4), _LT(7.0))]
            @test n2.lt_bounds == new_lt_bounds
            @test n3.lt_bounds == new_lt_bounds

            @test n1.gt_bounds == [Cerberus.BoundUpdate(_CVI(1), _GT(0.0))]
            @test n2.gt_bounds == new_gt_bounds
            @test n3.gt_bounds == new_gt_bounds
        end
    end
end

@testset "Tree" begin
    # BFS, branching first on 1 and then on 2
    n1 = Cerberus.Node()
    n2 = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
        Cerberus.BoundUpdate{_GT}[],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        1,
    )
    n3 = Cerberus.Node(
        Cerberus.BoundUpdate{_LT}[],
        [Cerberus.BoundUpdate(_CVI(1), _GT(1.0))],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        1,
    )
    n4 = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(2), _LT(0.0))],
        [Cerberus.BoundUpdate(_CVI(1), _GT(1.0))],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        2,
    )
    n5 = Cerberus.Node(
        [Cerberus.BoundUpdate(_CVI(1), _LT(0.0))],
        [Cerberus.BoundUpdate(_CVI(2), _GT(1.0))],
        Cerberus.AffineConstraint{_LT}[],
        Cerberus.AffineConstraint{_GT}[],
        2,
    )

    tree = @inferred Cerberus.Tree()
    @inferred Cerberus.push_node!(tree, n5)
    @inferred Cerberus.push_node!(tree, n4)
    @inferred Cerberus.push_node!(tree, n3)
    @inferred Cerberus.push_node!(tree, n2)
    @inferred Cerberus.push_node!(tree, n1)

    @test !isempty(tree)
    @test Cerberus.num_open_nodes(tree) == 5
    @test Cerberus.pop_node!(tree) == n1

    @test !isempty(tree)
    @test Cerberus.num_open_nodes(tree) == 4
    @test Cerberus.pop_node!(tree) == n2

    @test !isempty(tree)
    @test Cerberus.num_open_nodes(tree) == 3
    @test Cerberus.pop_node!(tree) == n3

    @test !isempty(tree)
    @test Cerberus.num_open_nodes(tree) == 2
    @test Cerberus.pop_node!(tree) == n4

    @test !isempty(tree)
    @test Cerberus.num_open_nodes(tree) == 1
    @test Cerberus.pop_node!(tree) == n5

    @test isempty(tree)
end


text:
"""
Reimplementation of the LeNet5 architecture from

LeCun, Bottou, Bengio & Haffner (1998),
Gradient-based learning applied to document recognition.
Procedings of the IEEE 86(11), 2278-2324.

and a version of LeNet5 with feedback connections.
"""
module LeNet5
using Flux
using ...StateFulls
using ...Splitters
using ...Mergers
using ...FeedbackChains

export lenet5, lenet5_fb, wrapfb_lenet5

"""
    lenet5(; σ=tanh, pad=2)

Generate the LeNet5 architecture from LeCun et al. (1998) with small modifications.

## Details

The implementation differs from the original LeNet5, as the output layer does not
compute radial basis functions, but is a normal `Dense` layer with a softmax.
The input image is padded. The network assumes a 32x32 input, so for MNIST digits
a `pad` of 2 is appropriate.
The non-linearity can be customized via the `σ` argument. The standard is `tanh`,
whereas the original LeNet5 used `x -> 1.7159 .* tanh(x)`.
"""
function lenet5(; σ=tanh, pad=2)
    return Chain(
        Conv((5,5), 1=>6, σ, pad=pad),
        MaxPool((2,2), stride=(2,2)),
        Conv((5,5), 6=>16, σ),
        MaxPool((2,2), stride=(2,2)),
        Conv((5,5), 16=>120, σ),
        x -> reshape(x, 120, size(x, 4)),
        Dense(120, 84, σ),
        Dense(84, 10),
        softmax
    )
end # function lenet5

"""
    lenet5_fb(; σ=tanh, pad=2)

Generate the LeNet5 architecture from LeCun et al. (1998) with feedback connections
and small modifications.
"""
function lenet5_fb(; σ=tanh, pad=2, op=+)
    return FeedbackChain(
        Conv((5,5), 1=>6, relu, pad=2),
        Merger("conv2", ConvTranspose((10,10), 16=>6, relu, stride=2), op),
        MaxPool((2,2), stride=(2,2)),
        Conv((5,5), 6=>16, relu),
        Splitter("conv2"),
        Merger("fc1", ConvTranspose((10,10), 120=>16, relu), op),
        MaxPool((2,2), stride=(2,2)),
        Conv((5,5), 16=>120, relu),
        Splitter("fc1"),
        x -> reshape(x, 120, size(x, 4)),
        Dense(120, 84, relu),
        Dense(84, 10),
        softmax
    )
end # function lenet5_fb

"""
    wrapfb_lenet5(net, batchsize; generator=zeros)

Wrap a `letnet5_fb` network in a `Flux.Recur`, assuming that batches are of size
`batchsize` and using the given `generator` to initialize the state.
"""
function wrapfb_lenet5(net, batchsize; generator=zeros)
    h = Dict(
        "conv2" => generator(10, 10, 16, batchsize),
        "fc1" => generator(1, 1, 120, batchsize)
    )
    return StateFull(net, h)
end # function wrapfb_lenet5

end # module LeNet5


text:
module DoubleStrain

using Random
using ProgressMeter
using CovidPOMDPs
using ..CovidPOMDPs:
    INFECTION_HORIZON,
    INF_DIST,
    VIRAL_LOADS,
    POS_TEST_PROBS,
    SYMPTOMATIC_ISOLATION_PROB,
    ASYMPTOMATIC_PROB,
    SYMPTOM_PROBS,
    DEFAULT_PARAMS,
    InfParams,
    CovidState,
    CovidAction,
    simplex_sample,
    RVsum

include("typedef.jl")
export DoubleCovidPOMDP, DoubleCovidState

include("simulate.jl")

include("POMDPsInterface.jl")

include("updater.jl")

end


text:
module QJuliaPLCG2

using QJuliaInterface
using QJuliaEnums
using QJuliaBlas
using QJuliaReduce
using QJuliaSolvers

using LinearAlgebra
using MPI

norm2   = QJuliaReduce.gnorm2
rdot    = QJuliaReduce.reDotProduct
verbose = true

const ε = eps(Float64)

function solver(x::AbstractArray, b::AbstractArray, Mat::Any, MatSloppy::Any, param::QJuliaSolvers.QJuliaSolverParam_qj)

    if verbose == true; println("Running LPCG in maximum " , param.nKrylov, " iters."); end

    if MPI.Initialized() == false; error("MPI was not inititalized, copy source field to the solution."); end

    if param.nKrylov > param.maxiter; println("Warning: search space exceeds requested maximum iter count. Some routines may fail."); end
    if param.pipeline < 1; println("Pipeline length was not set."); end

    if (param.maxiter == 0) || (param.nKrylov == 0)
      if param.use_init_guess == false
        x .=@. 0.0
      end
      return
    end #if

    mixed = (param.dtype_sloppy != param.dtype)

    σ = param.shift
    l = param.pipeline

    nkrylov = param.nKrylov
    zsize   = nkrylov+1+l
    data_type = typeof(x[1]) == BigFloat ? BigFloat : Float64

    local r    = zero(typeof(x)(undef, length(x)))
    local p    = zero(typeof(x)(undef, length(x)))
    # Basis vectors (incl. auxiliary vectors)
    local Z  = zero(Matrix{typeof(x[1])}(undef, length(x), (l+1)*zsize))

    Z1   = view(Z, :, 1:zsize)
    Z2   = view(Z, :, (zsize+1):2zsize)
    Zlp1 = view(Z, :, (l*zsize+1):(l+1)*zsize)

    # G matrix
    local g = zero(Matrix{data_type}(undef, (nkrylov+1+4l), nkrylov+4l))

    local γ = zero(Vector{data_type}(undef, (nkrylov+1)))
    local δ = zero(Vector{data_type}(undef, (nkrylov+1)))

    bnorm = sqrt(norm2(b))  #Save norm of b
    rnorm = 0.0     #if zero source then we will exit immediately doing no work

    z11 = view(Z1, :, 1)

    if param.use_init_guess == true
      #r = b - Ax0
      Mat(z11, x)
      z11 .=@. z11 - σ*x
      z11 .=@. b - z11
      rnorm = sqrt(norm2(z11))
    else
      rnorm = bnorm
      z11 .=@. b
    end

    if( rnorm == 0.0 ); return; end
    if verbose == true; println("LPCG: Initial residual = ", rnorm); end

    # Scale the first basis vector:
    z11 .=@. z11 / rnorm
    # Inititalize aux basis
    for i in 2:(l+1)
      Zi  = view(Z, :, ((i-1)*zsize+1):i*zsize)
      zi1 = view(Zi,:, 1)
      zi1 .=@. z11
    end

    #
    g[1,1] = 1.0
    # CG params:
    η = 0.0; ζ = 0.0; λ = 0.0

    #start iterations:
    for i in 1:(min(param.maxiter, nkrylov)+l)
      #
      zlp1ip1 = view(Zlp1,:,i+1)
      zlp1i   = view(Zlp1,:,i)
      #
      MatSloppy(zlp1ip1, zlp1i)
      #
      if i <= l; zlp1ip1 .=@. zlp1ip1 - σ * zlp1i; end # apply shift σ -> σ[i]

      if i <= (l-1)
        for k in (i+1):l
          Zk = view(Z, :, ((k-1)*zsize+1):k*zsize)
          zkip1 = view(Zk,:,i+1)
          zkip1 .=@. zlp1ip1
        end
      end

      #
      a  = i - l
      #
      if a >= 1
        for j in max(a-l+2, 1):a
          for k in max(a+1-2l, 1):j-1
            g[j,a+1] -= (g[k,j]*g[k,a+1])
          end
          g[j,a+1] /= g[j,j]
        end
        for k in max(a+1-2l, 1):a
          g[a+1,a+1] -=g[k,a+1]*g[k,a+1]
        end

        if g[a+1,a+1] <= 0.0
          println("SQRT breakdown detected at iteration ", i)
          break
        end
        #
        g[a+1,a+1] = sqrt(g[a+1,a+1])
        if a <= l
          third_term = a == 1 ? 0.0 : δ[a-1]*g[a-1,a]
          #
          γ[a] = (g[a,a+1] + σ*g[a,a] - third_term) / g[a,a] #Note σ -> σ[a]
          δ[a] = g[a+1,a+1] / g[a,a]
        else
          #first_term = a-l == 1 ? 0.0 : g[a,a-1]*δ[a-l-1]
          #
          #γ[a] = ( first_term + g[a,a]*γ[a-l] + g[a,a+1]*δ[a-l] - δ[a-1]*g[a-1,a] ) / g[a,a]
          γ[a] = ( g[a,a]*γ[a-l] + g[a,a+1]*δ[a-l] - δ[a-1]*g[a-1,a] ) / g[a,a]
          δ[a] = (g[a+1,a+1]*δ[a-l]) / g[a,a]
        end

        println("Result :: (a =  ",  a, " ) ", γ[a], " , ", δ[a], " ")

        third_coeff = (a == 1) ? 0.0 : δ[a-1]

        if a == 1
          Z1[:,a+1] .=@. Z2[:,a+1] + (σ - γ[a])*Z1[:,a]
        else
          Z1[:,a+1] .=@. Z2[:,a+1] + (σ - γ[a])*Z1[:,a] - δ[a-1]*Z1[:,a-1]
        end
        Z1[:,a+1] .=@. Z1[:,a+1] / δ[a]

        for k in 2:l
          Zk   = view(Z, : , ((k-1)*zsize+1):k*zsize)
          Zkp1 = view(Z, : , (k*zsize+1):(k+1)*zsize)

          Zk[:,a+k] .=@. Zkp1[:,a+k] + (σ - γ[a])*Zk[:,a+k-1] - third_coeff*Zk[:,a+k-2]
          Zk[:,a+k] .=@. Zk[:,a+k] / δ[a]
        end

        zlp1ip1 .=@. zlp1ip1 - γ[a]*Zlp1[:,i] - third_coeff*Zlp1[:,i-1]
        zlp1ip1 .=@. zlp1ip1 / δ[a]
      end # a >= 1

      for j in max(a-l+1,1):a+1; g[j,i+1] = dot(Zlp1[:, i+1], Z1[:, j]); end
      for j in max(a+2,1):i+1;   g[j,i+1] = dot(Zlp1[:, i+1], Zlp1[:, j]); end

      if a == 1
        η = γ[1]
        ζ = rnorm
        p .=@. z11 / η
      elseif a >= 2
        λ = δ[a-1] / η
        η = γ[a] - λ*δ[a-1]
        x .=@. x + ζ*p
        ζ = - λ*ζ
        p .=@. (Z1[:, a] - δ[a-1]*p) / η

      end

    end #for iters

    # Compute true residual
    Mat(r, x)

    r .=@. b - r
    rnorm = sqrt(norm2(r))

    param.true_res = rnorm / bnorm
    if verbose == true; println("GMRes: Relative residual: true = ", param.true_res); end

    #println("Hessenberg matrix\n\n", Hm1m)

end # solver

end #QJuliaPLCG2

max_stars_repo_name: alexstrel/QJulia

max_issues_repo_name: alexstrel/QJulia

max_forks_repo_name: alexstrel/QJulia


text:
module Profile

import Base: hash, ==

export @profile

macro profile(ex)
    quote
        try
            status = start_timer()
            if status < 0
                error(error_codes[status])
            end
            $(esc(ex))
        finally
            stop_timer()
        end
    end
end

####
#### User-level functions
####
function init(n::Integer, delay::Float64)
    status = ccall(:jl_profile_init, Cint, (Csize_t, Uint64), n, iround(10^9*delay))
    if status == -1
        error("could not allocate space for ", n, " instruction pointers")
    end
end

# init with default values
# Use a max size of 1M profile samples, and fire timer every 1ms
__init__() = init(1_000_000, 0.001)

clear() = ccall(:jl_profile_clear_data, Void, ())

function print{T<:Unsigned}(io::IO, data::Vector{T} = fetch(), lidict::Dict = getdict(data); format = :tree, C = false, combine = true, cols = Base.tty_size()[2])
    if format == :tree
        tree(io, data, lidict, C, combine, cols)
    elseif format == :flat
        flat(io, data, lidict, C, combine, cols)
    else
        error("output format ", format, " not recognized")
    end
end
print{T<:Unsigned}(data::Vector{T} = fetch(), lidict::Dict = getdict(data); kwargs...) = print(STDOUT, data, lidict; kwargs...)

function retrieve()
    data = fetch()
    copy(data), getdict(data)
end

function getdict(data::Vector{Uint})
    uip = unique(data)
    Dict(uip, [lookup(ip) for ip in uip])
end

####
#### Internal interface
####
immutable LineInfo
    func::ByteString
    file::ByteString
    line::Int
    fromC::Bool
end

const UNKNOWN = LineInfo("?", "?", -1, true)

==(a::LineInfo, b::LineInfo) = a.line == b.line && a.fromC == b.fromC && a.func == b.func && a.file == b.file

function hash(li::LineInfo, h::Uint)
    h += uint(0xf4fbda67fe20ce88)
    h = hash(li.line, h)
    h = hash(li.file, h)
    h = hash(li.func, h)
end

# C wrappers
start_timer() = ccall(:jl_profile_start_timer, Cint, ())

stop_timer() = ccall(:jl_profile_stop_timer, Void, ())

is_running() = bool(ccall(:jl_profile_is_running, Cint, ()))

get_data_pointer() = convert(Ptr{Uint}, ccall(:jl_profile_get_data, Ptr{Uint8}, ()))

len_data() = convert(Int, ccall(:jl_profile_len_data, Csize_t, ()))

maxlen_data() = convert(Int, ccall(:jl_profile_maxlen_data, Csize_t, ()))

function lookup(ip::Uint)
    info = ccall(:jl_lookup_code_address, Any, (Ptr{Void},Cint), ip, false)
    if length(info) == 4
        return LineInfo(string(info[1]), string(info[2]), int(info[3]), info[4])
    else
        return UNKNOWN
    end
end

error_codes = (Int=>ASCIIString)[
    -1=>"cannot specify signal action for profiling",
    -2=>"cannot create the timer for profiling",
    -3=>"cannot start the timer for profiling",
    -4=>"cannot unblock SIGUSR1"]

function fetch()
    len = len_data()
    maxlen = maxlen_data()
    if (len == maxlen)
        warn("The profile data buffer is full; profiling probably terminated\nbefore your program finished. To profile for longer runs, call Profile.init()\nwith a larger buffer and/or larger delay.")
    end
    pointer_to_array(get_data_pointer(), (len,))
end

# Number of backtrace "steps" that are triggered by taking the backtrace, e.g., inside profile_bt
# May be platform-specific?
#@unix_only const btskip = 2
#@windows_only const btskip = 0
const btskip = 0

## Print as a flat list
# Counts the number of times each line appears, at any nesting level
function count_flat{T<:Unsigned}(data::Vector{T})
    linecount = (T=>Int)[]
    toskip = btskip
    for ip in data
        if toskip > 0
            toskip -= 1
            continue
        end
        if ip == 0
            toskip = btskip
            continue
        end
        linecount[ip] = get(linecount, ip, 0)+1
    end
    iplist = Array(T, 0)
    n = Array(Int, 0)
    for (k,v) in linecount
        push!(iplist, k)
        push!(n, v)
    end
    return iplist, n
end

function parse_flat(iplist, n, lidict, C::Bool)
    # Convert instruction pointers to names & line numbers
    lilist = [lidict[ip] for ip in iplist]
    # Keep only the interpretable ones
    # The ones with no line number might appear multiple times in a single
    # backtrace, giving the wrong impression about the total number of backtraces.
    # Delete them too.
    keep = !Bool[x == UNKNOWN || x.line == 0 || (x.fromC && !C) for x in lilist]
    n = n[keep]
    lilist = lilist[keep]
    lilist, n
end

function flat{T<:Unsigned}(io::IO, data::Vector{T}, lidict::Dict, C::Bool, combine::Bool, cols::Integer)
    iplist, n = count_flat(data)
    if isempty(n)
        warning_empty()
        return
    end
    lilist, n = parse_flat(iplist, n, lidict, C)
    print_flat(io, lilist, n, combine, cols)
end

function print_flat(io::IO, lilist::Vector{LineInfo}, n::Vector{Int}, combine::Bool, cols::Integer)
    p = liperm(lilist)
    lilist = lilist[p]
    n = n[p]
    if combine
        j = 1
        for i = 2:length(lilist)
            if lilist[i] == lilist[j]
                n[j] += n[i]
                n[i] = 0
            else
                j = i
            end
        end
        keep = n .> 0
        n = n[keep]
        lilist = lilist[keep]
    end
    wcounts = max(6, ndigits(maximum(n)))
    maxline = 0
    maxfile = 0
    maxfunc = 0
    for li in lilist
        maxline = max(maxline, li.line)
        maxfile = max(maxfile, length(li.file))
        maxfunc = max(maxfunc, length(li.func))
    end
    wline = max(5, ndigits(maxline))
    ntext = cols - wcounts - wline - 3
    if maxfile+maxfunc <= ntext
        wfile = maxfile
        wfunc = maxfunc
    else
        wfile = ifloor(2*ntext/5)
        wfunc = ifloor(3*ntext/5)
    end
    println(io, lpad("Count", wcounts, " "), " ", rpad("File", wfile, " "), " ", rpad("Function", wfunc, " "), " ", lpad("Line", wline, " "))
    for i = 1:length(n)
        li = lilist[i]
        println(io, lpad(string(n[i]), wcounts, " "), " ", rpad(truncto(li.file, wfile), wfile, " "), " ", rpad(truncto(li.func, wfunc), wfunc, " "), " ", lpad(string(li.line), wline, " "))
    end
end

## A tree representation
# Identify and counts repetitions of all unique backtraces
function tree_aggregate{T<:Unsigned}(data::Vector{T})
    iz = find(data .== 0)  # find the breaks between backtraces
    treecount = (Vector{T}=>Int)[]
    istart = 1+btskip
    for iend in iz
        tmp = data[iend-1:-1:istart]
        treecount[tmp] = get(treecount, tmp, 0)+1
        istart = iend+1+btskip
    end
    bt = Array(Vector{T}, 0)
    counts = Array(Int, 0)
    for (k,v) in treecount
        if !isempty(k)
            push!(bt, k)
            push!(counts, v)
        end
    end
    bt, counts
end

tree_format_linewidth(x::LineInfo) = ndigits(x.line)+6

function tree_format(lilist::Vector{LineInfo}, counts::Vector{Int}, level::Int, cols::Integer)
    nindent = min(ifloor(cols/2), level)
    ndigcounts = ndigits(maximum(counts))
    ndigline = maximum([tree_format_linewidth(x) for x in lilist])
    ntext = cols-nindent-ndigcounts-ndigline-5
    widthfile = ifloor(0.4ntext)
    widthfunc = ifloor(0.6ntext)
    strs = Array(ByteString, length(lilist))
    showextra = false
    if level > nindent
        nextra = level-nindent
        nindent -= ndigits(nextra)+2
        showextra = true
    end
    for i = 1:length(lilist)
        li = lilist[i]
        if li != UNKNOWN
            base = " "^nindent
            if showextra
                base = string(base, "+", nextra, " ")
            end
            base = string(base,
                          rpad(string(counts[i]), ndigcounts, " "),
                          " ",
                          truncto(string(li.file), widthfile),
                          "; ",
                          truncto(string(li.func), widthfunc),
                          "; ")
            strs[i] = string(base, "line: ", li.line)
        else
            strs[i] = ""
        end
    end
    strs
end

# Print a "branch" starting at a particular level. This gets called recursively.
function tree{T<:Unsigned}(io::IO, bt::Vector{Vector{T}}, counts::Vector{Int}, lidict::Dict, level::Int, combine::Bool, cols::Integer)
    # Organize backtraces into groups that are identical up to this level
    if combine
        # Combine based on the line information
        d = (LineInfo=>Vector{Int})[]
        for i = 1:length(bt)
            ip = bt[i][level+1]
            key = lidict[ip]
            indx = Base.ht_keyindex(d, key)
            if indx == -1
                d[key] = [i]
            else
                push!(d.vals[indx], i)
            end
        end
        # Generate counts
        dlen = length(d)
        lilist = Array(LineInfo, dlen)
        group = Array(Vector{Int}, dlen)
        n = Array(Int, dlen)
        i = 1
        for (key, v) in d
            lilist[i] = key
            group[i] = v
            n[i] = sum(counts[v])
            i += 1
        end
    else
        # Combine based on the instruction pointer
        d = (T=>Vector{Int})[]
        for i = 1:length(bt)
            key = bt[i][level+1]
            indx = Base.ht_keyindex(d, key)
            if indx == -1
                d[key] = [i]
            else
                push!(d.vals[indx], i)
            end
        end
        # Generate counts, and do the code lookup
        dlen = length(d)
        lilist = Array(LineInfo, dlen)
        group = Array(Vector{Int}, dlen)
        n = Array(Int, dlen)
        i = 1
        for (key, v) in d
            lilist[i] = lidict[key]
            group[i] = v
            n[i] = sum(counts[v])
            i += 1
        end
    end
    # Order the line information
    if length(lilist) > 1
        p = liperm(lilist)
        lilist = lilist[p]
        group = group[p]
        n = n[p]
    end
    # Generate the string for each line
    strs = tree_format(lilist, n, level, cols)
    # Recurse to the next level
    len = Int[length(x) for x in bt]
    for i = 1:length(lilist)
        if !isempty(strs[i])
            println(io, strs[i])
        end
        idx = group[i]
        keep = len[idx] .> level+1
        if any(keep)
            idx = idx[keep]
            tree(io, bt[idx], counts[idx], lidict, level+1, combine, cols)
        end
    end
end

function tree{T<:Unsigned}(io::IO, data::Vector{T}, lidict::Dict, C::Bool, combine::Bool, cols::Integer)
    if !C
        data = purgeC(data, lidict)
    end
    bt, counts = tree_aggregate(data)
    if isempty(counts)
        warning_empty()
        return
    end
    level = 0
    len = Int[length(x) for x in bt]
    keep = len .> 0
    tree(io, bt[keep], counts[keep], lidict, level, combine, cols)
end

# Utilities
function truncto(str::ByteString, w::Int)
    ret = str;
    if length(str) > w
        ret = string("...", str[end-w+4:end])
    end
    ret
end

# Order alphabetically (file, function) and then by line number
function liperm(lilist::Vector{LineInfo})
    comb = Array(ByteString, length(lilist))
    for i = 1:length(lilist)
        li = lilist[i]
        if li != UNKNOWN
            comb[i] = @sprintf("%s:%s:%06d", li.file, li.func, li.line)
        else
            comb[i] = "zzz"
        end
    end
    sortperm(comb)
end

warning_empty() = warn("""
            There were no samples collected. Run your program longer (perhaps by
            running it multiple times), or adjust the delay between samples with
            Profile.init().""")

function purgeC(data, lidict)
    keep = Bool[d == 0 || lidict[d].fromC == false for d in data]
    data[keep]
end

end # module

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia