text:
using Distributed
using ThreadPools
using JLD2
using LinearMaps
#using IterativeSolvers

"""
    build_full_H(ci_vector::ClusteredState, cluster_ops, clustered_ham::ClusteredOperator)

Build full TPSCI Hamiltonian matrix in space spanned by `ci_vector`. This works in serial for the full matrix
"""
function build_full_H(ci_vector::ClusteredState, cluster_ops, clustered_ham::ClusteredOperator)
#={{{=#
    dim = length(ci_vector)
    H = zeros(dim, dim)

    zero_fock = TransferConfig([(0,0) for i in ci_vector.clusters])
    bra_idx = 0
    for (fock_bra, configs_bra) in ci_vector.data
        for (config_bra, coeff_bra) in configs_bra
            bra_idx += 1
            ket_idx = 0
            for (fock_ket, configs_ket) in ci_vector.data
                fock_trans = fock_bra - fock_ket

                # check if transition is connected by H
                if haskey(clustered_ham, fock_trans) == false
                    ket_idx += length(configs_ket)
                    continue
                end

                for (config_ket, coeff_ket) in configs_ket
                    ket_idx += 1
                    ket_idx <= bra_idx || continue

                    for term in clustered_ham[fock_trans]

                        check_term(term, fock_bra, config_bra, fock_ket, config_ket) || continue

                        me = contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_ket, config_ket)
                        H[bra_idx, ket_idx] += me
                    end

                    H[ket_idx, bra_idx] = H[bra_idx, ket_idx]

                end
            end
        end
    end
    return H
end
#=}}}=#

"""
    build_full_H_parallel(ci_vector::ClusteredState, cluster_ops, clustered_ham::ClusteredOperator)

Build full TPSCI Hamiltonian matrix in space spanned by `ci_vector`. This works in serial for the full matrix
"""
function build_full_H_parallel( ci_vector_l::ClusteredState{T,N,R}, ci_vector_r::ClusteredState{T,N,R},
                                cluster_ops, clustered_ham::ClusteredOperator;
                                sym=false) where {T,N,R}
#={{{=#
    dim_l = length(ci_vector_l)
    dim_r = length(ci_vector_r)
    H = zeros(dim_l, dim_r)

    dim_l == dim_r || sym == false || error(" dim_l!=dim_r yet sym==true")

    if (dim_l == dim_r) && sym == false
        @warn(" are you missing sym=true?")
    end
    jobs = []

    zero_fock = TransferConfig([(0,0) for i in 1:N])
    bra_idx = 0

    for (fock_bra, configs_bra) in ci_vector_l.data
        for (config_bra, coeff_bra) in configs_bra
            bra_idx += 1
            #push!(jobs, (bra_idx, fock_bra, config_bra) )
            #push!(jobs, (bra_idx, fock_bra, config_bra, H[bra_idx,:]) )
            push!(jobs, (bra_idx, fock_bra, config_bra, zeros(dim_r)) )
        end
    end

    function do_job(job)
        fock_bra = job[2]
        config_bra = job[3]
        Hrow = job[4]
        ket_idx = 0

        for (fock_ket, configs_ket) in ci_vector_r.data
            fock_trans = fock_bra - fock_ket

            # check if transition is connected by H
            if haskey(clustered_ham, fock_trans) == false
                ket_idx += length(configs_ket)
                continue
            end

            for (config_ket, coeff_ket) in configs_ket
                ket_idx += 1
                ket_idx <= job[1] || sym == false || continue

                for term in clustered_ham[fock_trans]

                    #length(term.clusters) <= 2 || continue
                    check_term(term, fock_bra, config_bra, fock_ket, config_ket) || continue

                    me = contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_ket, config_ket)
                    #if term isa ClusteredTerm4B
                    #    @btime contract_matrix_element($term, $cluster_ops, $fock_bra, $config_bra, $fock_ket, $config_ket)
                    #end
                    Hrow[ket_idx] += me
                    #H[job[1],ket_idx] += me
                end

            end

        end
    end

    # because @threads divides evenly the loop, let's distribute thework more fairly
    #mid = length(jobs) ÷ 2
    #r = collect(1:length(jobs))
    #perm = [r[1:mid] reverse(r[mid+1:end])]'[:]
    #jobs = jobs[perm]

    #for job in jobs
    Threads.@threads for job in jobs
        do_job(job)
        #@btime $do_job($job)
    end

    for job in jobs
        H[job[1],:] .= job[4]
    end

    if sym
        for i in 1:dim_l
            @simd for j in i+1:dim_l
                @inbounds H[i,j] = H[j,i]
            end
        end
    end

    return H
end
#=}}}=#

"""
    function tps_ci_direct( ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
                        H_old    = nothing,
                        v_old    = nothing,
                        verbose   = 0) where {T,N,R}

# Solve for eigenvectors/values in the basis defined by `ci_vector`. Use direct diagonalization.

If updating existing matrix, pass in H_old/v_old to avoid rebuilding that block
"""
function tps_ci_direct( ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
                        conv_thresh = 1e-5,
                        max_ss_vecs = 12,
                        max_iter    = 40,
                        shift       = nothing,
                        precond     = false,
                        H_old    = nothing,
                        v_old    = nothing,
                        verbose   = 0) where {T,N,R}
    #={{{=#
    println()
    @printf(" |== Tensor Product State CI =======================================\n")
    vec_out = deepcopy(ci_vector)
    e0 = zeros(T,R)
    @printf(" Hamiltonian matrix dimension = %5i: \n", length(ci_vector))
    dim = length(ci_vector)
    flush(stdout)

    precond == false || @warn("davidson preconditioning NYI")

    H = zeros(T, 1,1)

    if H_old != nothing
        v_old != nothing || error(" can't specify H_old w/out v_old")
        v_tot = deepcopy(ci_vector)
        v_new = deepcopy(ci_vector)

        project_out!(v_new, v_old)

        #v_tot = deepcopy(v_old)
        #add!(v_tot, v_new)

        dim_old = length(v_old)
        dim_new = length(v_new)

        # create indexing to find old indices in new space
        indices = OrderedDict{FockConfig{N}, OrderedDict{ClusterConfig{N}, Int}}()

        idx = 1
        for (fock,configs) in v_tot.data
            indices[fock] = OrderedDict{ClusterConfig{N}, Int}()
            for (config,coeff) in configs
                indices[fock][config] = idx
                idx += 1
            end
        end

        dim = dim_old + dim_new

        dim == length(v_tot) || error(" not adding up", dim_old, " ", dim_new, " ", length(v_tot))

        H = zeros(T, dim, dim)

        # add old H elements
        @printf(" %-50s", "Fill old/old Hamiltonian: ")
        flush(stdout)
        @time _fill_H_block!(H, H_old, v_old, v_old, indices)

        @printf(" %-50s", "Build old/new Hamiltonian matrix with dimension: ")
        flush(stdout)
        @time Htmp = build_full_H_parallel(v_old, v_new, cluster_ops, clustered_ham)
        _fill_H_block!(H, Htmp, v_old, v_new, indices)
        _fill_H_block!(H, Htmp', v_new, v_old, indices)

        @printf(" %-50s", "Build new/new Hamiltonian matrix with dimension: ")
        flush(stdout)
        @time Htmp = build_full_H_parallel(v_new, v_new, cluster_ops, clustered_ham, sym=true)
        _fill_H_block!(H, Htmp, v_new, v_new, indices)

        vec_out = deepcopy(v_tot)
    else
        @printf(" %-50s", "Build full Hamiltonian matrix with dimension: ")
        @time H = build_full_H_parallel(ci_vector, ci_vector, cluster_ops, clustered_ham, sym=true)
    end

    flush(stdout)
    if length(vec_out) > 10
        @printf(" Now diagonalize\n")
        #@time e0,v = Arpack.eigs(H, nev = R, which=:SR)
        #@time e0,v = Arpack.eigs(H, nev = R, v0=get_vector(v_tot,root=1), which=:SR)
        davidson = FermiCG.Davidson(H, v0=get_vectors(ci_vector),
                                max_iter=max_iter, max_ss_vecs=max_ss_vecs, nroots=R, tol=conv_thresh)

        time = @elapsed e0,v = FermiCG.solve(davidson);
        @printf(" %-50s", "Diagonalization time: ")
        @printf("%10.6f seconds\n",time)
    else
        @printf(" %-50s", "Diagonalize: \n")
        @time F = eigen(H)
        e0 = F.values[1:R]
        v = F.vectors[:,1:R]
    end
    set_vector!(vec_out, v)

    clustered_S2 = extract_S2(ci_vector.clusters)
    @printf(" %-50s", "Compute S2 expectation values: ")
    @time s2 = compute_expectation_value_parallel(vec_out, cluster_ops, clustered_S2)
    flush(stdout)
    @printf(" %5s %12s %12s\n", "Root", "Energy", "S2")
    for r in 1:R
        @printf(" %5s %12.8f %12.8f\n",r, e0[r], abs(s2[r]))
    end

    if verbose > 1
        for r in 1:R
            display(vec_out, root=r)
        end
    end

    @printf(" ==================================================================|\n")
    return e0, vec_out, H
end
#=}}}=#

function _fill_H_block!(H_big, H_small, v_l,v_r, indices)
    #={{{=#
    # Fill H_big with elements from H_small
    idx_l = 1

    idx_l = zeros(Int,length(v_l))
    idx_r = zeros(Int,length(v_r))

    idx = 1
    for (fock,configs) in v_l.data
        for (config,coeff) in configs
            idx_l[idx] = indices[fock][config]
            idx += 1
        end
    end

    idx = 1
    for (fock,configs) in v_r.data
        for (config,coeff) in configs
            idx_r[idx] = indices[fock][config]
            idx += 1
        end
    end

    for (il,iil) in enumerate(idx_l)
        for (ir,iir) in enumerate(idx_r)
            H_big[iil,iir] = H_small[il,ir]
        end
    end
#    for (fock_l,configs_l) in v_l.data
#        for (config_l,coeff_l) in configs_l
#            idx_l_tot = indices[fock_l][config_l]
#
#            idx_r = 1
#            for (fock_r,configs_r) in v_r.data
#                for (config_r,coeff_r) in configs_r
#                    idx_r_tot = indices[fock_r][config_r]
#
#                    H_big[idx_l_tot, idx_r_tot] = H_small[idx_l, idx_r]
#
#                    idx_r += 1
#                end
#            end
#
#            idx_l += 1
#        end
#    end
end
#=}}}=#

"""
    tps_ci_davidson(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}

# Solve for eigenvectors/values in the basis defined by `ci_vector`. Use iterative davidson solver.
"""
function tps_ci_davidson(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
                        v0 = nothing,
                        conv_thresh = 1e-5,
                        max_ss_vecs = 12,
                        max_iter    = 40,
                        shift       = nothing,
                        precond     = false,
                        verbose     = 0) where {T,N,R}
    #={{{=#
    println()
    @printf(" |== Tensor Product State CI =======================================\n")
    vec_out = deepcopy(ci_vector)
    e0 = zeros(T,R)

    dim = length(ci_vector)
    iters = 0

    function matvec(v::AbstractMatrix)
        iters += 1
        #in = deepcopy(ci_vector)
        in = ClusteredState(ci_vector, R=size(v,2))
        set_vector!(in, v)
        #sig = deepcopy(in)
        #zero!(sig)
        #build_sigma!(sig, ci_vector, cluster_ops, clustered_ham, cache=cache)
        return tps_ci_matvec(in, cluster_ops, clustered_ham)
    end

    Hmap = FermiCG.LinOp(matvec, dim)

    davidson = FermiCG.Davidson(Hmap, v0=get_vectors(ci_vector),
                                max_iter=max_iter, max_ss_vecs=max_ss_vecs, nroots=R, tol=conv_thresh)

    e = nothing
    v = nothing
    if precond
        @printf(" %-50s", "Compute diagonal: ")
        clustered_ham_0 = extract_1body_operator(clustered_ham, op_string = "Hcmf")
        @time Hd = compute_diagonal(ci_vector, cluster_ops, clustered_ham_0)
        @printf(" %-50s", "Compute <0|H0|0>: ")
        @time E0 = compute_expectation_value_parallel(ci_vector, cluster_ops, clustered_ham_0)[1]
        @time Eref = compute_expectation_value_parallel(ci_vector, cluster_ops, clustered_ham)[1]
        Hd .+= Eref - E0
        @printf(" Now iterate: \n")
        flush(stdout)
        @time e,v = FermiCG.solve(davidson, Adiag=Hd);
    else
        @time e,v = FermiCG.solve(davidson);
    end
    set_vector!(vec_out, v)

    clustered_S2 = extract_S2(ci_vector.clusters)
    @printf(" %-50s", "Compute S2 expectation values: ")
    @time s2 = compute_expectation_value_parallel(vec_out, cluster_ops, clustered_S2)
    flush(stdout)
    @printf(" %5s %12s %12s\n", "Root", "Energy", "S2")
    for r in 1:R
        @printf(" %5s %12.8f %12.8f\n",r, e[r], abs(s2[r]))
    end

    if verbose > 1
        for r in 1:R
            display(vec_out, root=r)
        end
    end

    @printf(" ==================================================================|\n")
    return e, vec_out
end
#=}}}=#

"""
    tps_ci_matvec(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}

# Compute the action of `clustered_ham` on `ci_vector`.
"""
function tps_ci_matvec(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}
    #={{{=#

    jobs = []

    bra_idx = 0
    for (fock_bra, configs_bra) in ci_vector.data
        for (config_bra, coeff_bra) in configs_bra
            bra_idx += 1
            push!(jobs, (bra_idx, fock_bra, config_bra, coeff_bra, zeros(T,R)) )
        end
    end

    function do_job(job)
        fock_bra = job[2]
        config_bra = job[3]
        coeff_bra = job[4]
        sig_out = job[5]

        for (fock_trans, terms) in clustered_ham
            fock_ket = fock_bra - fock_trans

            haskey(ci_vector.data, fock_ket) || continue

            configs_ket = ci_vector[fock_ket]

            for (config_ket, coeff_ket) in configs_ket
                for term in clustered_ham[fock_trans]
                    check_term(term, fock_bra, config_bra, fock_ket, config_ket) || continue

                    #norm(term.ints)*maximum(abs.(coeff_ket)) > 1e-5 || continue
                    #@btime norm($term.ints)*maximum(abs.($coeff_ket)) > 1e-12

                    me = contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_ket, config_ket)
                    #if term isa ClusteredTerm4B
                    #    @btime contract_matrix_element($term, $cluster_ops, $fock_bra, $config_bra, $fock_ket, $config_ket)
                    #end
                    @simd for r in 1:R
                        @inbounds sig_out[r] += me * coeff_ket[r]
                    end
                    #@btime $sig_out .+= $me .* $ci_vector[$fock_ket][$config_ket]
                end

            end

        end
    end

    #for job in jobs
    Threads.@threads for job in jobs
        do_job(job)
        #@btime $do_job($job)
    end

    sigv = zeros(size(ci_vector))
    for job in jobs
        #for r in 1:R
        #    sigv[job[1],r] += job[5][r]
        #end
        sigv[job[1],:] .+= job[5]
    end

    return sigv
end
#=}}}=#

"""
    tpsci_ci(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
            thresh_cipsi = 1e-2,
            thresh_foi   = 1e-6,
            thresh_asci  = 1e-2,
            thresh_var   = -1.0,
            max_iter     = 10,
            conv_thresh  = 1e-4,
            nbody        = 4,
            incremental  = true,
            ci_conv      = 1e-5,
            ci_max_iter  = 50,
            ci_max_ss_vecs = 12,
            davidson     = false,
            max_mem_ci   = 20.0,
            matvec       = 1) where {T,N,R}

# Run TPSCI
- `thresh_cipsi`: threshold for which configurations to include in the variational space. Add if |c^{(1)}| > `thresh_cipsi`
- `thresh_foi`  : threshold for which terms to keep in the H|0> vector used to form the first order wavefunction
- `thresh_asci` : threshold for determining from which variational configurations  ``|c^{(0)}_i|`` > `thresh_asci`
- `thresh_var`  : threshold for clipping the result of the variational wavefunction. Not really needed default set to -1 (off)
- `max_iter`    : maximum selected CI iterations
- `conv_thresh` : stop selected CI iterations when energy change is smaller than `conv_thresh`
- `nbody`       : only consider up to `nbody` terms when searching for new configurations
- `incremental` : for the sigma vector incrementally between iterations
- `matvec`      : which implementation of the matrix vector code
- `ci_conv`     : convergence threshold for the inner CI step (only needed when davidson is used)
- `ci_max_iter` : max iterations for inner CI step (only needed when davidson is used)
- `ci_max_ss_vecs`: max subspace size for inner CI step (only needed when davidson is used)
- `davidson`    : use davidson? changes to true after needing more than max_mem_ci
- `max_mem_ci`  : maximum memory (Gb) allowed for storing full H. If more is needed, do Davidson.
"""
function tpsci_ci(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
    thresh_cipsi    = 1e-2,
    thresh_foi      = 1e-6,
    thresh_asci     = 1e-2,
    thresh_var      = -1.0,
    max_iter        = 10,
    conv_thresh     = 1e-4,
    nbody           = 4,
    incremental     = true,
    ci_conv         = 1e-5,
    ci_max_iter     = 50,
    ci_max_ss_vecs  = 12,
    davidson        = false,
    max_mem_ci      = 20.0,
    matvec          = 1) where {T,N,R}
#={{{=#
    vec_var = deepcopy(ci_vector)
    vec_pt = deepcopy(ci_vector)
    length(ci_vector) > 0 || error(" input vector has zero length")
    zero!(vec_pt)
    e0 = zeros(T,R)
    e2 = zeros(T,R)
    e0_last = zeros(T,R)

    clustered_S2 = extract_S2(ci_vector.clusters)

    println(" ci_vector     : ", size(ci_vector) )
    println(" thresh_cipsi  : ", thresh_cipsi   )
    println(" thresh_foi    : ", thresh_foi     )
    println(" thresh_asci   : ", thresh_asci    )
    println(" thresh_var    : ", thresh_var     )
    println(" max_iter      : ", max_iter       )
    println(" conv_thresh   : ", conv_thresh    )
    println(" nbody         : ", nbody          )
    println(" incremental   : ", incremental    )
    println(" ci_conv       : ", ci_conv        )
    println(" ci_max_iter   : ", ci_max_iter    )
    println(" ci_max_ss_vecs: ", ci_max_ss_vecs )
    println(" davidson      : ", davidson       )
    println(" max_mem_ci    : ", max_mem_ci     )
    println(" matvec        : ", matvec         )

    vec_asci_old = ClusteredState(ci_vector.clusters, R=R)
    sig = ClusteredState(ci_vector.clusters, R=R)
    clustered_ham_0 = extract_1body_operator(clustered_ham, op_string = "Hcmf")

    H = zeros(T,size(ci_vector))

    v_old = deepcopy(ci_vector)

    for it in 1:max_iter

        println()
        println()
        println(" ===================================================================")
        @printf("     Selected CI Iteration: %4i epsilon: %12.8f\n", it,thresh_cipsi)
        println(" ===================================================================")

        if it > 1
            l1 = length(vec_var)
            clip!(vec_var, thresh=thresh_var)
            l2 = length(vec_var)
            @printf(" Clip values < %8.1e         %6i → %6i\n", thresh_var, l1, l2)

            project_out!(vec_pt, vec_var)

            v_old = deepcopy(vec_var)
            v_new = deepcopy(vec_pt)

            l1 = length(vec_var)
            zero!(vec_pt)
            add!(vec_var, vec_pt)
            l2 = length(vec_var)

            @printf(" Add pt vector to current space %6i → %6i\n", l1, l2)
            length(v_old) + length(v_new) == l2 || error(" not adding up", length(v_old)+length(v_new), " ", l2)
        else
            rand!(vec_var)
        end

        e0 = nothing
        mem_needed = sizeof(T)*length(vec_var)*length(vec_var)*1e-9
        @printf(" Memory needed to hold full CI matrix: %12.8f (Gb)\n",mem_needed)
        flush(stdout)
        if (mem_needed > max_mem_ci) || davidson == true
            orthonormalize!(vec_var)
            e0, vec_var = tps_ci_davidson(vec_var, cluster_ops, clustered_ham,
                                   conv_thresh = ci_conv,
                                   max_iter = ci_max_iter,
                                   max_ss_vecs = ci_max_ss_vecs)
        else
            if it > 1
                # just update matrix
                e0, vec_var, H = tps_ci_direct(vec_var, cluster_ops, clustered_ham,
                                               H_old = H,
                                               v_old = v_old)
            else
                e0, vec_var, H = tps_ci_direct(vec_var, cluster_ops, clustered_ham)
            end
            #v_new = vec_var
#            e0, vec_var, H = tps_ci_direct(vec_var, cluster_ops, clustered_ham)
        end
        flush(stdout)

        # get barycentric energy <0|H0|0>
        Efock = compute_expectation_value_parallel(vec_var, cluster_ops, clustered_ham_0)
        #Efock = nothing
        flush(stdout)
        vec_asci = deepcopy(vec_var)
        l1 = length(vec_asci)
        clip!(vec_asci, thresh=thresh_asci)
        l2 = length(vec_asci)
        @printf(" Length of ASCI vector %8i → %8i \n", l1, l2)

        #
        # |sig_i> = H|v_i> = H|v_i-1> + H ( |v_i> - |v_i-1> )
        #                  = |sig_i-1> + H|del_i>

        del_v0 = deepcopy(vec_asci_old)
        ovlps = []
        for r in 1:R
            if dot(vec_asci_old, vec_asci, r, r) < 0
                scale!(vec_asci, -1.0, root=r)
            end
        end
        scale!(del_v0,-1.0)
        add!(del_v0, vec_asci)
        println(" Norm of delta v:")
        [@printf(" %12.8f\n",i) for i in norm(del_v0)]

        vec_asci_old = deepcopy(vec_asci)

        #_, vec_pt_del = compute_pt2(del_v0, cluster_ops, clustered_ham, E0=Efock, thresh_foi=thresh_foi, matvec=matvec, nbody=nbody)
        #add!(vec_pt, vec_pt_del)
        if incremental
            if matvec == 1
                del_sig_it = open_matvec_serial2(del_v0, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
            elseif matvec == 2
                del_sig_it = open_matvec_thread(del_v0, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
            elseif matvec == 3
                del_sig_it = open_matvec_thread2(del_v0, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
            else
                error("wrong matvec")
            end
            flush(stdout)
            add!(sig, del_sig_it)

            println(" Length of FOIS vector: ", length(sig))

            @printf(" %-50s", "Compute diagonal: ")
            flush(stdout)
            @time Hd = compute_diagonal(sig, cluster_ops, clustered_ham_0)

            sig_v = get_vectors(sig)
            v_pt  = zeros(size(sig_v))

            println()
            @printf(" %5s %12s %12s\n", "Root", "E(0)", "E(2)")
            for r in 1:R
                denom = 1.0 ./ (Efock[r] .- Hd)
                v_pt[:,r] .= denom .* sig_v[:,r]
                e2[r] = sum(sig_v[:,r] .* v_pt[:,r])

                @printf(" %5s %12.8f %12.8f\n",r, e0[r], e0[r] + e2[r])
            end
            flush(stdout)

            vec_pt = deepcopy(sig)
            set_vector!(vec_pt,v_pt)
        else
            e2, vec_pt = compute_pt2(vec_asci, cluster_ops, clustered_ham, E0=Efock, thresh_foi=thresh_foi, matvec=matvec, nbody=nbody)
        end

        l1 = length(vec_pt)
        clip!(vec_pt, thresh=thresh_cipsi)
        l2 = length(vec_pt)

        project_out!(sig, vec_asci)

        @printf(" Length of PT1  vector %8i → %8i \n", l1, l2)
        #add!(vec_var, vec_pt)

        if maximum(abs.(e0_last .- e0)) < conv_thresh
            print_tpsci_iter(vec_var, it, e0, true)
            break
        else
            print_tpsci_iter(vec_var, it, e0, false)
            e0_last .= e0
        end
    end
    return e0, vec_var
end
#=}}}=#

function print_tpsci_iter(ci_vector::ClusteredState{T,N,R}, it, e0, converged) where {T,N,R}
#={{{=#
    if converged
        @printf("*TPSCI Iter %-3i Dim: %-6i", it, length(ci_vector))
    else
        @printf(" TPSCI Iter %-3i Dim: %-6i", it, length(ci_vector))
    end
    @printf(" E(var): ")
    for i in 1:R
        @printf("%13.8f ", e0[i])
    end
#    @printf(" E(pt2): ")
#    for i in 1:R
#        @printf("%13.8f ", e2[i])
#    end
    println()
end
#=}}}=#

"""
    compute_pt2(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
        nbody=4,
        H0="Hcmf",
        E0=nothing, #pass in <0|H0|0>, or compute it
        thresh_foi=1e-8,
        prescreen=false,
        verbose=1,
        matvec=3) where {T,N,R}
"""
function compute_pt2(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator;
        nbody=4,
        H0="Hcmf",
        E0=nothing, #pass in <0|H0|0>, or compute it
        thresh_foi=1e-8,
        prescreen=false,
        verbose=1,
        matvec=3) where {T,N,R}
    #={{{=#

    println()
    println(" |............................do PT2................................")
    println(" thresh_foi    :", thresh_foi   )
    println(" prescreen     :", prescreen   )
    println(" H0            :", H0   )
    println(" nbody         :", nbody   )

    e2 = zeros(T,R)

    norms = norm(ci_vector);
    println(" Norms of input states")
    [@printf(" %12.8f\n",i) for i in norms]
    println(" Compute FOIS vector")

    if matvec == 1
        #@time sig = open_matvec(ci_vector, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
        sig = open_matvec_serial2(ci_vector, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi, prescreen=prescreen)
    elseif matvec == 2
        sig = open_matvec_thread(ci_vector, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
    elseif matvec == 3
        sig = open_matvec_thread2(ci_vector, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi, prescreen=prescreen)
    else
        error("wrong matvec")
    end
    #@time sig = open_matvec_parallel(ci_vector, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)
    #@btime sig = open_matvec_parallel($ci_vector, $cluster_ops, $clustered_ham, nbody=$nbody, thresh=$thresh_foi)
    println(" Length of FOIS vector: ", length(sig))

    clustered_ham_0 = extract_1body_operator(clustered_ham, op_string = H0)

    project_out!(sig, ci_vector)
    println(" Length of FOIS vector: ", length(sig))

    @printf(" %-50s", "Compute diagonal")
    @time Hd = compute_diagonal(sig, cluster_ops, clustered_ham_0)

    if E0 == nothing
        @printf(" %-50s", "Compute <0|H0|0>:")
        @time E0 = compute_expectation_value_parallel(ci_vector, cluster_ops, clustered_ham_0)
        #E0 = diag(E0)
        flush(stdout)
    end

    @printf(" %-50s", "Compute <0|H|0>:")
    @time Evar = compute_expectation_value_parallel(ci_vector, cluster_ops, clustered_ham)
    #Evar = diag(Evar)
    flush(stdout)

    sig_v = get_vectors(sig)
    v_pt  = zeros(size(sig_v))

    println()
    @printf(" %5s %12s %12s\n", "Root", "E(0)", "E(2)")
    for r in 1:R
        denom = 1.0 ./ (E0[r]/(norms[r]*norms[r]) .- Hd)
        v_pt[:,r] .= denom .* sig_v[:,r]
        e2[r] = sum(sig_v[:,r] .* v_pt[:,r])

        @printf(" %5s %12.8f %12.8f\n",r, Evar[r]/norms[r], Evar[r]/(norms[r]*norms[r]) + e2[r])
    end

    set_vector!(sig,v_pt)
    println(" ..................................................................|")

    return e2, sig
end
#=}}}=#

"""
    compute_expectation_value(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator; nbody=4) where {T,N,R}

Compute expectation value of a `ClusteredOperator` (`clustered_ham`) for state `ci_vector`
"""
function compute_expectation_value(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator; nbody=4) where {T,N,R}
    #={{{=#

    out = zeros(T,R)

    for (fock_bra, configs_bra) in ci_vector.data

        for (fock_ket, configs_ket) in ci_vector.data
            fock_trans = fock_bra - fock_ket

            # check if transition is connected by H
            haskey(clustered_ham, fock_trans) || continue

            for (config_bra, coeff_bra) in configs_bra
                for (config_ket, coeff_ket) in configs_ket

                    me = 0.0
                    for term in clustered_ham[fock_trans]

                        length(term.clusters) <= nbody || continue
                        check_term(term, fock_bra, config_bra, fock_ket, config_ket) || continue

                        me += contract_matrix_element(term, cluster_ops,
                                                      fock_bra, config_bra,
                                                      fock_ket, config_ket)
                    end

                    #out .+= coeff_bra .* coeff_ket .* me
                    for r in 1:R
                        out[r] += coeff_bra[r] * coeff_ket[r] * me
                    end

                end

            end
        end
    end

    return out
end
#=}}}=#

"""
    function compute_expectation_value_parallel(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}
"""
function compute_expectation_value_parallel(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}
    #={{{=#

    #
    # This will be were we collect our results
    evals = zeros(T,R)

    jobs = []

    for (fock_bra, configs_bra) in ci_vector.data
        for (config_bra, coeff_bra) in configs_bra
            push!(jobs, (fock_bra, config_bra, coeff_bra, zeros(T,R)) )
        end
    end

    function _add_val!(eval_job, me, coeff_bra, coeff_ket)
        for ri in 1:R
            #for rj in ri:R
            #    @inbounds eval_job[ri,rj] += me * coeff_bra[ri] * coeff_ket[rj]
            #    #eval_job[rj,ri] = eval_job[ri,rj]
            #end
            @inbounds eval_job[ri] += me * coeff_bra[ri] * coeff_ket[ri]
        end
    end

    function do_job(job)
        fock_bra = job[1]
        config_bra = job[2]
        coeff_bra = job[3]
        eval_job = job[4]
        ket_idx = 0

        for (fock_ket, configs_ket) in ci_vector.data
            fock_trans = fock_bra - fock_ket

            # check if transition is connected by H
            if haskey(clustered_ham, fock_trans) == false
                ket_idx += length(configs_ket)
                continue
            end

            for (config_ket, coeff_ket) in configs_ket
                #ket_idx += 1
                #ket_idx <= job[1] || continue

                me = 0.0
                for term in clustered_ham[fock_trans]

                    #length(term.clusters) <= 2 || continue
                    check_term(term, fock_bra, config_bra, fock_ket, config_ket) || continue

                    me += contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_ket, config_ket)
                    #if term isa ClusteredTerm4B
                    #    @btime contract_matrix_element($term, $cluster_ops, $fock_bra, $config_bra, $fock_ket, $config_ket)
                    #end
                    #Hrow[ket_idx] += me
                    #H[job[1],ket_idx] += me
                end
                #
                # now add the results
                #@inbounds for ri in 1:R
                #    @simd for rj in ri:R
                _add_val!(eval_job, me, coeff_bra, coeff_ket)
                #for ri in 1:R
                #    for rj in ri:R
                #        eval_job[ri,rj] += me * coeff_bra[ri] * coeff_ket[rj]
                #        #eval_job[rj,ri] = eval_job[ri,rj]
                #    end
                #end
            end
        end
    end

    #for job in jobs
    #Threads.@threads for job in jobs
    @qthreads for job in jobs
        do_job(job)
        #@btime $do_job($job)
    end

    for job in jobs
        evals .+= job[4]
    end

    return evals
end
#=}}}=#

"""
    open_matvec(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}

Compute the action of the Hamiltonian on a tpsci state vector. Open here, means that we access the full FOIS
(restricted only by thresh), instead of the action of H on v within a subspace of configurations.
This is essentially used for computing a PT correction outside of the subspace, or used for searching in TPSCI.
"""
function open_matvec(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}
#={{{=#
    println(" In open_matvec")
    sig = deepcopy(ci_vector)
    zero!(sig)
    clusters = ci_vector.clusters
    #sig = ClusteredState(clusters)
    #sig = OrderedDict{FockConfig{N}, OrderedDict{NTuple{N,Int16}, MVector{T} }}()
    for (fock_ket, configs_ket) in ci_vector.data
        for (ftrans, terms) in clustered_ham
            fock_bra = ftrans + fock_ket

            #
            # check to make sure this fock config doesn't have negative or too many electrons in any cluster
            all(f[1] >= 0 for f in fock_bra) || continue
            all(f[2] >= 0 for f in fock_bra) || continue
            all(f[1] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue
            all(f[2] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue

            #if haskey(sig, fock_bra) == false
            #    sig[fock_bra] = OrderedDict{NTuple{N,Int16}, MVector{T}}()
            #end
            haskey(sig, fock_bra) || add_fockconfig!(sig, fock_bra)
            for term in terms

                length(term.clusters) <= nbody || continue

                for (config_ket, coeff_ket) in configs_ket

                    sig_i = contract_matvec(term, cluster_ops, fock_bra, fock_ket, config_ket, coeff_ket, thresh=thresh)
                    #if term isa ClusteredTerm2B
                    #    @btime sig_i = contract_matvec($term, $cluster_ops, $fock_bra, $fock_ket, $config_ket, $coeff_ket, thresh=$thresh)
                    #end
                    #typeof(sig_i) == typeof(sig[fock_bra]) || println(typeof(sig_i), "\n",  typeof(sig[fock_bra]), "\n")

                    merge!(+, sig[fock_bra], sig_i)
                    #sig[fock_bra] = merge(+, sig[fock_bra], sig_i)

                    #for (config,coeff) in sig_i
                    #    #display(coeff[1])
                    #    #display(sig[fock_bra][config][1])
                    #    sig[fock_bra][config][1] += coeff[1]
                    #    #sig[fock_bra][config] = sig[fock_bra][config] + coeff
                    #end
                end
            end
        end
    end

    return sig
end
#=}}}=#

"""
    open_matvec_thread(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}

Compute the action of the Hamiltonian on a tpsci state vector. Open here, means that we access the full FOIS
(restricted only by thresh), instead of the action of H on v within a subspace of configurations.
This is essentially used for computing a PT correction outside of the subspace, or used for searching in TPSCI.

This parallellizes over FockConfigs in the output state, so it's not the most fine-grained, but it avoids data races in
filling the final vector
"""
function open_matvec_thread(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}
#={{{=#
    println(" In open_matvec_thread")
    sig = deepcopy(ci_vector)
    zero!(sig)
    clusters = ci_vector.clusters
    jobs = Dict{FockConfig{N},Vector{Tuple}}()
    #sig = ClusteredState(clusters)
    #sig = OrderedDict{FockConfig{N}, OrderedDict{NTuple{N,Int16}, MVector{T} }}()

    for (fock_ket, configs_ket) in ci_vector.data
        for (ftrans, terms) in clustered_ham
            fock_bra = ftrans + fock_ket

            #
            # check to make sure this fock config doesn't have negative or too many electrons in any cluster
            all(f[1] >= 0 for f in fock_bra) || continue
            all(f[2] >= 0 for f in fock_bra) || continue
            all(f[1] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue
            all(f[2] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue

            job_input = (terms, fock_ket, configs_ket)
            if haskey(jobs, fock_bra)
                push!(jobs[fock_bra], job_input)
            else
                jobs[fock_bra] = [job_input]
            end

        end
    end

    jobs_vec = []
    for (fock_bra, job) in jobs
        push!(jobs_vec, (fock_bra, job))
    end

    jobs_out = Vector{ClusteredState{T,N,R}}()
    for tid in 1:Threads.nthreads()
        push!(jobs_out, ClusteredState(clusters, T=T, R=R))
    end

    #println(" Number of jobs:    ", length(jobs))
    #println(" Number of threads: ", Threads.nthreads())
    BLAS.set_num_threads(1)
    #Threads.@threads for job in jobs_vec

    #for job in jobs_vec
    #@qthreads for job in jobs_vec
    Threads.@threads for job in jobs_vec
        fock_bra = job[1]
        sigi = _open_matvec_job(job[2], fock_bra, cluster_ops, nbody, thresh, N, R, T)
        tmp = jobs_out[Threads.threadid()]
        jobs_out[Threads.threadid()][fock_bra] = sigi
    end

    for threadid in 1:Threads.nthreads()
        #display(size(jobs_out[threadid]))
        add!(sig, jobs_out[threadid])
    end

    #BLAS.set_num_threads(Threads.nthreads())
    return sig
end
#=}}}=#

function _open_matvec_job(job, fock_bra, cluster_ops, nbody, thresh, N, R, T)
#={{{=#
    sigfock = OrderedDict{ClusterConfig{N}, MVector{R, T} }()

    for jobi in job

        terms, fock_ket, configs_ket = jobi

        for term in terms

            length(term.clusters) <= nbody || continue

            for (config_ket, coeff_ket) in configs_ket

                sig_i = contract_matvec(term, cluster_ops, fock_bra, fock_ket, config_ket, coeff_ket, thresh=thresh)
                #if term isa ClusteredTerm2B
                #    @btime sig_i = contract_matvec($term, $cluster_ops, $fock_bra, $fock_ket, $config_ket, $coeff_ket, thresh=$thresh)
                #    error("here")
                #end
                merge!(+, sigfock, sig_i)
            end
        end
    end
    return sigfock
end
#=}}}=#

"""
    open_matvec_parallel(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}

Compute the action of the Hamiltonian on a tpsci state vector. Open here, means that we access the full FOIS
(restricted only by thresh), instead of the action of H on v within a subspace of configurations.
This is essentially used for computing a PT correction outside of the subspace, or used for searching in TPSCI.

This parallellizes over FockConfigs in the output state, so it's not the most fine-grained, but it avoids data races in
filling the final vector
"""
function open_matvec_parallel(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}
#={{{=#

    sig = deepcopy(ci_vector)
    zero!(sig)
    clusters = ci_vector.clusters
    jobs = Dict{FockConfig{N},Vector{Tuple}}()

    println(" Copy data to each worker")
    @sync for pid in procs()
        @spawnat pid eval(:(ci_vector = deepcopy($ci_vector)))
        @spawnat pid eval(:(sig_job = ClusteredState($clusters, R=$R)))
        @spawnat pid eval(:(cluster_ops = $cluster_ops))
        @spawnat pid eval(:(clusters = $clusters))
        @spawnat pid eval(:(thresh = $thresh))
    end
    flush(stdout)

    println(" Collect jobs")
    @time for (fock_ket, configs_ket) in ci_vector.data
        for (ftrans, terms) in clustered_ham
            fock_bra = ftrans + fock_ket

            #
            # check to make sure this fock config doesn't have negative or too many electrons in any cluster
            all(f[1] >= 0 for f in fock_bra) || continue
            all(f[2] >= 0 for f in fock_bra) || continue
            all(f[1] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue
            all(f[2] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue

            job_input = (terms, fock_ket, configs_ket)
            if haskey(jobs, fock_bra)
                push!(jobs[fock_bra], job_input)
            else
                jobs[fock_bra] = [job_input]
            end

        end
    end

    jobs_vec = []
    for (fock_bra, job) in jobs
        push!(jobs_vec, (fock_bra, job))
    end

    jobs_out = Dict{Int, ClusteredState{T,N,R}}()
    for pid in procs()
        jobs_out[pid] = ClusteredState(clusters, T=T, R=R)
    end

    println(" Number of jobs:    ", length(jobs))

    #@sync @distributed for job in jobs_vec

    futures = []

    println(" Compute all jobs")
    @time @sync begin
        for job in jobs_vec
            fock_bra = job[1]
            future_sigi = @spawnat :any _open_matvec_job_parallel(job[2], fock_bra, nbody, thresh, N, R, T)
            #jobs_out[myid()][fock_bra] = sigi
            push!(futures, future_sigi)
        end
    end

    println(" Combine results")
    flush(stdout)
    @time for pid in procs()
        add!(sig, @fetchfrom pid sig_job)
    end

    #BLAS.set_num_threads(Threads.nthreads())
    return sig
end
#=}}}=#

function _open_matvec_job_parallel(job, fock_bra, nbody, thresh, N, R, T)
#={{{=#
    #sigfock = OrderedDict{ClusterConfig{N}, MVector{R, T} }()

    #sig = ClusteredState(clusters,R=R)
    add_fockconfig!(sig_job, fock_bra)

    for jobi in job

        terms, fock_ket, configs_ket = jobi

        for term in terms

            length(term.clusters) <= nbody || continue

            for (config_ket, coeff_ket) in configs_ket

                sig_i = contract_matvec(term, cluster_ops, fock_bra, fock_ket, config_ket, coeff_ket, thresh=thresh)
                #if term isa ClusteredTerm4B
                #    @btime sig_i = contract_matvec($term, $cluster_ops, $fock_bra, $fock_ket, $config_ket, $coeff_ket, thresh=$thresh)
                #    error("here")
                #end
                merge!(+, sig_job[fock_bra], sig_i)
            end
        end
    end
    return
end
#=}}}=#

"""
    open_matvec_parallel2(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}
"""
function open_matvec_parallel2(ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham; thresh=1e-9, nbody=4) where {T,N,R}
#={{{=#
    sig = deepcopy(ci_vector)
    zero!(sig)
    clusters = ci_vector.clusters

    println(" create empty sig vector on each worker")
    @sync for pid in procs()
        @spawnat pid eval(:(ci_vector = deepcopy($ci_vector)))
        @spawnat pid eval(:(sig_job = ClusteredState($clusters, R=$R)))
        @spawnat pid eval(:(cluster_ops = $cluster_ops))
        @spawnat pid eval(:(clusters = $clusters))
        @spawnat pid eval(:(thresh = $thresh))
    end
    println("done")
    flush(stdout)

    #jobs = Vector{Tuple{TransferConfig{N}, ClusteredTerm}}()
    #println(" Number of jobs: ", length(jobs))

    println(" Compute all jobs")
    futures = []
    @time @sync for (ftrans, terms) in clustered_ham
        for term in terms
            length(term.clusters) <= nbody || continue

            future = @spawnat :any _do_job(ftrans, term)
            #future = do_job(ftrans, term)
            push!(futures, future)
        end
    end

    println(" combine results")
    flush(stdout)
    @time @sync for pid in procs()
        add!(sig, @fetchfrom pid sig_job)
    end

#    n = length(futures)
#    @elapsed while n > 0 # print out results
#        add!(sig, take!(futures))
#        n = n - 1
#    end

    return sig
end
#=}}}=#

function _do_job(ftrans, term)
    for (fock_ket, configs_ket) in ci_vector.data
        fock_bra = ftrans + fock_ket

        #
        # check to make sure this fock config doesn't have negative or too many electrons in any cluster
        all(f[1] >= 0 for f in fock_bra) || continue
        all(f[2] >= 0 for f in fock_bra) || continue
        all(f[1] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue
        all(f[2] <= length(clusters[fi]) for (fi,f) in enumerate(fock_bra)) || continue

        haskey(sig_job, fock_bra) || add_fockconfig!(sig_job, fock_bra)

        for (config_ket, coeff_ket) in configs_ket

            sig_i = contract_matvec(term, cluster_ops, fock_bra, fock_ket, config_ket, coeff_ket, thresh=thresh)

            merge!(+, sig_job[fock_bra], sig_i)
        end
    end
    return
end

"""
    compute_diagonal(vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham) where {T,N,R}

Form the diagonal of the hamiltonan, `clustered_ham`, in the basis defined by `vector`
"""
function compute_diagonal(vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham) where {T,N,R}
    #={{{=#
    Hd = zeros(size(vector)[1])
    idx = 0
    zero_trans = TransferConfig([(0,0) for i in 1:N])
    for (fock_bra, configs_bra) in vector.data
        for (config_bra, coeff_bra) in configs_bra
            idx += 1
            for term in clustered_ham[zero_trans]
		    try
			    Hd[idx] += contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_bra, config_bra)
		    catch
			    display(term)
			    display(fock_bra)
			    display(config_bra)
			    error()
		    end

            end
        end
    end
    return Hd
end
#=}}}=#

"""
    compute_diagonal!(Hd, vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham) where {T,N,R}

Form the diagonal of the hamiltonan, `clustered_ham`, in the basis defined by `vector`
"""
function compute_diagonal!(Hd, vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham) where {T,N,R}
    #={{{=#
    idx = 0
    zero_trans = TransferConfig([(0,0) for i in 1:N])
    for (fock_bra, configs_bra) in vector.data
        for (config_bra, coeff_bra) in configs_bra
            idx += 1
            for term in clustered_ham[zero_trans]
		    try
			    Hd[idx] += contract_matrix_element(term, cluster_ops, fock_bra, config_bra, fock_bra, config_bra)
		    catch
			    display(term)
			    display(fock_bra)
			    display(config_bra)
			    error()
		    end

            end
        end
    end
    return
end
#=}}}=#

"""
    expand_each_fock_space!(s::ClusteredState{T,N,R}, bases::Vector{ClusterBasis}) where {T,N,R}

For each fock space sector defined, add all possible basis states
- `basis::Vector{ClusterBasis}`
"""
function expand_each_fock_space!(s::ClusteredState{T,N,R}, bases::Vector{ClusterBasis}) where {T,N,R}
    # {{{
    println("\n Make each Fock-Block the full space")
    # create full space for each fock block defined
    for (fblock,configs) in s.data
        #println(fblock)
        dims::Vector{UnitRange{Int16}} = []
        #display(fblock)
        for c in s.clusters
            # get number of vectors for current fock space
            dim = size(bases[c.idx][fblock[c.idx]], 2)
            push!(dims, 1:dim)
        end
        for newconfig in product(dims...)
            #display(newconfig)
            #println(typeof(newconfig))
            #
            # this is not ideal - need to find a way to directly create key
            config = ClusterConfig(collect(newconfig))
            s.data[fblock][config] = zeros(SVector{R,T})
            #s.data[fblock][[i for i in newconfig]] = 0
        end
    end
end
# }}}

"""
    expand_to_full_space!(s::AbstractState, bases::Vector{ClusterBasis}, na, nb)

Define all possible fock space sectors and add all possible basis states
- `basis::Vector{ClusterBasis}`
- `na`: Number of alpha electrons total
- `nb`: Number of alpha electrons total
"""
function expand_to_full_space!(s::AbstractState, bases::Vector{ClusterBasis}, na, nb)
    # {{{
    println("\n Expand to full space")
    ns = []

    for c in s.clusters
        nsi = []
        for (fspace,basis) in bases[c.idx]
            push!(nsi,fspace)
        end
        push!(ns,nsi)
    end
    for newfock in product(ns...)
        nacurr = 0
        nbcurr = 0
        for c in newfock
            nacurr += c[1]
            nbcurr += c[2]
        end
        if (nacurr == na) && (nbcurr == nb)
            config = FockConfig(collect(newfock))
            add_fockconfig!(s,config)
        end
    end
    expand_each_fock_space!(s,bases)

    return
end
# }}}

"""
    project_out!(v::ClusteredState, w::ClusteredState)

Project w out of v
    |v'> = |v> - |w><w|v>
"""
function project_out!(v::ClusteredState, w::ClusteredState)
    for (fock,configs) in w.data
        if haskey(v.data, fock)
            for (config, coeff) in configs
                if haskey(v.data[fock], config)
                    delete!(v.data[fock], config)
                end
            end
            if length(v[fock]) == 0
                delete!(v.data, fock)
            end
        end
    end
    # I'm not sure why this is necessary
    idx = 0
    for (fock,configs) in v.data
        for (config, coeffs) in v.data[fock]
            idx += 1
        end
    end
end

"""
    hosvd(ci_vector::ClusteredState{T,N,R}, cluster_ops; hshift=1e-8, truncate=-1) where {T,N,R}

Peform HOSVD aka Tucker Decomposition of ClusteredState
"""
function hosvd(ci_vector::ClusteredState{T,N,R}, cluster_ops; hshift=1e-8, truncate=-1) where {T,N,R}
#={{{=#

    cluster_rotations = []
    for ci in ci_vector.clusters
        println()
        println(" --------------------------------------------------------")
        println(" Density matrix: Cluster ", ci.idx)
        println()
        println(" Compute BRDM")
        println(" Hshift = ",hshift)

        dims = Dict()
        for (fock, mat) in cluster_ops[ci.idx]["H"]
            fock[1] == fock[2] || error("?")
            dims[fock[1]] = size(mat,1)
        end

        rdms = build_brdm(ci_vector, ci, dims)
        norm = 0
        entropy = 0
        rotations = Dict{Tuple,Matrix{T}}()
        for (fspace,rdm) in rdms
            fspace_norm = 0
            fspace_entropy = 0
            @printf(" Diagonalize RDM for Cluster %2i in Fock space: ",ci.idx)
            println(fspace)
            F = eigen(Symmetric(rdm))

            idx = sortperm(F.values, rev=true)
            n = F.values[idx]
            U = F.vectors[:,idx]

            # Either truncate the unoccupied cluster states, or remix them with a hamiltonian to be unique
            if truncate < 0
                remix = []
                for ni in 1:length(n)
                    if n[ni] < 1e-8
                        push!(remix, ni)
                    end
                end
                U2 = U[:,remix]
                Hlocal = U2' * cluster_ops[ci.idx]["H"][(fspace,fspace)] * U2

                F = eigen(Symmetric(Hlocal))
                n2 = F.values
                U2 = U2 * F.vectors

                U[:,remix] .= U2[:,:]

            else
                keep = []
                for ni in 1:length(n)
                    if abs(n[ni]) > truncate
                        push!(keep, ni)
                    end
                end
                @printf(" Truncated Tucker space. Starting: %5i Ending: %5i\n" ,length(n), length(keep))
                U = U[:,keep]
            end

            n = diag(U' * rdm * U)
            Elocal = diag(U' * cluster_ops[ci.idx]["H"][(fspace,fspace)] * U)

            norm += sum(n)
            fspace_norm = sum(n)
            @printf("                 %4s:    %12s    %12s\n", "","Population","Energy")
            for (ni_idx,ni) in enumerate(n)
                if abs(ni/norm) > 1e-16
                    fspace_entropy -= ni*log(ni/norm)/norm
                    entropy -=  ni*log(ni)
                    @printf("   Rotated State %4i:    %12.8f    %12.8f\n", ni_idx,ni,Elocal[ni_idx])
                end
           end
           @printf("   ----\n")
           @printf("   Entanglement entropy:  %12.8f\n" ,fspace_entropy)
           @printf("   Norm:                  %12.8f\n" ,fspace_norm)

           #
           # let's just be careful that our vectors remain orthogonal
           F = svd(U)
           U = F.U * F.Vt
           check_orthogonality(U)
           rotations[fspace] = U
        end
        @printf(" Final entropy:.... %12.8f\n",entropy)
        @printf(" Final norm:....... %12.8f\n",norm)
        @printf(" --------------------------------------------------------\n")

        flush(stdout)

        #ci.rotate_basis(rotations)
        #ci.check_basis_orthogonality()
        push!(cluster_rotations, rotations)
    end
    return cluster_rotations
end
#=}}}=#

"""
    build_brdm(ci_vector::ClusteredState, ci, dims)

Build block reduced density matrix for `Cluster`,  `ci`
- `ci_vector::ClusteredState` = input state
- `ci` = Cluster type for whihch we want the BRDM
- `dims` = list of dimensions for each fock sector
"""
function build_brdm(ci_vector::ClusteredState, ci, dims)
    # {{{
    rdms = OrderedDict()
    for (fspace, configs) in ci_vector.data
        curr_dim = dims[fspace[ci.idx]]
        rdm = zeros(curr_dim,curr_dim)
        for (configi,coeffi) in configs
            for cj in 1:curr_dim

                configj = [configi...]
                configj[ci.idx] = cj
                configj = ClusterConfig(configj)

                if haskey(configs, configj)
                    rdm[configi[ci.idx],cj] += sum(coeffi.*configs[configj])
                end
            end
        end

        if haskey(rdms, fspace[ci.idx])
            rdms[fspace[ci.idx]] += rdm
        else
            rdms[fspace[ci.idx]] = rdm
        end

    end
    return rdms
end
# }}}

function dump_tpsci(filename::AbstractString, ci_vector::ClusteredState{T,N,R}, cluster_ops, clustered_ham::ClusteredOperator) where {T,N,R}
    @save filename ci_vector cluster_ops clustered_ham
end

#function load_tpsci(filename::AbstractString)
#    a = @load filename
#    return eval.(a)
#end

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
function demo()
    println("+-------------------------------------------+")
    println("|                                           |")
    println("|     LightML.jl demo for digits dataset    |")
    println("|                                           |")
    println("+-------------------------------------------+")
    # ...........
    #  LOAD DATA
    # ...........
    data = dat.load_digits()
    digit1 = 1
    digit2 = 8
    idx = vcat(find(data["target"] .== digit1), find(data["target"] .== digit2))
    y = data["target"][idx]
    # Change labels to {0, 1}
    y[y .== digit1] = -1
    y[y .== digit2] = 1

    X = data["data"][idx,:]
    X = normalize_(X)

    println("Dataset: The Digit Dataset (digits $(digit1) and $(digit2))" )
    # ..........................
    #  DIMENSIONALITY REDUCTION
    # ..........................
    pca = PCA(n_components=5)
    train!(pca, X) # Reduce to 5 dimensions
    X = transform(pca, X)

    # ..........................
    #  TRAIN / TEST SPLIT
    # ..........................
    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)
    rescaled_y_train = (y_train + 1) ./ 2
    rescaled_y_test = (y_test + 1) ./2
    # .......
    #  SETUP
    # .......
    adaboost = Adaboost(n_clf = 10)
    naive_bayes = NaiveBayes()
    knn = KnnClassifier(k=4)
    logistic_regression = LogisticRegression()
    #mlp = NeuralNetwork(hidden = 3)
    decision_tree = ClassificationTree()
    random_forest = randomForest(n_estimators=20)
    support_vector_machine_rbf = svm(X_train, y_train, max_iter=500, kernel="rbf", C=0.6)
    support_vector_machine_linear = svm(X_train, y_train, max_iter=500, kernel="linear", C=0.6)
    #lda = LDA()
    gbc = GradientBoostingClassifier(n_clf=20, learning_rate=.9, max_depth=2)
    #xgboost = XGBoost(n_estimators=50, learning_rate=0.5)

    # ........
    #  TRAIN
    # ........

    println("\tk-nearst")
    train!(knn, X_train, y_train)
    println("\tAdaboost")
    train!(adaboost,X_train, y_train)
    println("\tDecision Tree")
    train!(decision_tree,X_train, y_train)
    println("\tGradient Boosting")
    train!(gbc,X_train, y_train)
    #println("\tLDA")
    #train!(lda,X_train, y_train)
    println("\tLogistic Regression")
    train!(logistic_regression,X_train, rescaled_y_train)
    #println("\tMultilayer Perceptron")
    #train!(mlp,X_train, y_train)
    println("\tNaive Bayes")
    train!(naive_bayes,X_train, y_train)
    #println("\tPerceptron")
    #train!(perceptron,X_train, y_train)
    println("\tRandom Forest")
    train!(random_forest,X_train, y_train)
    println("\tSupport Vector Machine with rbf kernel")
    train!(support_vector_machine_rbf)
    println("\tSupport Vector Machine with linear kernel")
    train!(support_vector_machine_linear)
    #println("\tXGBoost")
    #xgboost.fit(X_train, y_train)

    # .........
    #  PREDICT
    # .........
    y_pred = Dict()
    y_pred["Adaboost"] = predict(adaboost, X_test)
    y_pred["Gradient Boosting"] = predict(gbc,X_test)
    y_pred["Naive Bayes"] = predict(naive_bayes, X_test)
    y_pred["K Nearest Neighbors"] = predict(knn, X_test)
    y_pred["Logistic Regression"] = predict(logistic_regression, X_test)
    #y_pred["LDA"] = predict(lda, X_test)
    #y_pred["Multilayer Perceptron"] = sign(predict(mlp, X_test))
    #y_pred["Perceptron"] = perceptron.predict(X_test)
    y_pred["Decision Tree"] = predict(decision_tree, X_test)
    y_pred["Random Forest"] = predict(random_forest, X_test)
    y_pred["Support Vector Machine, rbf kernel"] = predict(support_vector_machine_rbf, X_test)
    y_pred["Support Vector Machine, linear kernel"] = predict(support_vector_machine_linear, X_test)

    #y_pred["XGBoost"] = xgboost.predict(X_test)

    # ..........
    #  ACCURACY
    # ..........
    println("Accuracy:")
    for clf in keys(y_pred)
        if clf == "Logistic Regression"
            temp = accuracy(rescaled_y_test,y_pred[clf])
            st = temp * 100
            st = @sprintf "%.3f" st
            println("\t$(clf): $(st) %")
        else
            temp = accuracy(y_test,y_pred[clf])
            st = temp * 100
            st = @sprintf "%.3f" st
            println("\t$(clf): $(st) %")
        end
    end

    # .......
    #  PLOT
    # .......

    x1 = X_test[:, 1]
    x2 = X_test[:, 2]
    df = DataFrame(x = x1, y = x2, clu = y_test)
    Gadfly.plot(df, x = "x", y = "y", color = "clu", Geom.point, Guide.title("PCA"))
end

ext: jl


text:
"""
  chunkSharedSV!(f,philatent,Y,T,A,D,nj,nwinc,ninc,nw,n,alpha,beta,indt,indnw,indn; verbose=false)

Iterate over chunk `p` of data on a SharedArray, updating copies of sufficient stats: `nj`,`nwinc`,`n`, and assignments: `A` for all tokens.

Sample the topic index `z ~ Z` from the sufficient stats of token from document `j` corresponding to lexicon term index `i`.
```math
p(z=k) \\propto \\frac{nw^{\\lnot ij}+\\beta[i]}{n^{\\lnot ij} + W\\beta[i]} (nj^{\\lnot ij} + \\alpha)
```
The function \$ f(A_{\\lnot i},Y_{\\lnot j})\$ is applied to the assignments (`A`) and response data `Y` during sampling of \$ a_{i,j}\$ for term \$i\$ of document \$j\$.

The function `philatent(\beta,nw)` denotes prior knowledge of the word/topic associations.

# Examples
```
chunkSharedSV!(f,philatent,Y,T,A,D,nj,nwinc,ninc,nw,n,alpha,beta,indt,indnw,indn; verbose=false)
```
"""
function chunkSharedSV!(philatent,T,A,D,nj,nwinc,nw,alpha,beta,indt,indnw,nlex; verbose=false)
  indnw_linear = LinearIndices(nw)[indnw...]
  for i in 1:length(indt)
    t = T[indt[i]]
    a = A[indt[i]]
    d = D[indt[i]]

    # take current doc/term complement of nj,nw
    njI = LinearIndices(nj)
    nji = njI[CartesianIndex(a,d)]
    nj_old = nj[nji] # copy old for debug
    if verbose
      println("attempting to access a:$a, t:$t of LinearIndices(nw)[indnw...]:")
    end
    nwi = indnw_linear[CartesianIndex(a,t)]
    nw_old = nw[nwi] # copy old for debug
    if verbose
      println("\n \n")
      println("sampling chunk token $i of ",length(indt)," from doc $d")
      println("nji=$nji,nwi=$nwi")
      println("nj[nji]=",nj[nji])
      println("nw[nwi]=",nw[nwi])
      println("nj[:,d]=",nj[:,d])
      println("nw[:,t]=",nw[:,t])
      println("indnw:\n",indnw)
      println("size(nw[indnw...]):",size(nw[indnw...]))
      println("size(nw):",size(nw))
      println("beta[:,t]:",beta[:,t])
    end
    @assert nj_old > 0 "nji must be positive"
    @assert nw_old > 0 "nwi must be positive"

    nj[nji] -= 1
    nw[nwi] -= 1

    # sample the new assignment
    @assert philatent in [2,1,0] "latent phi needs to be specified"
    if philatent == 0
      wprob = beta[:,t]
    elseif philatent == 1
      wprob = beta[:,t] ./ (beta[:,t] .* nlex)
    elseif philatent == 2
      wprob = (nw[indnw...][:,t] .+ beta[:,t]) ./ sum(nw[indnw...] .+ beta[:,t],dims=2)
    end

    dprob = (nj[:,d] .+ vec(alpha)) ./ sum(nj[:,d] .+ vec(alpha))
    z = dprob .* wprob

    if verbose
      println("\n \n")
      println("z:",z)
      println("wprob:",wprob)
      println("dprob:",dprob)
    end

    samp = rand()*sum(z)
    a_new = findfirst(x->x>0,vec(cumsum(z,dims=1)) .- samp)

    if verbose
      println("\n \n")
      println("a_new:",a_new)
      println("a:",a)
      println("d:",d)
      println("njI:",njI)
      println("CartesianIndex(a_new,d):",CartesianIndex(a_new,d))
    end

    if isnothing(a_new)
      a_new = a
    end

    # update nj, nw
    nji_new = njI[CartesianIndex(a_new,d)]
    nwi_new = indnw_linear[CartesianIndex(a_new,t)]
    nj[nji_new] += 1
    nw[nwi_new] += 1
    A[indt[i]] = a_new

    # record assignment nwinc
    nwinc_old = nwinc[nwi_new]
    nwinc[nwi_new] += 1

    if verbose
      println("\n \n")
      println("A=$a->$a_new")
      println("nwi=$nwi,size(nw[indnw...])=",size(nw[indnw...]))
      println("nwinc update $nwi_new: $nwinc_old->",nwinc[nwi_new])
    end

    if verbose
      println("nj update: doc=$d, topic $nj_old -> ",nj[nji])
      println("nw update: term=$t, topic $nw_old -> ",nw[nwi])
      # sleep(0.2)
    end
  end
end

"""
  sampleSharedSV!(f,philatent,Y,T,A,D,nj,nwinc,ninc,nw,n,alpha,beta,indt,indnw,indn,k,l,wids; verbose=false)

Parallel loop over SharedArray data, iterating locally on each worker in `wids`.
See `initShared` for parameter definitions: `T`,`A`,`D`,`nj`,`nwinc`,`ninc`,`nw`,`n`,`alpha`,`beta`.
Additional parameters: number of topics `k`, number of lexicon terms `l`.

The function \$ f(A_{\\lnot i},Y_{\\lnot j})\$ is applied to the assignments (`A`) and orthogonal data `Y` during sampling of \$ a_{i,j}\$ for term \$i\$ of document \$j\$.

The function `philatent(\beta,nw)` denotes prior knowledge of the word/topic associations.

# Examples
```
sampleSharedSV!(f,philatent,Y,T,A,D,nj,nwinc,ninc,nw,n,alpha,beta,indt,indnw,indn,k,l,wids; verbose=false)
```
"""
function sampleSharedSV!(philatent,T,A,D,nj,nwinc,nw,alpha,beta,indt,indnw,k,nlex,wids; verbose=false)
  nw_global = zeros(k,nlex)
  # sample the assignments

  # test
  if verbose
    for i in 1:length(wids)
      println("i:",i)
      println("indt[i]:",indt[i])
      println("indnw[i]:",indnw[i])
      chunkSharedSV!(philatent,T,A,D,nj,nwinc,nw,alpha,beta,indt[i],indnw[i],nlex,verbose=verbose)
    end
  else
    @sync begin
      for i in 1:length(wids)
        @async remotecall_wait(chunkSharedSV!,wids[i],philatent,T,A,D,nj,nwinc,nw,alpha,beta,indt[i],indnw[i],nlex,verbose=verbose)
      end
    end
  end

  # synchronize global nw,n
  for i in 1:length(wids)
    nw_global += nwinc[indnw[i]...]
  end
  for i in 1:length(wids)
      nw[indnw[i]...] .= nw_global
  end
end

"""
  runChainSharedSV(corpus,partition,partmap,betalocal,alphalocal,k,l,drnm,nlex,stateint,statefloat,philatent,f; verbose=false)

Parallel loop over the data, iterating locally on each worker.
See `initDistributed` for parameter definitions: `T`,`A`,`D`,`nj`,`nwinc`,`ninc`,`nw`,`n`,`alpha`,`beta`.
Additional parameters: number of topics `k`, number of lexicon terms `l`.

# Examples
```
runChainSharedSV(corpus,partition,partmap,betalocal,alphalocal,k,l,drnm,nlex,stateint,statefloat,philatent,f; verbose=false)
```
"""

function runChainSharedSV(corpus0,partition,partmap,betalocal,alphalocal,
                            k,nlex,philatent,
                            stateint,statefloat,nworkers_local,chainblockn,n_iter,
                            chain_i,drnm,initcorp,Thetas,indthin; verbose=false)

  ccopy = deepcopy(corpus0)
  if initcorp == "beta"
    corpus = initcorpBeta(corpus0,betalocal)
  elseif initcorp == "chain"
    corpus = initcorpChain(corpus0,chain_i,k)
  elseif initcorp == "unif"
    corpus = initcorpUnif(corpus0,k)
  end

  @assert ccopy != corpus "re-initialization failed"
  println("starting chain $chain_i")

  println("begin initializing state arrays: T,A,D,Z,J,W")
  T,A,D,Z,J,W,nj,nwinc,nw,alpha,beta,indt,indnw = initShared(corpus,partition,k,nlex,alphalocal,betalocal,sstatint=Int64)

  @assert size(nw) == (k,nlex*nworkers_local) "nw: incorrect shape or number of distributed copies, topics"
  @assert size(nj) == (k,length(corpus)) "nj: incorrect shape or number of documents, topics"

  wids = workers()[1:nworkers_local]

  thin_ind = 1 # the index of the saved thinned samples
  @showprogress 1 "Running $n_iter samples..." for i in 1:Int(n_iter)
      # println("iteration $i")
      sampleSharedSV!(philatent,T,A,D,nj,nwinc,nw,alpha,beta,indt,indnw,k,nlex,wids; verbose=verbose)
      if i in indthin
        @sync @distributed for doc_id = 1:length(partmap)
          thetaIter!(Thetas,A,J,doc_id,alpha,thin_ind,chain_i,k)
        end
        thin_ind+=1
      end
  end
  println("sampling: chain $chain_i mcmc has finished")
end


text:
using Random
using StaticArrays: SVector

function test_interpolation(ord::BSplineOrder, ::Type{Ty} = Float64) where {Ty}
    rng = MersenneTwister(42)
    ndata = 40
    xs = sort(randn(rng, ndata))

    # This is Int   if Ty = Int
    #         Int32 if Ty = SVector{N, Int32}
    Tdata = eltype(Ty)

    ys = if Tdata <: Integer
        [rand(rng, Ty) .% Ty(16) for _ = 1:ndata]  # values in -15:15
    else
        randn(rng, Ty, ndata)
    end

    k = order(ord)

    @testset "No BCs" begin
        S = @inferred interpolate(xs, ys, ord)
        @test S.(xs) ≈ ys
    end

    if iseven(k)
        @testset "Natural BCs" begin
            S = @inferred interpolate(xs, ys, ord, Natural())
            @test S.(xs) ≈ ys
            ts = @inferred knots(S)
            Nt = length(ts)
            ts_unique = view(ts, k:(Nt - k + 1))
            @test allunique(ts_unique)

            # Check that unique knots are the same as data locations (this is not
            # the case without natural BCs)
            @test length(ts_unique) == length(xs)
            @test ts_unique == xs

            # Check that some derivatives are zero.
            for n = 2:(k ÷ 2)
                Sder = Derivative(n) * S
                for x ∈ boundaries(basis(S))
                    @test norm(Sder(x)) < norm(S(x)) * 1e-7
                end
            end
        end
    end

    nothing
end

@testset "Interpolation" begin
    @testset "k = $k" for k ∈ (3, 4, 6, 8)
        test_interpolation(BSplineOrder(k), Float64)
    end
    types = (Int32, ComplexF32, SVector{2, Float32})
    @testset "T = $T" for T ∈ types
        test_interpolation(BSplineOrder(4), T)
    end
end


text:
module ForwardOptimalControl

import ..DynamicsModelInterface
import ..JuMPUtils
import Ipopt
import JuMP

using JuMP: @variable, @constraint, @objective
using UnPack: @unpack

export forward_quadratic_objective, linear_dynamics_constraints, solve_lqr, solve_optimal_control

#============================================= LQ case =============================================#

"The performance index for the forward optimal control problem."
function forward_quadratic_objective(x, u; Q, R)
    T = last(size(x))
    sum(x[:, t]' * Q * x[:, t] + u[:, t]' * R * u[:, t] for t in 1:T)
end

function linear_dynamics_constraints(x, u; A, B)
    reduce(hcat, ((x[:, t + 1] - A[t] * x[:, t] - B[t] * u[:, t]) for t in axes(x)[2][1:(end - 1)]))
end

"Solves a forward LQR problem using JuMP."
function solve_lqr(
    A,
    B,
    Q,
    R,
    x0,
    T;
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
)
    n_states, n_controls = size(only(unique(B)))
    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])
    @constraint(opt_model, linear_dynamics_constraints(x, u; A, B) .== 0)
    @constraint(opt_model, x[:, 1] .== x0)
    @objective(opt_model, Min, forward_quadratic_objective(x, u; Q, R))
    @time JuMP.optimize!(opt_model)
    JuMPUtils.isconverged(opt_model), JuMPUtils.get_values(; x, u), opt_model
end

#=========================================== Non-LQ-Case ===========================================#

"Solves a forward optimal control problem with protentially nonlinear dynamics and nonquadratic
costs using JuMP."
function solve_optimal_control(
    control_system,
    cost_model,
    x0,
    T;
    fixed_inputs = (),
    init = (),
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
    verbose = false,
)
    @unpack n_states, n_controls = control_system

    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    # decision variables
    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])

    # initial guess
    JuMPUtils.init_if_hasproperty!(x, init, :x)
    JuMPUtils.init_if_hasproperty!(u, init, :u)

    # fix certain inputs
    for i in fixed_inputs
        JuMP.fix.(u[i, :], init.u[i, :])
    end

    DynamicsModelInterface.add_dynamics_constraints!(control_system, opt_model, x, u)
    if !isnothing(x0)
        @constraint(opt_model, x[:, 1] .== x0)
    end
    cost_model.add_objective!(opt_model, x, u; cost_model.weights)
    time = @elapsed JuMP.optimize!(opt_model)
    verbose && @info time

    JuMPUtils.isconverged(opt_model), JuMPUtils.get_values(; x, u), opt_model
end

end


text:
export register,
    protocol,
    protocol!,
    PROTOCOLS,
    key,
    astuple,
    DefaultStatus

immutable type DefaultStatus
    default::Integer
    bad_request::Integer
    not_found::Integer
    system_error::Integer
end

DefaultStatus(;kwargs...) = begin
    o = DefaultStatus(-1, -1, -1, -1)
    for (k, v) in kwargs
        setfield(o, k, v)
    end
    o
end

immutable type Protocol{SCK<:Number,SCV<:STR_TYPE}
    name::STR_TYPE
    version::STR_TYPE
    status_codes::Associative{SCK, SCV}
    status_wo_msg::Array{SCK}
    supports::Array{STR_TYPE}
    prepare_res_headers::Function
    default_status::DefaultStatus
    header_msg_chunk::(STR_TYPE, STR_TYPE)
    header_msg_type::STR_TYPE
    header_msg_size::STR_TYPE
    req_post_init::Function
    res_post_init::Function
end

key(p::Protocol) = (p.name, p.version)

immutable type Protocols
    data::Associative{(STR_TYPE,STR_TYPE), Protocol}
end

Protocols() = Protocols(Dict{(STR_TYPE,STR_TYPE), Protocol}())

const PROTOCOLS = Protocols()

status{K<:Integer, V<:STR_TYPE}(args::(K, V)...) = Dict{Integer, STR_TYPE}(args)
status{K<:Integer, V<:STR_TYPE}(d::Dict{K, V}) = convert(Dict{Integer, STR_TYPE}, d)

astuple(p::Protocol) = (p.name, p.version)

protocol(d::Protocols, name_version::(STR_TYPE,STR_TYPE)) = d.data[name_version]
protocol(name_version::(STR_TYPE,STR_TYPE)) = protocol(PROTOCOLS, name_version)

protocol!(d::Protocols, p::Protocol) = (d.data[(p.name, p.version)] = p)
protocol!(p::Protocol) = protocol!(PROTOCOLS, p)

protocol!(d::Protocols, name_version::(STR_TYPE,STR_TYPE), p::Protocol) = (d.data[name_version] = p)
protocol!(name_version::(STR_TYPE,STR_TYPE), p::Protocol) = protocol!(PROTOCOLS, name_version, p)


text:
import OnlineStats: OnlineStat, Extrema, fit!, value
import ...Cubes.Axes: CategoricalAxis, RangeAxis
import IterTools
using WeightedOnlineStats
import ProgressMeter: next!

import WeightedOnlineStats: WeightedOnlineStat
abstract type TableAggregator end
struct OnlineAggregator{O,S}<:TableAggregator
    o::O
end
function OnlineAggregator(O::OnlineStat,s::Symbol)
    OnlineAggregator{typeof(O),s}(copy(O))
end
cubeeltype(t::OnlineAggregator)=Float64
function fitrow!(o::OnlineAggregator{T,S},r) where {T<:OnlineStat,S}
    v = getproperty(r,S)
    !ismissing(v) && fit!(o.o,v)
end
value(o::OnlineAggregator)=value(o.o)
struct WeightOnlineAggregator{O,S,W}<:TableAggregator
    o::O
    w::W
end
function WeightOnlineAggregator(O::WeightedOnlineStat,s::Symbol,w)
    WeightOnlineAggregator{typeof(O),s,typeof(w)}(copy(O),w)
end
cubeeltype(t::WeightOnlineAggregator{T}) where T = cubeeltype(T)
value(o::WeightOnlineAggregator)=value(o.o)
function fitrow!(o::WeightOnlineAggregator{T,S},r) where {T<:OnlineStat,S}
    v = getproperty(r,S)
    w = o.w(r)
    if !checkmiss(v) && !ismissing(w)
      fit!(o.o,v,w)
    end
end
checkmiss(v) = ismissing(v)
checkmiss(v::AbstractVector) = any(ismissing,v)
struct GroupedOnlineAggregator{O,S,BY,W,C}<:TableAggregator
    d::O
    w::W
    by::BY
    cloneobj::C
end
value(o::GroupedOnlineAggregator)=Dict(zip(keys(o.d),map(value,(values(o.d)))))
struct SymType{S}
end
SymType(s::Symbol)=SymType{s}()
(f::SymType{S})(x) where S=getproperty(x,S)
getbytypes(et,by) = Tuple{map(i->Base.nonmissingtype(Base.return_types(i,Tuple{et})[1]),by)...}
cubeeltype(t::GroupedOnlineAggregator{T}) where T=cubeeltype(T)
cubeeltype(t::Type{<:Dict{<:Any,T}}) where T = cubeeltype(T)
cubeeltype(t::Type{<:WeightedOnlineStat{T}}) where T = T
cubeeltype(t::Type{<:OnlineStat{Number}}) = Float64
cubeeltype(t::Type{<:WeightedCovMatrix{T}}) where T = T
cubeeltype(t::Type{<:Extrema{T}}) where T = T

function GroupedOnlineAggregator(O::OnlineStat,s::Symbol,by,w,iter)
    ost = typeof(O)
    et = eltype(iter)
    bytypes = Tuple{map(i->Base.nonmissingtype(Base.return_types(i,Tuple{et})[1]),by)...}
    d = Dict{bytypes,ost}()
    GroupedOnlineAggregator{typeof(d),s,typeof(by),typeof(w),ost}(d,w,by,O)
end

dicteltype(::Type{<:Dict{K,V}}) where {K,V} = V
dictktype(::Type{<:Dict{K,V}}) where {K,V} = K
actval(v) = v
actval(v::SentinelMissings.SentinelMissing) = v[]
function fitrow!(o::GroupedOnlineAggregator{T,S,BY,W},r) where {T,S,BY,W}
    v = getproperty(r,S)
    if !ismissing(v)
        w = o.w(r)
        if w===nothing
            bykey = map(i->actval(i(r)),o.by)
            if !any(ismissing,bykey)
                if haskey(o.d,bykey)
                    fit!(o.d[bykey],v)
                else
                    o.d[bykey] = copy(o.cloneobj)
                    fit!(o.d[bykey],v)
                end
            end
        else
           if !ismissing(w)
                bykey = map(i->actval(i(r)),o.by)
                if !any(ismissing,bykey)
                    if haskey(o.d,bykey)
                        fit!(o.d[bykey],v,w)
                    else
                        o.d[bykey] = copy(o.cloneobj)
                        fit!(o.d[bykey],v,w)
                    end
                end
            end
        end
    end
end
export TableAggregator, fittable, cubefittable
function TableAggregator(iter,O,fitsym;by=(),weight=nothing)
  !isa(by,Tuple) && (by=(by,))
  if !isempty(by)
    weight===nothing && (weight=(i->nothing))
    by = map(i->isa(i,Symbol) ? (SymType(i)) : i,by)
    GroupedOnlineAggregator(O,fitsym,by,weight,iter)
  else
    if weight===nothing
      OnlineAggregator(O,fitsym)
    else
      WeightOnlineAggregator(O,fitsym,weight)
    end
  end
end

function tooutaxis(::SymType{s},iter::CubeIterator{<:Any,<:Any,<:Any,<:Any,<:Any,S},k,ibc) where {s,S}
    ichosen = findfirst(i->i===s,fieldnames(S))
    if ichosen<=length(iter.inars)
        bycube = iter.dc.incubes[ichosen].cube
        if haskey(bycube.properties,"labels")
            idict=bycube.properties["labels"]
            axname=get(bycube.properties,"name","Label")
            outAxis=CategoricalAxis(axname,collect(String,values(idict)))
            convertdict=Dict(k=>i for (i,k) in enumerate(keys(idict)))
        else
            sort!(k)
            outAxis=CategoricalAxis("Label$(ibc)",k)
            convertdict=Dict(k=>i for (i,k) in enumerate(k))
        end
    else
       iax = findAxis(string(s),iter.dc.LoopAxes)
        outAxis=iter.dc.LoopAxes[iax]
        convertdict=Dict(k=>i for (i,k) in enumerate(outAxis.values))
    end
    outAxis,convertdict
end
function tooutaxis(f,iter::CubeIterator{<:Any,<:Any,<:Any,<:Any,<:Any,S},k,ibc) where S
    sort!(k)
    outAxis=CategoricalAxis("Category$(ibc)",k)
    convertdict=Dict(k=>i for (i,k) in enumerate(k))
    outAxis,convertdict
end

varsym(::WeightOnlineAggregator{<:Any,S}) where S = S
varsym(::OnlineAggregator{<:Any,S}) where S = S
varsym(::GroupedOnlineAggregator{<:Any,S}) where S = S
axt(::CategoricalAxis)=CategoricalAxis
axt(::RangeAxis)=RangeAxis
getStatType(::WeightOnlineAggregator{T}) where T = T
getStatType(::OnlineAggregator{T}) where T = T
getStatType(t::GroupedOnlineAggregator{T}) where T=getStatType(T)
getStatType(t::Type{<:Dict{<:Any,T}}) where T = T

getStatOutAxes(tab,agg)=getStatOutAxes(tab,agg,getStatType(agg))
getStatOutAxes(tab,agg,::Type{<:OnlineStat})=()
function getStatOutAxes(tab, agg, ::Type{<:Extrema})
    ( CategoricalAxis(:Extrema, ["min", "max"]), )
end
function getStatOutAxes(tab,agg,::Type{<:WeightedCovMatrix})
    varn = fieldnames(eltype(tab))
    s = varsym(agg)
    icube = findfirst(isequal(s),varn)
    ax = tab.dc.incubes[icube].axesSmall[1]
    oldname = YAXArrays.Cubes.Axes.axname(ax)
    coname = string("Co",oldname)
    v = ax.values
    axtype = axt(ax)
    a1 = axtype(oldname,copy(v))
    a2 = axtype(coname,copy(v))
    (a1,a2)
end
function getStatOutAxes(tab,agg,::Type{<:WeightedAdaptiveHist})
    nbin = getnbins(agg)
    a1 = RangeAxis("Bin",1:nbin)
    a2 = CategoricalAxis("Hist",["MidPoints","Frequency"])
    (a1,a2)
end
function getByAxes(iter,agg::GroupedOnlineAggregator)
    by = agg.by
    ntuple(length(by)) do ibc
      bc=agg.by[ibc]
      tooutaxis(bc,iter,unique(map(i->i[ibc],collect(keys(agg.d)))),ibc)
    end
end
getByAxes(iter,agg)=()

function tooutcube(
    agg,
    iter,
    post
  )
  outaxby = getByAxes(iter,agg)
  axby = map(i->i[1],outaxby)
  convdictall = map(i->i[2],outaxby)

  outaxstat = getStatOutAxes(iter,agg)
  outax = (outaxstat...,axby...)
  snew = map(length,outax)
  aout = fill!(zeros(Union{cubeeltype(agg),Missing},snew),missing)
  filloutar(aout,convdictall,agg,map(i->1:length(i),outaxstat),post)

  YAXArray(collect(outax),aout)
end
function filloutar(aout,convdictall,agg::GroupedOnlineAggregator,s,post)
    for (k,v) in agg.d
        i = CartesianIndices((s...,map((i,d)->d[convert(keytype(d),i)]:d[convert(keytype(d),i)],k,convdictall)...))
        aout[i.indices...].=post(v)
    end
end
function filloutar(aout,convdictall,agg,g,post)
    copyto!(aout,post(agg.o))
end
"""
    fittable(tab,o,fitsym;by=(),weight=nothing)

Loops through an iterable table `tab` and thereby fitting an OnlineStat `o` with the values
specified through `fitsym`. Optionally one can specify a field (or tuple) to group by.
Any groupby specifier can either be a symbol denoting the entry to group by or an anynymous
function calculating the group from a table row.

For example the following would caluclate a weighted mean over a cube weighted by grid cell
area and grouped by country and month:

````julia
fittable(iter,WeightedMean,:tair,weight=(i->abs(cosd(i.lat))),by=(i->month(i.time),:country))
````
"""
function fittable(tab,o,fitsym;by=(),weight=nothing,showprog=false)
  agg = TableAggregator(tab,o,fitsym,by=by,weight=weight)
  if showprog
    runfitrows_progress(agg,tab)
  else
    foreach(i->fitrow!(agg,i),tab)
  end
  agg
end
fittable(tab,o::Type{<:OnlineStat},fitsym;kwargs...)=fittable(tab,o(),fitsym;kwargs...)

@noinline function runfitrows_progress(agg,tab)
  p = Progress(length(tab)÷100,1)
  every = 0
  for row in tab
    fitrow!(agg,row)
    every += 1
    if every == 100
      next!(p)
      every=0
    end
  end
end

struct collectedValue{V,S,SY}
    value::V
    laststruct::S
end

function Base.getproperty(s::collectedValue{<:Any,<:Any,SY},sy::Symbol) where SY
    if sy == SY
        getfield(s,:value)
    else
        getproperty(getfield(s,:laststruct),sy)
    end
end

function collectval(row::Union{Tuple, Vector},::Val{SY}) where SY
    nvars = length(row)
    v = ntuple(i->getfield(row[i],SY),nvars) |> collect
    val = collectedValue{typeof(v),typeof(row[end]),SY}(v, row[end])
end

getpostfunction(s::OnlineStat)=getpostfunction(typeof(s))
getpostfunction(::Type{<:OnlineStat})=value
function getpostfunction(w::WeightedAdaptiveHist)
  nb = w.alg.b
  i->begin
    r = hcat(value(i)...)
    if size(r,1)<nb
      r = vcat(r,zeros(eltype(r),nb-size(r,1),2))
    end
    r
end
end
getnbins(f::GroupedOnlineAggregator)=f.cloneobj.alg.b
getnbins(f::TableAggregator)=f.o.alg.b

fitfun(o) = fitfun(typeof(o))
fitfun(::Type{<:Any}) = fittable

"""
    cubefittable(tab,o,fitsym;post=getpostfunction(o),kwargs...)

Executes [`fittable`](@ref) on the [`@CubeTable`](@ref) `tab` with the
(Weighted-)OnlineStat `o`, looping through the values specified by `fitsym`.
Finally, writes the results from the `TableAggregator` to an output data cube.
"""
function cubefittable(tab,o,fitsym;post=getpostfunction(o),kwargs...)
  agg=fitfun(o)(tab,o,fitsym;showprog=true,kwargs...)
  tooutcube(agg,tab,post)
end

function tupleeltypebyname(::Type{NamedTuple{names,tt}},s::Symbol) where {names, tt}
    i=findfirst(isequal(s),names)
    fieldtype(tt,i)
end


text:
#runs the basic panel form regressions
function olscontributionreg(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  push!(specs,
    xspec=Meta.parse("Lpcontributions"),
    xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
    xspec=Meta.parse("Lpprogramexpense"),
    xnames = [:intercept, :LagExp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
  xspec=Meta.parse("Lpcontributions"),
  xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
  xspec=Meta.parse("Lpprogramexpense"),
  xnames = [:intercept, :LagExp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
  xspec=Meta.parse("Lpcontributions"),
  xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  push!(specs,
  xspec=Meta.parse("Lpprogramexpense"),
  xnames = [:intercept, :LagExp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olscontributionreghist(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  push!(specs,
    xspec=Meta.parse("pcontributions"),
    xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn3yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpense"),
    xnames = [:intercept, :Exp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn3yr)

  push!(specs,
  xspec=Meta.parse("pcontributions"),
  xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
  xspec=Meta.parse("pprogramexpense"),
  xnames = [:intercept, :Exp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
  xspec=Meta.parse("pcontributions"),
  xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  push!(specs,
  xspec=Meta.parse("pprogramexpense"),
  xnames = [:intercept, :Exp],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olscontributionregcomb(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  push!(specs,
    xspec=Meta.parse("pcontributions"),
    xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn3yr)

  push!(specs,
    xspec=Meta.parse("Lpcontributions"),
    xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
  xspec=Meta.parse("pcontributions"),
  xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
  xspec=Meta.parse("Lpcontributions"),
  xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
  xspec=Meta.parse("pcontributions"),
  xnames = [:intercept, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  push!(specs,
  xspec=Meta.parse("Lpcontributions"),
  xnames = [:intercept, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olssensitivitypanelregPE(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
    xspec=Meta.parse("LpprogramexpenseXLpcont + Lpcontributions + Lpprogramexpense"),
    xnames = [:intercept, :LagContribExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
    xspec=Meta.parse("LpprogramexpenseXLpcont + Lpcontributions + Lpprogramexpense"),
    xnames = [:intercept, :LagContribExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  push!(specs,
    xspec=Meta.parse("LpprogramexpenseXLpcont + Lpcontributions + Lpprogramexpense"),
    xnames = [:intercept, :LagContribExp, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olssensitivitypanelregPEhist(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXpcont + pcontributions"),
    xnames = [:intercept, :ContribAvgExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn3yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpenseXpcont + pcontributions + pprogramexpense"),
    xnames = [:intercept, :ContribExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn3yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXpcont + pcontributions"),
    xnames = [:intercept, :ContribAvgExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpenseXpcont + pcontributions + pprogramexpense"),
    xnames = [:intercept, :ContribExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5003yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avgXpcont + pcontributions"),
    xnames = [:intercept, :ContribAvgExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  push!(specs,
    xspec=Meta.parse("pprogramexpenseXpcont + pcontributions + pprogramexpense"),
    xnames = [:intercept, :ContribExp, :Contrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff33yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olssensitivitypanelregPS(df::AbstractDataFrame)
  specs::FMSpecs = FMSpecs(6)

  #set up the results
  #NOTE: delete the l before public support on next pass
  push!(specs,
    xspec=Meta.parse("publicsupport_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgSup, :LagContrib],
    withinspec = :ein,
    clusterspec = :ein,
    yspec=:lreturn)

  push!(specs,
    xspec=Meta.parse("LpublicsupportXLpcont + Lpcontributions + Lpublicsupport"),
    xnames = [:intercept, :LagContribSup, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturn)

  push!(specs,
    xspec=Meta.parse("publicsupport_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgSup, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
    xspec=Meta.parse("LpublicsupportXLpcont + Lpcontributions + Lpublicsupport"),
    xnames = [:intercept, :LagContribSup, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmsp5001yr)

  push!(specs,
    xspec=Meta.parse("publicsupport_avgXLpcont + Lpcontributions"),
    xnames = [:intercept, :LagContribAvgSup, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  push!(specs,
    xspec=Meta.parse("LpublicsupportXLpcont + Lpcontributions + Lpublicsupport"),
    xnames = [:intercept, :LagContribSup, :LagContrib],
    clusterspec = :ein,
    withinspec = :ein,
    yspec=:lreturnnetbmff31yr)

  #run the regressions
  computeFMLMresults!(df, specs)

  return specs
end

#runs the basic panel form regressions
function olssensitivity2pass(df::DataFrame;
  rhs2ndpassfields_avg::Vector{Symbol} = [:pprogramexpense_avg],#, :publicsupport_avg=#],
  rhs2ndpassfields::Vector{Symbol} = [:Lpprogramexpense#=, :Lpublicsupport_avg=#],
  winsorizeprop::Float64 = WINSORIZE_PROP)

  specs::FMSpecs = FMSpecs(6)

  local βdf::DataFrame
  local βdfcross::DataFrame
  local focalfields::Vector{Symbol}
  #run the first pass regressions to get the betas
  #println("q expense complete: ",
  #  sum(completecases(df[[:lreturn, :Llcontributions, :lprogramexpense_avg]])))

  regressbynonprofit!(df,
    YField = :lreturn,
    XFields = [:Lpcontributions],
    βNames = [:b0temp1, :b1lRetLCon],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  regressbynonprofit!(df,
    YField = :lreturnnetbmsp5001yr,
    XFields = [:Lpcontributions],
    βNames = [:b0temp2, :b1lRetNetSP500LCon],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  regressbynonprofit!(df,
    YField = :lreturnnetbmff31yr,
    XFields = [:Lpcontributions],
    βNames = [:b0temp3, :b1lRetNetFF3LCon],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  deletecols!(df, [:b0temp1, :b0temp2, :b0temp3])

  βfields::Vector{Symbol} = [:b1lRetLCon, :b1lRetNetSP500LCon,
    :b1lRetNetFF3LCon]

  (f->winsorize!(df[f], prop=winsorizeprop)).(βfields)
  focalfields = [:ein; βfields]

  #now fold in the 2nd pass regression terms through repeated inner joins and uniqueness filters
  βdf = unique(df[[focalfields; :categoryfiltered]], focalfields)
  for f ∈ rhs2ndpassfields_avg
    subdf::DataFrame = view(df, completecases(df[:,[:ein,f]]), [:ein, f])
    βdf = join(βdf, unique(subdf), on=:ein, kind=:left)
  end

  plotβtab!(βdf, tabfield = :b1lRetLCon)
  #CSV.write("$WORKING_PATH\\betatest.csv", βdf)

  #make sure the joins went as planned
  @assert size(unique(βdf),1) == size(βdf,1)
  focalfields = [:ein; :fisyr; βfields; :categoryfiltered]

  βdfcross = df[focalfields]

  for f ∈ rhs2ndpassfields
    subdf::DataFrame = view(df, completecases(df[:,[:ein,:fisyr,f]]), [:ein,:fisyr,f])
    βdfcross = join(βdfcross, unique(subdf), on=[:ein, :fisyr], kind=:left)
  end

  #make sure all of the betas within each entity are unique
  @assert size(unique(βdfcross),1) == size(βdfcross,1)

  #make the fisyrs into symbols
  rename!(βdfcross, [:fisyr=>:fisyrold])
  βdfcross[:fisyr] = (y::Int->Symbol("y", y)).(βdfcross[:fisyrold])
  deletecols!(βdfcross, :fisyrold)

  #Regress the betas on the focal variables
  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    withinspec = :categoryfiltered,
    yspec=:b1lRetLCon)

  push!(specs,
    xspec=Meta.parse("Lpprogramexpense"),
    xnames = [:intercept, :LagExp],
    withinspec = :categoryfiltered,
    yspec=:b1lRetLCon)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    withinspec = :categoryfiltered,
    yspec=:b1lRetNetSP500LCon)

  push!(specs,
    xspec=Meta.parse("Lpprogramexpense"),
    xnames = [:intercept, :LagExp],
    withinspec = :categoryfiltered,
    yspec=:b1lRetNetSP500LCon)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    withinspec = :categoryfiltered,
    yspec=:b1lRetNetFF3LCon)

  push!(specs,
    xspec=Meta.parse("Lpprogramexpense"),
    xnames = [:intercept, :LagExp],
    #withinspec = :fisyr,
    withinspec = :categoryfiltered,
    yspec=:b1lRetNetFF3LCon)

  println(describe(βdf))
  println(describe(βdfcross))

  #run the regressions
  βdfs::Vector{DataFrame} = [βdf, βdfcross, βdf, βdfcross, βdf, βdfcross]
  computeFMLMresults!(βdfs, specs)

  return specs
end

#runs the basic panel form regressions
function olssensitivity2passhist(df::DataFrame;
  rhs2ndpassfields_avg::Vector{Symbol} = [:pprogramexpense_avg],#, :publicsupport_avg=#],
  rhs2ndpassfields::Vector{Symbol} = [:pprogramexpense#=, :Lpublicsupport_avg=#],
  winsorizeprop::Float64 = WINSORIZE_PROP)

  specs::FMSpecs = FMSpecs(6)
  df = prefilterforolssensitivity(deepcopy(df), returnfield = :lreturn3yr)

  local βdf::DataFrame
  local βdfcross::DataFrame
  local focalfields::Vector{Symbol}
  #run the first pass regressions to get the betas
  #println("q expense complete: ",
  #  sum(completecases(df[[:lreturn, :Llcontributions, :lprogramexpense_avg]])))

  regressbynonprofit!(df,
    YField = :lreturn3yr,
    XFields = [:pcontributions],
    βNames = [:b0temp1, :b1lRet3yrCon],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  regressbynonprofit!(df,
    YField = :lreturnnetbmsp5003yr,
    XFields = [:pcontributions],
    βNames = [:b0temp2, :b1lRet3yrNetSP500Con],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  regressbynonprofit!(df,
    YField = :lreturnnetbmff33yr,
    XFields = [:pcontributions],
    βNames = [:b0temp3, :b1lRet3yrNetFF3Con],
    groupbyfield = :ein,
    predictwithresults = false,
    checkunique=false)

  deletecols!(df, [:b0temp1, :b0temp2, :b0temp3])

  βfields::Vector{Symbol} = [:b1lRet3yrCon, :b1lRet3yrNetSP500Con,
    :b1lRet3yrNetFF3Con]

  (f->winsorize!(df[f], prop=winsorizeprop)).(βfields)
  focalfields = [:ein; βfields]

  #now fold in the 2nd pass regression terms through repeated inner joins and uniqueness filters
  βdf = unique(df[focalfields])
  for f ∈ rhs2ndpassfields_avg
    subdf::DataFrame = view(df, completecases(df[:,[:ein,f]]), [:ein, f])
    βdf = join(βdf, unique(subdf), on=:ein, kind=:left)
  end

  #CSV.write("$WORKING_PATH\\betatest.csv", βdf)

  #make sure the joins went as planned
  @assert size(unique(βdf),1) == size(βdf,1)
  focalfields = [:ein; :fisyr; βfields]

  βdfcross = df[focalfields]

  for f ∈ rhs2ndpassfields
    subdf::DataFrame = view(df, completecases(df[:,[:ein,:fisyr,f]]), [:ein,:fisyr,f])
    βdfcross = join(βdfcross, unique(subdf), on=[:ein, :fisyr], kind=:left)
  end

  #make sure all of the betas within each entity are unique
  @assert size(unique(βdfcross),1) == size(βdfcross,1)

  #make the fisyrs into symbols
  rename!(βdfcross, [:fisyr=>:fisyrold])
  βdfcross[:fisyr] = (y::Int->Symbol("y", y)).(βdfcross[:fisyrold])
  deletecols!(βdfcross, :fisyrold)

  plotβtab!(βdf)

  #Regress the betas on the focal variables
  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    yspec=:b1lRet3yrCon)

  push!(specs,
    xspec=Meta.parse("pprogramexpense"),
    xnames = [:intercept, :Exp],
    withinspec = :fisyr,
    yspec=:b1lRet3yrCon)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    yspec=:b1lRet3yrNetSP500Con)

  push!(specs,
    xspec=Meta.parse("pprogramexpense"),
    xnames = [:intercept, :Exp],
    withinspec = :fisyr,
    yspec=:b1lRet3yrNetSP500Con)

  push!(specs,
    xspec=Meta.parse("pprogramexpense_avg"),
    xnames = [:intercept, :AvgExp],
    yspec=:b1lRet3yrNetFF3Con)

  push!(specs,
    xspec=Meta.parse("pprogramexpense"),
    xnames = [:intercept, :Exp],
    withinspec = :fisyr,
    yspec=:b1lRet3yrNetFF3Con)

  println(describe(βdf))
  println(describe(βdfcross))

  #run the regressions
  βdfs::Vector{DataFrame} = [βdf, βdfcross, βdf, βdfcross, βdf, βdfcross]
  computeFMLMresults!(βdfs, specs)

  return specs
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using SixelGraphics

# inspiration and some code from
# http://mathemartician.blogspot.nl/2012/07/julia-set-in-julia.html

function julia(z, maxiter::Int64)
    c=-0.75+0.11im
    for n = 1:maxiter
        if abs(z) > 2
            return n-1
        end
        z = z^2 + c
    end
    return maxiter
end

w = 384
h = 288

s = Screen(w,h)

clear(s.b,63)

# for every pixel
for y=1:h, x=1:w
    # translate numbers [1:w, 1:h] -> -2.5:2.5 + -1:1 im
    c = complex((x-w/2)/(h/2.5), (y-h/2)/(h/2))
    # call the julia function and use the distance as the color
    putpixel(s.b, div(julia(c,255),16), x-1, y-1)
end

sixels(s)