text:
# include("core.jl")

"""
    Aggregate{T,C1,C2}(molecules, coupling)
    Aggregate{T,C1,C2}(molecules)

Mutable stuct which purpose is to store important information about the whole system.

# Arguments
* `molecules`: Vector of molecules ([`Molecule`](@ref)).
* `coupling`: Matrix of couplings ``J_{nm}`` between molecules.
"""
mutable struct Aggregate{T<:Integer,C1<:ComputableType,C2<:ComputableType}
    molecules::Vector{Molecule{T,C1,C2}}
    coupling::Matrix{C1}
    function Aggregate{T,C1,C2}(
        molecules::Vector{Molecule{T,C1,C2}},
        coupling::Matrix{C1},
    ) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
        new(molecules, coupling)
    end
end

Aggregate(molecules::Vector{Molecule{T,C1,C2}}, coupling::Matrix{C1}) where {T,C1,C2} =
    Aggregate{T,C1,C2}(molecules, coupling)
Aggregate(molecules::Vector{Molecule{T,C1,C2}}) where {T,C1,C2} =
    Aggregate{T,C1,C2}(molecules, zeros(C1, (length(molecules) + 1, length(molecules) + 1)))

"""
    getNvib(agg)

Get maximum number of vibrational states of each molecule in the [`Aggregate`](@ref).

"""
function getNvib(
    agg::Aggregate{T,C1,C2},
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    NvibMols = Array{Array{T,1},1}(undef, 0)
    for mol in agg.molecules
        push!(NvibMols, fill(mol.Nvib, length(mol.modes)))
    end
    return NvibMols
end

"""
    getShifts(agg)

Get shifts for every mode on each molecule in the [`Aggregate`](@ref).

"""
function getShifts(
    agg::Aggregate{T,C1,C2},
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    shifts = Array{Array{C1,1},1}(undef, 0)
    for mol in agg.molecules
        push!(shifts, getMolShifts(mol))
    end
    return shifts
end

"""
    getFrequencies(agg)

Get frequencies for every mode on each molecule in the [`Aggregate`](@ref).

"""
function getFrequencies(
    agg::Aggregate{T,C1,C2},
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    frequencies = Array{Array{C1,1},1}(undef, 0)
    for mol in agg.molecules
        push!(frequencies, getMolFrequencies(mol))
    end
    return frequencies
end

"""
    vibrationalIndices(agg)

Get all vibrational indices of the [`Aggregate`](@ref).

"""
function vibrationalIndices(
    agg::Aggregate{T,C1,C2},
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    NvibMols = getNvib(agg)
    molLen = length(agg.molecules)
    NvibIndMols = Array{Array{Array{T,1},1},1}(undef, 0)
    for mol_i = 1:molLen
        push!(NvibIndMols, vibrationalIndices(NvibMols[mol_i]))
    end
    molNvib = map((vibInds) -> length(vibInds), NvibIndMols)
    molInds = vibrationalIndices(molNvib)
    aggInds = Array{Array{Array{T,1},1},1}(undef, 0)
    for molInd in molInds
        aggInd = Array{Array{T,1},1}(undef, 0)
        for mol_i = 1:molLen
            push!(aggInd, NvibIndMols[mol_i][molInd[mol_i]])
        end
        push!(aggInds, aggInd)
    end
    return aggInds
end

"""
    vibrationalIndices(agg)

Get all electronic indices of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function electronicIndices(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    vibInds = Array{Array{T,1},1}(undef, 0)
    molCount = length(agg.molecules)
    currentInds = fill(1, (molCount))
    push!(vibInds, copy(currentInds))
    for indPos = 1:molCount
        currentInds[indPos] += 1
        push!(vibInds, copy(currentInds))
        currentInds[indPos] -= 1
    end
    return vibInds
end

"""
    getIndices(agg)

Get all indices (electronic and vibrational) of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getIndices(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    vibInds = vibrationalIndices(agg)
    vibIndsLen = length(vibInds)
    elInds = electronicIndices(agg)
    elIndsLen = length(elInds)
    indices = Array{Array{Array{T,1} where T,1},1}(undef, 0)
    for el_i = 1:elIndsLen
        for vib_i = 1:vibIndsLen
            push!(indices, [elInds[el_i], vibInds[vib_i]])
        end
    end
    return indices
end

"""
    getVibIndices(agg)

Get pointers (integers) to the indices of the [`Aggregate`](@ref) separated by
electronic states (e.g. [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getVibIndices(agg, aggIndices)
    aggIndLen = length(aggIndices)
    vibindices = Array{Array{Int,1},1}(undef, 0)
    elLen = length(agg.molecules) + 1
    for el_i = 1:elLen
        push!(vibindices, Array{Int,1}(undef, 0))
    end

    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)
        push!(vibindices[elOrder1], I)
    end
    return vibindices
end

"""
    getFranckCondonFactors(agg, aggIndices)
    getFranckCondonFactors(agg)

Get Frack-Condon factors of the [`Aggregate`](@ref) in a form of matrix.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
"""
function getFranckCondonFactors(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    FC = zeros(C2, (aggIndLen, aggIndLen))
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            if elind1 == elind2
                if vibind1 == vibind2
                    FC[I, J] = 1.0
                end
            else
                fc = 1.0::C2
                for mi = 1:molLen
                    mol = agg.molecules[mi]
                    fc *=
                        getMolStateFC(mol, elind1[mi], vibind1[mi], elind2[mi], vibind2[mi])
                end
                FC[I, J] = fc
            end
        end
    end
    return FC
end

getFranckCondonFactors(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} =
    getFranckCondonFactors(agg::Aggregate{T,C1,C2}, nothing)

"""
    getAggStateEnergy(agg, aggElState, aggVibState)

Get Hamiltonian of the [`Aggregate`](@ref) in a form of DenseOperator.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggElState`: Aggregate electric state (e.g. [1, 1, 2]).
* `aggVibState`: Aggregate vibrational state (e.g. [[9, 1], [2, 5, 3], [10]]).
"""
function getAggStateEnergy(
    agg::Aggregate{T,C1,C2},
    aggElState::Vector{U},
    aggVibState::Vector{Vector{U}},
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType,U<:Integer}
    energy = 0.0::C1
    for mol_i = 1:length(agg.molecules)
        molElState = aggElState[mol_i]
        molVibState = aggVibState[mol_i]
        energy += getMolStateEnergy(agg.molecules[mol_i], molElState, molVibState)
    end
    return energy
end

function elIndOrder(elInd::Vector{T}) where {T<:Integer}
    len = length(elInd)
    ind = 1::T
    for i = 1:len
        ind += (elInd[i] - 1) * i
    end
    return ind
end

"""
    getAggHamSystemSmall(agg)

Get Hamiltonian of the system, ``H_S`` only for ``\\mathcal{H}_S``.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getAggHamSystemSmall(
    agg::Aggregate{T,C1,C2};
    groundEnergy::Bool = true,
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    molLen = length(agg.molecules)
    Ham_sys = zeros(C2, (molLen + 1, molLen + 1))

    agg_shifts = getShifts(agg)
    agg_frequencies = getShifts(agg)
    reorganisation_energies = []
    for mol_i = 1:length(agg.molecules)
        mol = agg.molecules[mol_i]
        reorganisation_energy = 0.
        for mode_i = 1:length(mol.modes)
            reorganisation_energy += agg_frequencies[mol_i][mode_i] * agg_shifts[mol_i][mode_i]^2 / 2.
        end
        push!(reorganisation_energies, reorganisation_energy)
    end

    elInds = electronicIndices(agg)
    E_agg = zeros(C1, (2, molLen))
    E_agg[1, :] = map(mol -> mol.E[1], agg.molecules)
    E_agg[2, :] = map(mol -> mol.E[2], agg.molecules)
    for elInd in elInds
        ind = OpenQuantumSystems.elIndOrder(elInd)
        E_state = 0
        for mol_i = 1:molLen
            E_state += E_agg[elInd[mol_i], mol_i]
        end
        # add reorganisation energy
        if ind > 1
            E_state += reorganisation_energies[ind-1]
        end
        Ham_sys[ind, ind] = E_state
    end

    Ham_sys[:, :] += agg.coupling[:, :]
    if !groundEnergy
        E0 = Ham_sys[1, 1]
        for i = 1:size(Ham_sys, 1)
            Ham_sys[i, i] -= E0
        end
    end
    b = GenericBasis([size(Ham_sys, 1)])
    return DenseOperator(b, b, Ham_sys)
end

"""
    getAggHamSystemBig(agg)

Get Hamiltonian of the system, ``H_S`` for ``\\mathcal{H}_S \\otimes \\mathcal{H}_B``.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getAggHamSystemBig(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any,
    franckCondonFactors::Any;
    groundEnergy::Bool = true,
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    if franckCondonFactors === nothing
        franckCondonFactors = getFranckCondonFactors(agg, aggIndices)
    end
    Ham_sys = getAggHamSystemSmall(agg; groundEnergy = groundEnergy)
    Ham = zeros(C2, (aggIndLen, aggIndLen))
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        elOrder1 = elIndOrder(elind1)
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            elOrder2 = elIndOrder(elind2)

            Ham[I, J] = Ham_sys.data[elOrder1, elOrder2] * franckCondonFactors[I, J]
        end
    end
    if !groundEnergy
        E0 = Ham[1, 1]
        for I = 1:aggIndLen
            Ham[I, I] -= E0
        end
    end
    b = GenericBasis([aggIndLen])
    return DenseOperator(b, b, Ham)
end

"""
    getAggHamBathSmall(agg)

Get Hamiltonian of the bath, ``H_B`` only for ``\\mathcal{H}_B``.

"""
function getAggHamBathSmall(
    agg::Aggregate{T,C1,C2};
    groundEnergy::Bool = true
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    molLen = length(agg.molecules)
    vibInds = vibrationalIndices(agg)
    vibIndsLen = length(vibInds)
    Ham_bath = zeros(C2, (vibIndsLen, vibIndsLen))

    agg2 = deepcopy(agg)
    for mol_i in molLen
        agg2.molecules[mol_i].E[1] = 0.0
    end
    elind = fill(1, (molLen + 1))
    Ham_sys = getAggHamSystemSmall(agg)
    for I = 1:vibIndsLen
        vibind = vibInds[I]
        Ham_bath[I, I] = getAggStateEnergy(agg, elind, vibind) - Ham_sys.data[1, 1]
    end
    b = GenericBasis([vibIndsLen])
    if !groundEnergy
        E0 = Ham_bath[1, 1]
        for i = 1:size(Ham_bath, 1)
            Ham_bath[i, i] -= E0
        end
    end
    return DenseOperator(b, b, Ham_bath)
end

"""
    getAggHamBathBig(agg)

Get Hamiltonian of the bath, ``H_B`` only for ``\\mathcal{H}_S \\otimes \\mathcal{H}_B``.

"""
function getAggHamBathBig(
    agg::Aggregate{T,C1,C2};
    groundEnergy::Bool = true
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    aggIndices = getIndices(agg)
    aggIndLen = length(aggIndices)
    b = GenericBasis([aggIndLen])

    Ham_bath = getAggHamBathSmall(agg)
    Ham_sys = getAggHamSystemSmall(agg)
    b_sys = GenericBasis([size(Ham_sys, 1)])
    one_sys = OneDenseOperator(b_sys)
    # TODO: optimise types in future
    one_sys = DenseOperator(b_sys, b_sys, real(one_sys.data))
    Ham_bath = tensor(Ham_bath, one_sys)
    if !groundEnergy
        E0 = Ham_bath.data[1, 1]
        for I = 1:size(Ham_bath.data, 1)
            Ham_bath.data[I, I] -= E0
        end
    end
    return DenseOperator(b, b, Ham_bath.data)
end

"""
    getAggHamSystemBath(agg, aggIndices,
    \tfranckCondonFactors; groundEnergy = false)
    getAggHamSystemBath(agg, aggIndices; groundEnergy = false)
    getAggHamSystemBath(agg; groundEnergy = false)

Get system and bath Hamiltonian of the [`Aggregate`](@ref) in a form of DenseOperator.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamSystemBath(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any,
    franckCondonFactors::Any;
    groundEnergy::Bool = true,
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    if franckCondonFactors === nothing
        franckCondonFactors = getFranckCondonFactors(agg, aggIndices)
    end

    Ham_B = getAggHamBathBig(agg; groundEnergy=true)
    Ham_S = getAggHamSystemBig(agg, aggIndices, franckCondonFactors; groundEnergy=true)
    Ham_0 = Ham_B + Ham_S
    if !groundEnergy
        E0 = Ham_0.data[1, 1]
        for I = 1:aggIndLen
            Ham_0.data[I, I] -= E0
        end
    end
    return Ham_0
end

getAggHamSystemBath(
    agg::Aggregate{T,C1,C2};
    groundEnergy::Bool = true,
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamSystemBath(
    agg::Aggregate{T,C1,C2},
    nothing,
    nothing;
    groundEnergy = groundEnergy,
)

getAggHamSystemBath(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any;
    groundEnergy::Bool = true,
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamSystemBath(
    agg::Aggregate{T,C1,C2},
    aggIndices,
    nothing;
    groundEnergy = groundEnergy,
)

"""
    getAggHamInteraction(agg, aggIndices, franckCondonFactors)
    getAggHamInteraction(agg, aggIndices)
    getAggHamInteraction(agg)

Get interation Hamiltonian of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamInteraction(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any,
    franckCondonFactors::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    if franckCondonFactors === nothing
        franckCondonFactors = getFranckCondonFactors(agg, aggIndices)
    end
    Ham_I = zeros(C2, (aggIndLen, aggIndLen))
    agg_shifts = getShifts(agg)
    agg_frequencies = getShifts(agg)
    agg_coeffs = deepcopy(agg_frequencies)
    for mol_i = 1:length(agg.molecules)
        mol = agg.molecules[mol_i]
        for mode_i = 1:length(mol.modes)
            agg_coeffs[mol_i][mode_i] = agg_shifts[mol_i][mode_i] * sqrt(agg_frequencies[mol_i][mode_i] / 2)
        end
    end

    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        elOrder1 = elIndOrder(elind1)
        if elOrder1 == 1
            continue
        end
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            elOrder2 = elIndOrder(elind2)
            if elOrder2 == 1 || elOrder1 != elOrder2
                continue
            end
            diff_num = 0; mol_j = 0; mode_j = 0
            for mol_i = 1:length(agg.molecules)
                mol = agg.molecules[mol_i]
                for mode_i = 1:length(mol.modes)
                    diff_vib = abs(vibind1[mol_i][mode_i] - vibind2[mol_i][mode_i])
                    diff_num += diff_vib
                    if diff_vib == 1 && diff_num == 1
                        mol_j = mol_i; mode_j = mode_i;
                    end
                end
            end

            if diff_num != 1 || mol_j != elOrder1-1
                continue
            end
            vib_n = vibind1[mol_j][mode_j]
            vib_m = vibind2[mol_j][mode_j]
            coeff = agg_coeffs[mol_j][mode_j]
            Ham_I[I, J] = - coeff * sqrt(min(vib_n, vib_m))
        end
    end
    b = GenericBasis([aggIndLen])
    return DenseOperator(b, b, Ham_I)
end

getAggHamInteraction(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamInteraction(
    agg::Aggregate{T,C1,C2},
    nothing,
    nothing
)

getAggHamInteraction(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamInteraction(
    agg::Aggregate{T,C1,C2},
    aggIndices,
    nothing
)

"""
    getAggHamiltonian(agg, aggIndices, franckCondonFactors;
    \t, groundEnergy = true)
    getAggHamiltonian(agg, aggIndices; groundEnergy = true)
    getAggHamiltonian(agg; groundEnergy = true)

Get Hamiltonian of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamiltonian(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any,
    franckCondonFactors::Any;
    groundEnergy::Bool = true
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    if franckCondonFactors === nothing
        franckCondonFactors = getFranckCondonFactors(agg, aggIndices)
    end
    Ham_I = getAggHamInteraction(agg, aggIndices, franckCondonFactors)
    Ham_0 = getAggHamSystemBath(agg, aggIndices, franckCondonFactors; groundEnergy = groundEnergy)
    return Ham_0 + Ham_I
end

getAggHamiltonian(
    agg::Aggregate{T,C1,C2};
    groundEnergy::Bool = true
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonian(
    agg::Aggregate{T,C1,C2},
    nothing,
    nothing;
    groundEnergy = groundEnergy,
)

getAggHamiltonian(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any;
    groundEnergy::Bool = true
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonian(
    agg::Aggregate{T,C1,C2},
    aggIndices,
    nothing;
    groundEnergy = groundEnergy,
)

### Sparse versions

"""
    getFranckCondonFactorsSparse(agg, aggIndices)
    getFranckCondonFactorsSparse(agg)

Sparse version of [`getFranckCondonFactors`](@ref).

"""
function getFranckCondonFactorsSparse(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    FC = spzeros(C2, aggIndLen, aggIndLen)
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            if elind1 == elind2
                if vibind1 == vibind2
                    FC[I, J] = 1.0
                end
            else
                fc = 1.0::C2
                for mi = 1:molLen
                    mol = agg.molecules[mi]
                    fc *=
                        getMolStateFC(mol, elind1[mi], vibind1[mi], elind2[mi], vibind2[mi])
                end
                if fc != 0.0
                    FC[I, J] = fc
                end
            end
        end
    end
    return FC
end

getFranckCondonFactorsSparse(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getFranckCondonFactorsSparse(
    agg::Aggregate{T,C1,C2},
    nothing
)

"""
    getAggHamiltonianSparse(agg, aggIndices)
    getAggHamiltonianSparse(agg)

Sparse version of [`getAggHamiltonian`](@ref).

"""
function getAggHamiltonianSparse(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any,
    franckCondonFactors::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    if franckCondonFactors === nothing
        franckCondonFactors = getFranckCondonFactors(agg, aggIndices)
    end
    Ham = spzeros(C2, aggIndLen, aggIndLen)
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        elOrder1 = elIndOrder(elind1)
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            elOrder2 = elIndOrder(elind2)
            if I == J
                val = getAggStateEnergy(agg, elind1, vibind1)
                if val != 0.0
                    Ham[I, J] = val
                end
            else
                if elind1 != elind2
                    val = agg.coupling[elOrder1, elOrder2] * franckCondonFactors[I, J]
                    if val != 0.0
                        Ham[I, J] = val
                    end
                end
            end
        end
    end
    E0 = Ham[1, 1]
    for I = 1:aggIndLen
        Ham[I, I] -= E0
    end
    b = GenericBasis([aggIndLen])
    return SparseOperator(b, b, Ham)
end

getAggHamiltonianSparse(
    agg::Aggregate{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonianSparse(
    agg::Aggregate{T,C1,C2},
    nothing,
    nothing
)

getAggHamiltonianSparse(
    agg::Aggregate{T,C1,C2},
    aggIndices::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonianSparse(
    agg::Aggregate{T,C1,C2},
    aggIndices,
    nothing
)

"""
    setupAggregate(agg; groundEnergy=true, verbose=false)

Generate all basic data from the [`Aggregate`](@ref). Returns
`aggInds, vibindices, aggIndLen, basis, FCFact, FCProd, Ham, Ham_0, Ham_I`.

"""
function setupAggregate(agg; groundEnergy = true, verbose = false)
    if verbose
        println("aggInds")
        @time aggInds = getIndices(agg)
        println("vibindices")
        @time vibindices = getVibIndices(agg, aggInds)
        aggIndLen = length(aggInds)
        println("aggIndLen: ", aggIndLen)
        println("basis")
        @time basis = GenericBasis([aggIndLen])
        println("FCFact")
        @time FCFact = getFranckCondonFactors(agg, aggInds)
        println("FCProd")
        @time FCProd =
            getFCProd(agg, FCFact, aggInds, vibindices)
        println("Ham")
        @time Ham = getAggHamiltonian(
            agg,
            aggInds,
            FCFact;
            groundEnergy = groundEnergy,
        )
        println("Ham_0")
        @time Ham_0 = getAggHamSysBath(
            agg,
            aggInds;
            groundEnergy = groundEnergy,
        )
        println("Ham_I")
        @time Ham_I = Ham - Ham_0
    else
        aggInds = getIndices(agg)
        vibindices = getVibIndices(agg, aggInds)
        aggIndLen = length(aggInds)
        basis = GenericBasis([aggIndLen])
        FCFact = getFranckCondonFactors(agg, aggInds)
        FCProd = getFCProd(agg, FCFact, aggInds, vibindices)
        Ham = getAggHamiltonian(
            agg,
            aggInds,
            FCFact;
            groundEnergy = groundEnergy,
        )
        Ham_0 = getAggHamSysBath(
            agg,
            aggInds;
            groundEnergy = groundEnergy,
        )
        Ham_I = Ham - Ham_0
    end
    return (aggInds, vibindices, aggIndLen, basis, FCFact, FCProd, Ham, Ham_0, Ham_I)
end


text:
struct SampleDistribution{T,K} <: Distribution{Multivariate, Discrete}
    dist::Categorical{T}
    map::Dict{K,Int}
    inverse_map::Dict{Int,K}
end

function SampleDistribution(p::AbstractVector{T}, dict::Dict{K,<:Integer}) where {T<:Real,K}
    SampleDistribution{T,K}(Categorical{T}(p), dict, inversedict(dict))
end

# this is a dict with probabilities as values
function SampleDistribution(dict::Dict{K,T}) where {K,T<:AbstractFloat}
    ks = collect(keys(dict))
    p = getindex.(dict, ks)
    SampleDistribution(p, Dict(ks[i]=>i for i ∈ 1:length(ks)))
end

# this is a dict with counts as values
function SampleDistribution(dict::Dict{K,T}) where {K,T<:Integer}
    total = sum(values(dict))
    SampleDistribution(Dict{K,Float64}(k=>Float64(v/total) for (k,v) ∈ dict))
end

export SampleDistribution

# this will do counts for you
SampleDistribution(v::AbstractVector) = SampleDistribution(proportionmap(v))

supportkeys(d::SampleDistribution) = keys(d.map)
support(d::SampleDistribution) = collect(supportkeys(d))

getsample(d::SampleDistribution, i::Integer) = d.inverse_map[i]

pdfbyindex(d::SampleDistribution, i::Integer) = pdf(d.dist, i)

"""
    expectation(d::SampleDistribution, f::Function)

Compute the expectation value of `f` over distribution `d`.  `f` should take a single argument
which is in the support of `d`.  Also supports `do` syntax.
"""
function expectation(d::SampleDistribution, f::Function)
    sum(pdfbyindex(d, i)*f(getsample(d, i)) for i ∈ 1:ncategories(d))
end
expectation(f::Function, d::SampleDistribution) = expectation(d, f)
export expectation

ext: jl


text:
# Sites Space
export SiteSpace

"""
    SiteSpace{NFlips, S} <: AbstractSpace{Sites, S}

A space of sites contains all possible configurations of the sites.
`NFlips` indicated how many site will flip according to certain rules
when `shake!` is called.

    SiteSpace(label, shape; nflips)

constructs a `SiteSpace` with label in type `SiteLabel`, and
`shape` with an optional keyword `nflips` (default to be 1).

```julia-repl
julia> SiteSpace(Bit, 2, 2)
QMTK.SiteSpace{1,QMTK.Randomized}(Int8[1 1; 0 0])
```
"""
mutable struct SiteSpace{NFlips, S} <: AbstractSpace{Sites, S}
    data::Sites
end

SiteSpace(::Type{L}, shape::Tuple; nflips=1) where {L <: SiteLabel} =
    SiteSpace{nflips, Randomized}(rand(L, shape))
SiteSpace(::Type{L}, shape...; nflips=1) where {L <: SiteLabel} =
    SiteSpace{nflips, Randomized}(rand(L, shape))
SiteSpace(::Type{S}, space::SiteSpace{N}) where {N, S} = SiteSpace{N, S}(space.data)

# Properties
data(space::SiteSpace) = space.data

# Basics
function reset!(space::SiteSpace{N, Randomized}) where N
    reset!(space.data)
    return SiteSpace(UnRandomized, space)
end

copy(space::SiteSpace{N, S}) where {N, S} = SiteSpace(space.data)
acquire(space::SiteSpace) = copy(space.data)

# Randomized methods
shake!(space::SiteSpace{1, Randomized}) = randflip!(space.data)

function shake!(space::SiteSpace{N, Randomized}; MAX_FLIP=1000_000) where N
    tape = Int[]
    for i = 1:MAX_FLIP
        offset = rand(1:length(space.data))
        if offset in tape
            continue
        else
            flip!(space.data, offset)
            push!(tape, offset)
        end

        if length(tape) == N
            break
        end
    end

    return space.data
end

function randomize!(space::SiteSpace{N, UnRandomized}) where N
    rand!(space.data)
    return SiteSpace(Randomized, space)
end

struct SiteTraverser{L <: SiteLabel, T, N} <: AbstractSpaceTraverser
    data::Sites{L, T, N}
end

SiteTraverser(space::SiteSpace) = SiteTraverser(reset!(space.data))

traverse(space::SiteSpace) = SiteTraverser(space)

import Base: start, next, done, length, eltype

eltype(::Type{SiteTraverser{L, T, N}}) where {L, T, N} = Sites{L, T, N}
length(itr::SiteTraverser) = 1 << length(itr.data)

start(itr::SiteTraverser{L}) where L = 1

function next(itr::SiteTraverser, state)
    r = copy(itr.data)
    itr.data << 1
    return r, state + 1
end

done(itr::SiteTraverser, state) = state > length(itr)


text:
include("6.7.6 (c+d x)^m hyper(a+b x)^n hyper(a+b x)^p.jl")
include("6.7.7 F^(c (a+b x)) hyper(d+e x)^n.jl")
include("6.7.8 u hyper(a+b log(c x^n))^p.jl")
include("6.7.9 Active hyperbolic functions.jl")

max_stars_repo_path: src/rules/6 Hyperbolic functions/6.7 Miscellaneous/.jl


text:
using Primes

ispernicious(n::Integer) = isprime(count_ones(n))
nextpernicious(n::Integer) = begin n += 1; while !ispernicious(n) n += 1 end; return n end
function perniciouses(n::Int)
    rst = Vector{Int}(n)
    rst[1] = 3
    for i in 2:n
        rst[i] = nextpernicious(rst[i-1])
    end
    return rst
end
perniciouses(a::Integer, b::Integer) = filter(ispernicious, a:b)

println("First 25 pernicious numbers: ", join(perniciouses(25), ", "))
println("Perniciouses in [888888877, 888888888]: ", join(perniciouses(888888877, 888888888), ", "))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# Read in the data for the model

function get_Kenyadata(filename)
    KenyaTbl = readtable(filename)
    KenyaTbl = KenyaTbl[1:(end-1),:]
    n, = size(KenyaTbl)
    #Equator conversion 1 deg lat = 110.57km, 1 deg long = 111.32 km
    #Parse the string into an x,y tuple (in kms)
    KenyaTbl[:Loc] = [(0.,0.) for i = 1:n]
    KenyaTbl[:UrbanIndex] = ones(n)
    KenyaTbl[:Phase] = zeros(n)
    for i = 1:n
        x = join(split(KenyaTbl[:Location_1][i],"("))
        x = join(split(x,")"))
        y = split(x,",")
        z = parse(Float64,y[2])*111.32,parse(Float64,y[1])*110.57
        KenyaTbl[:Loc][i] = z
        if ~ismissing(KenyaTbl[:Rural][i])
            KenyaTbl[:UrbanIndex][i] = KenyaTbl[:Urban][i]/(KenyaTbl[:Urban][i]+KenyaTbl[:Rural][i])
        end
        KenyaTbl[:Phase] = 30.
    end
    return KenyaTbl,n
end

function get_flightdata(filename)
    flights=readtable(filename)
    into_mombassa=convert(Vector,flights[:,1])
    into_nairobi=convert(Vector,flights[:,2])
    return into_mombassa,into_nairobi
end

function get_prevdata(filename)
    prev_table=readtable(filename)
    global_prev=convert(Vector,prev_table[:,1])
    return global_prev
end


text:
"""
**AnalyticVI**

Variational Inference solver for conjugate or conditionally conjugate likelihoods (non-gaussian are made conjugate via augmentation)
All data is used at each iteration (use AnalyticSVI for Stochastic updates)

```julia
AnalyticVI(;ϵ::T=1e-5)
```
**Keywords arguments**

    - `ϵ::T` : convergence criteria
"""
mutable struct AnalyticVI{T<:Real} <: Inference{T}
    ϵ::T #Convergence criteria
    nIter::Integer #Number of steps performed
    optimizer_η₁::LatentArray{Optimizer} #Learning rate for stochastic updates
    optimizer_η₂::LatentArray{Optimizer} #Learning rate for stochastic updates
    Stochastic::Bool #Use of mini-batches
    nSamples::Int64
    nSamplesUsed::Int64 #Size of mini-batches
    MBIndices::Vector{Int64} #Indices of the minibatch
    ρ::T #Stochastic Coefficient
    HyperParametersUpdated::Bool #To know if the inverse kernel matrix must updated
    ∇η₁::LatentArray{Vector{T}}
    ∇η₂::LatentArray{Matrix{T}} #Stored as a matrix since symmetric sums do not help for the moment WARNING
    function AnalyticVI{T}(ϵ::T,nIter::Integer,optimizer_η₁::AbstractVector{<:Optimizer},optimizer_η₂::AbstractVector{<:Optimizer},Stochastic::Bool,nSamples::Integer,nSamplesUsed::Integer,MBIndices::AbstractVector,ρ::T,flag::Bool) where T
        return new{T}(ϵ,nIter,optimizer_η₁,optimizer_η₂,Stochastic,nSamples,nSamplesUsed,MBIndices,ρ,flag)
    end
    function AnalyticVI{T}(ϵ::T,nIter::Integer,optimizer_η₁::AbstractVector{<:Optimizer},optimizer_η₂::AbstractVector{<:Optimizer},Stochastic::Bool,nSamples::Integer,nSamplesUsed::Integer,MBIndices::AbstractVector,ρ::T,flag::Bool,∇η₁::AbstractVector{<:AbstractVector},
    ∇η₂::AbstractVector{<:AbstractMatrix}) where T
        return new{T}(ϵ,nIter,optimizer_η₁,optimizer_η₂,Stochastic,nSamples,nSamplesUsed,MBIndices,ρ,flag,∇η₁,∇η₂)
    end
end

function AnalyticVI(;ϵ::T=1e-5) where {T<:Real}
    AnalyticVI{Float64}(ϵ,0,[VanillaGradDescent(η=1.0)],[VanillaGradDescent(η=1.0)],false,1,1,[1],1.0,true)
end

"""
**AnalyticSVI**
Stochastic Variational Inference solver for conjugate or conditionally conjugate likelihoods (non-gaussian are made conjugate via augmentation)

```julia
AnalyticSVI(nMinibatch::Integer;ϵ::T=1e-5,optimizer::Optimizer=ALRSVI())
```
    - `nMinibatch::Integer` : Number of samples per mini-batches

**Keywords arguments**

    - `ϵ::T` : convergence criteria
    - `optimizer::Optimizer` : Optimizer used for the variational updates. Should be an Optimizer object from the [GradDescent.jl](https://github.com/jacobcvt12/GradDescent.jl) package. Default is `ALRSVI()` (Adaptive Learning Rate for Stochastic Variational Inference)
"""
function AnalyticSVI(nMinibatch::Integer;ϵ::T=1e-5,optimizer::Optimizer=ALRSVI()) where {T<:Real}
    AnalyticVI{T}(ϵ,0,[optimizer],[optimizer],true,1,nMinibatch,1:nMinibatch,1.0,true)
end

function Base.show(io::IO,inference::AnalyticVI{T}) where T
    print(io,"Analytic$(inference.Stochastic ? " Stochastic" : "") Variational Inference")
end

"""Initialize the final version of the inference object"""
function init_inference(inference::AnalyticVI{T},nLatent::Integer,nFeatures::Integer,nSamples::Integer,nSamplesUsed::Integer) where {T<:Real}
    inference.nSamples = nSamples
    inference.nSamplesUsed = nSamplesUsed
    inference.MBIndices = 1:nSamplesUsed
    inference.ρ = nSamples/nSamplesUsed
    inference.optimizer_η₁ = [copy(inference.optimizer_η₁[1]) for _ in 1:nLatent]
    inference.optimizer_η₂ = [copy(inference.optimizer_η₂[1]) for _ in 1:nLatent]
    inference.∇η₁ = [zeros(T,nFeatures) for _ in 1:nLatent];
    inference.∇η₂ = [Matrix(Diagonal(ones(T,nFeatures))) for _ in 1:nLatent]
    return inference
end

"""Generic method for variational updates using analytical formulas"""
function variational_updates!(model::AbstractGP{L,AnalyticVI{T}}) where {L<:Likelihood,T}
    local_updates!(model)
    natural_gradient!(model)
    global_update!(model)
end

"""Coordinate ascent updates on the natural parameters"""
function natural_gradient!(model::VGP{L,AnalyticVI{T}}) where {T<:Real,L<:Likelihood{T}}
    model.η₁ .= ∇μ(model) .+ model.invKnn.*model.μ₀
    model.η₂ .= -0.5*Symmetric.(Diagonal{T}.(∇Σ(model)).+model.invKnn)
end

"""Computation of the natural gradient for the natural parameters"""
function natural_gradient!(model::SVGP{L,AnalyticVI{T}}) where {T<:Real,L<:Likelihood{T}}
    model.inference.∇η₁ .= model.inference.ρ.*transpose.(model.κ).*∇μ(model) .+ model.invKmm.*model.μ₀ .- model.η₁
    model.inference.∇η₂ .= -0.5*(model.inference.ρ.*transpose.(model.κ).*Diagonal{T}.(∇Σ(model)).*model.κ.+model.invKmm) .- model.η₂
end

"""Conversion from natural to standard distribution parameters"""
function global_update!(model::VGP{L,AnalyticVI{T}}) where {L<:Likelihood,T}
    model.Σ .= -0.5.*inv.(model.η₂)
    model.μ .= model.Σ.*model.η₁
end

"""Update of the natural parameters and conversion from natural to standard distribution parameters"""
function global_update!(model::SVGP{L,AnalyticVI{T}}) where {L<:Likelihood,T}
    if model.inference.Stochastic
        model.η₁ .= model.η₁ .+ GradDescent.update.(model.inference.optimizer_η₁,model.inference.∇η₁)
        model.η₂ .= Symmetric.(model.η₂ .+ GradDescent.update.(model.inference.optimizer_η₂,model.inference.∇η₂))
    else
        model.η₁ .+= model.inference.∇η₁
        model.η₂ .= Symmetric.(model.inference.∇η₂ .+ model.η₂)
    end
    model.Σ .= -0.5.*inv.(model.η₂)
    model.μ .= model.Σ.*model.η₁
end


text:
# TODO: should be in StaticLint. visited check is costly.
resolve_shadow_binding(b) = b
function resolve_shadow_binding(b::StaticLint.Binding, visited=StaticLint.Binding[])
    if b in visited
        throw(LSInfiniteLoop("Inifinite loop in bindings."))
    else
        push!(visited, b)
    end
    if b.val isa StaticLint.Binding
        return resolve_shadow_binding(b.val, visited)
    else
        return b
    end
end

function get_definitions(x, tls, env, locations) end # Fallback

function get_definitions(x::SymbolServer.ModuleStore, tls, env, locations)
    if haskey(x.vals, :eval) && x[:eval] isa SymbolServer.FunctionStore
        get_definitions(x[:eval], tls, env, locations)
    end
end

function get_definitions(x::Union{SymbolServer.FunctionStore,SymbolServer.DataTypeStore}, tls, env, locations)
    StaticLint.iterate_over_ss_methods(x, tls, env, function (m)
        if safe_isfile(m.file)
            push!(locations, Location(filepath2uri(m.file), Range(m.line - 1, 0, m.line - 1, 0)))
        end
        return false
    end)
end

function get_definitions(b::StaticLint.Binding, tls, env, locations)
    if !(b.val isa EXPR)
        get_definitions(b.val, tls, env, locations)
    end
    if b.type === StaticLint.CoreTypes.Function || b.type === StaticLint.CoreTypes.DataType
        for ref in b.refs
            method = StaticLint.get_method(ref)
            if method !== nothing
                get_definitions(method, tls, env, locations)
            end
        end
    elseif b.val isa EXPR
        get_definitions(b.val, tls, env, locations)
    end
end

function get_definitions(x::EXPR, tls::StaticLint.Scope, env, locations)
    doc1, o = get_file_loc(x)
    if doc1 isa Document
        push!(locations, Location(doc1._uri, Range(doc1, o .+ (0:x.span))))
    end
end

safe_isfile(s::Symbol) = safe_isfile(string(s))
safe_isfile(::Nothing) = false
function safe_isfile(s::AbstractString)
    try
        !occursin("\0", s) && isfile(s)
    catch err
        isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
        false
    end
end

function textDocument_definition_request(params::TextDocumentPositionParams, server::LanguageServerInstance, conn)
    locations = Location[]
    doc = getdocument(server, URI2(params.textDocument.uri))
    offset = get_offset(doc, params.position)
    x = get_expr1(getcst(doc), offset)
    if x isa EXPR && StaticLint.hasref(x)
        # Replace with own function to retrieve references (with loop saftey-breaker)
        b = refof(x)
        b = resolve_shadow_binding(b)
        (tls = StaticLint.retrieve_toplevel_scope(x)) === nothing && return locations
        get_definitions(b, tls, getenv(doc, server), locations)
    elseif x isa EXPR && CSTParser.isstringliteral(x)
        # TODO: move to its own function
        if valof(x) isa String && sizeof(valof(x)) < 256 # AUDIT: OK
            try
                if isabspath(valof(x)) && safe_isfile(valof(x))
                    push!(locations, Location(filepath2uri(valof(x)), Range(0, 0, 0, 0)))
                elseif !isempty(getpath(doc)) && safe_isfile(joinpath(_dirname(getpath(doc)), valof(x)))
                    push!(locations, Location(filepath2uri(joinpath(_dirname(getpath(doc)), valof(x))), Range(0, 0, 0, 0)))
                end
            catch err
                isa(err, Base.IOError) ||
                    isa(err, Base.SystemError) ||
                    (VERSION == v"1.2.0" && isa(err, ErrorException) && err.msg == "type Nothing has no field captures ") ||
                    rethrow()
            end
        end
    end

    return locations
end

function get_file_loc(x::EXPR, offset=0, c=nothing)
    if c !== nothing
        for a in x
            a == c && break
            offset += a.fullspan
        end
    end
    if parentof(x) !== nothing
        return get_file_loc(parentof(x), offset, x)
    elseif headof(x) === :file && StaticLint.hasmeta(x)
        return x.meta.error, offset
    else
        return nothing, offset
    end
end

function textDocument_formatting_request(params::DocumentFormattingParams, server::LanguageServerInstance, conn)
    doc = getdocument(server, URI2(params.textDocument.uri))
    newcontent = DocumentFormat.format(get_text(doc), server.format_options)
    end_l, end_c = get_position_at(doc, sizeof(get_text(doc))) # AUDIT: OK
    lsedits = TextEdit[TextEdit(Range(0, 0, end_l, end_c), newcontent)]

    return lsedits
end

function find_references(textDocument::TextDocumentIdentifier, position::Position, server)
    locations = Location[]
    doc = getdocument(server, URI2(textDocument.uri))
    offset = get_offset(doc, position)
    x = get_expr1(getcst(doc), offset)
    x === nothing && return locations
    for_each_ref(x) do r, doc1, o
        push!(locations, Location(doc1._uri, Range(doc1, o .+ (0:r.span))))
    end
    return locations
end

function for_each_ref(f, identifier::EXPR)
    if identifier isa EXPR && StaticLint.hasref(identifier) && refof(identifier) isa StaticLint.Binding
        for r in refof(identifier).refs
            if r isa EXPR
                doc1, o = get_file_loc(r)
                if doc1 isa Document
                    f(r, doc1, o)
                end
            end
        end
    end
end

function textDocument_references_request(params::ReferenceParams, server::LanguageServerInstance, conn)
    return find_references(params.textDocument, params.position, server)
end

function textDocument_rename_request(params::RenameParams, server::LanguageServerInstance, conn)
    tdes = Dict{String,TextDocumentEdit}()
    locations = find_references(params.textDocument, params.position, server)

    for loc in locations
        if loc.uri in keys(tdes)
            push!(tdes[loc.uri].edits, TextEdit(loc.range, params.newName))
        else
            doc = getdocument(server, URI2(loc.uri))
            tdes[loc.uri] = TextDocumentEdit(VersionedTextDocumentIdentifier(loc.uri, doc._version), [TextEdit(loc.range, params.newName)])
        end
    end

    return WorkspaceEdit(missing, collect(values(tdes)))
end

function textDocument_prepareRename_request(params::PrepareRenameParams, server::LanguageServerInstance, conn)
    doc = getdocument(server, URI2(params.textDocument.uri))
    x = get_expr1(getcst(doc), get_offset(doc, params.position))
    x isa EXPR || return nothing
    _, x_start_offset = get_file_loc(x)
    x_range = Range(doc, x_start_offset .+ (0:x.span))
    return x_range
end

is_valid_binding_name(name) = false
function is_valid_binding_name(name::EXPR)
    (headof(name) === :IDENTIFIER && valof(name) isa String && !isempty(valof(name))) ||
    CSTParser.isoperator(name) ||
    (headof(name) === :NONSTDIDENTIFIER && length(name.args) == 2 && valof(name.args[2]) isa String && !isempty(valof(name.args[2])))
end
function get_name_of_binding(name::EXPR)
    if headof(name) === :IDENTIFIER
        valof(name)
    elseif CSTParser.isoperator(name)
        string(Expr(name))
    elseif headof(name) === :NONSTDIDENTIFIER
        valof(name.args[2])
    else
        ""
    end
end

function textDocument_documentSymbol_request(params::DocumentSymbolParams, server::LanguageServerInstance, conn)
    uri = params.textDocument.uri
    doc = getdocument(server, URI2(uri))

    return collect_document_symbols(getcst(doc), server, doc)
end

function collect_document_symbols(x::EXPR, server::LanguageServerInstance, doc, pos=0, symbols=DocumentSymbol[])
    if bindingof(x) !== nothing
        b =  bindingof(x)
        if b.val isa EXPR && is_valid_binding_name(b.name)
            ds = DocumentSymbol(
                get_name_of_binding(b.name), # name
                missing, # detail
                _binding_kind(b), # kind
                false, # deprecated
                Range(doc, (pos .+ (0:x.span))), # range
                Range(doc, (pos .+ (0:x.span))), # selection range
                DocumentSymbol[] # children
            )
            push!(symbols, ds)
            symbols = ds.children
        end
    end
    if length(x) > 0
        for a in x
            collect_document_symbols(a, server, doc, pos, symbols)
            pos += a.fullspan
        end
    end
    return symbols
end

function collect_bindings_w_loc(x::EXPR, pos=0, bindings=Tuple{UnitRange{Int},StaticLint.Binding}[])
    if bindingof(x) !== nothing
        push!(bindings, (pos .+ (0:x.span), bindingof(x)))
    end
    if length(x) > 0
        for a in x
            collect_bindings_w_loc(a, pos, bindings)
            pos += a.fullspan
        end
    end
    return bindings
end

function collect_toplevel_bindings_w_loc(x::EXPR, pos=0, bindings=Tuple{UnitRange{Int},StaticLint.Binding}[]; query="")
    if bindingof(x) isa StaticLint.Binding && valof(bindingof(x).name) isa String && bindingof(x).val isa EXPR && startswith(valof(bindingof(x).name), query)
        push!(bindings, (pos .+ (0:x.span), bindingof(x)))
    end
    if scopeof(x) !== nothing && !(headof(x) === :file || CSTParser.defines_module(x))
        return bindings
    end
    if length(x) > 0
        for a in x
            collect_toplevel_bindings_w_loc(a, pos, bindings, query=query)
            pos += a.fullspan
        end
    end
    return bindings
end

function _binding_kind(b)
    if b isa StaticLint.Binding
        if b.type === nothing
            return 13
        elseif b.type == StaticLint.CoreTypes.Module
            return 2
        elseif b.type == StaticLint.CoreTypes.Function
            return 12
        elseif b.type == StaticLint.CoreTypes.String
            return 15
        elseif b.type == StaticLint.CoreTypes.Int || b.type == StaticLint.CoreTypes.Float64
            return 16
        elseif b.type == StaticLint.CoreTypes.DataType
            return 23
        else
            return 13
        end
    elseif b isa SymbolServer.ModuleStore
        return 2
    elseif b isa SymbolServer.MethodStore
        return 6
    elseif b isa SymbolServer.FunctionStore
        return 12
    elseif b isa SymbolServer.DataTypeStore
        return 23
    else
        return 13
    end
end

function julia_getModuleAt_request(params::VersionedTextDocumentPositionParams, server::LanguageServerInstance, conn)
    uri = URI2(params.textDocument.uri)

    if hasdocument(server, uri)
        doc = getdocument(server, uri)
        if doc._version == params.version
            offset = get_offset2(doc, params.position.line, params.position.character, true)
            x, p = get_expr_or_parent(getcst(doc), offset, 1)
            if x isa EXPR
                if x.head === :MODULE || x.head === :IDENTIFIER || x.head === :END
                    if x.parent !== nothing && x.parent.head === :module
                        x = x.parent
                        if CSTParser.defines_module(x)
                            x = x.parent
                        end
                    end
                end
                if CSTParser.defines_module(x) && p <= offset <= p + x[1].fullspan + x[2].fullspan
                    x = x.parent
                end

                scope = StaticLint.retrieve_scope(x)
                if scope !== nothing
                    return get_module_of(scope)
                end
            end
        else
            return mismatched_version_error(uri, doc, params, "getModuleAt")
        end
    else
        return nodocuemnt_error(uri)
    end
    return "Main"
end

function get_module_of(s::StaticLint.Scope, ms=[])
    if CSTParser.defines_module(s.expr) && CSTParser.isidentifier(s.expr.args[2])
        pushfirst!(ms, StaticLint.valofid(s.expr.args[2]))
    end
    if parentof(s) isa StaticLint.Scope
        return get_module_of(parentof(s), ms)
    else
        return isempty(ms) ? "Main" : join(ms, ".")
    end
end

function julia_getDocAt_request(params::VersionedTextDocumentPositionParams, server::LanguageServerInstance, conn)
    uri = URI2(params.textDocument.uri)
    hasdocument(server, uri) || return nodocuemnt_error(uri)

    doc = getdocument(server, uri)
    env = getenv(doc, server)
    if doc._version !== params.version
        return mismatched_version_error(uri, doc, params, "getDocAt")
    end

    x = get_expr1(getcst(doc), get_offset(doc, params.position))
    x isa EXPR && CSTParser.isoperator(x) && resolve_op_ref(x, env)
    documentation = get_hover(x, "", server)

    return documentation
end

# TODO: handle documentation resolving properly, respect how Documenter handles that
function julia_getDocFromWord_request(params::NamedTuple{(:word,),Tuple{String}}, server::LanguageServerInstance, conn)
    exact_matches = []
    approx_matches = []
    word_sym = Symbol(params.word)
    traverse_by_name(getsymbols(getenv(server))) do sym, val
        is_exact_match = sym === word_sym
        # this would ideally use the Damerau-Levenshtein distance or even something fancier:
        is_match = is_exact_match || REPL.levenshtein(string(sym), string(word_sym)) <= 1
        if is_match
            val = get_hover(val, "", server)
            if !isempty(val)
                push!(is_exact_match ? exact_matches : approx_matches, val)
            end
        end
    end
    if isempty(exact_matches) && isempty(approx_matches)
        return "No results found."
    else
        return join(isempty(exact_matches) ? approx_matches[1:min(end, 10)] : exact_matches, "\n---\n")
    end
end

function textDocument_selectionRange_request(params::SelectionRangeParams, server::LanguageServerInstance, conn)
    doc = getdocument(server, URI2(params.textDocument.uri))
    ret = map(params.positions) do position
        offset = get_offset(doc, position)
        x = get_expr1(getcst(doc), offset)
        get_selection_range_of_expr(x)
    end
    return ret isa Vector{SelectionRange} ?
        ret :
        nothing
end

# Just returns a selection for each parent EXPR, should be more selective
get_selection_range_of_expr(x) = missing
function get_selection_range_of_expr(x::EXPR)
    doc, offset = get_file_loc(x)
    l1, c1 = get_position_at(doc, offset)
    l2, c2 = get_position_at(doc, offset + x.span)
    SelectionRange(Range(l1, c1, l2, c2), get_selection_range_of_expr(x.parent))
end


text:
abstract type
    AbstractIDSequence{S<:ValueSupport} <: MultivariateDistribution{S}
end

length(X::AbstractIDSequence) = length(distributions(X))

function insupport(X::AbstractIDSequence,x::AbstractVector)
  length(x)==length(X) && all(dx->insupport(dx[1],dx[2]),zip(distributions(X),x))
end

function _rand!{T<:Real}(X::AbstractIDSequence, x::AbstractVector{T})
  for (i,d_i) in enumerate(distributions(X))
    x[i] = rand(d_i)
  end
  return x
end

function _logpdf{T<:Real}(X::AbstractIDSequence, x::AbstractVector{T})
  return sum(arg->logpdf(arg[1],arg[2]),zip(distributions(X),x))
end

mean(X::AbstractIDSequence) = Float64[mean(d_i) for d_i in distributions(X)]
var(X::AbstractIDSequence) = Float64[var(d_i) for d_i in distributions(X)]
cov(X::AbstractIDSequence) = Diagonal(var(X))
entropy(X::AbstractIDSequence) = sum(entropy,distributions(X))

function mgf(X::AbstractIDSequence, t::AbstractVector)
  prod([mgf(d_i,t_i) for (d_i,t_i) in zip(distributions(X),t)])
end
function cf(X::AbstractIDSequence, t::AbstractVector)
  prod([cf(d_i,t_i) for (d_i,t_i) in zip(distributions(X),t)])
end


text:
function plot_spectrogram(S,titre)
    S = abs(S[1:Int(size(S)[1]/2),:])
    S = log.(S + 1e-4)
    imshow(S, cmap = get_cmap("jet"), interpolation = "nearest")
    title(titre)
end


text:
module test_emojisymbols_emoji

using Test

legacy_emoji = ["☺", "☎", "🅰", "▫", "㊗", "◼", "⤵", "✖", "❄", "‼", "↔", "☑", "♥", "✉", "♣", "✒", "✈", "®", "〰", "🅱", "↪", "⬅", "⤴", "◀", "❇", "♨", "◻", "✴", "✔", "🈷", "▶", "☁", "♠", "ℹ", "☀", "⁉", "©", "⬇", "↙", "⚠", "〽", "▪", "✂", "✌", "↗", "™", "Ⓜ", "↘", "➡", "↖", "㊙", "☝", "↕", "⬆", "♻", "🅿", "🈂", "♦", "↩", "✏", "🅾", "❤", "✳"]

using REPL
if VERSION < v"1.7.0-DEV.849"
    @test !haskey(REPL.REPLCompletions.emoji_symbols, "\\:thinking_face:")
else
    @test haskey(REPL.REPLCompletions.emoji_symbols, "\\:thinking_face:")
end
s1 = values(copy(REPL.REPLCompletions.emoji_symbols))
@test intersect(legacy_emoji, s1) == legacy_emoji

using EmojiSymbols
s2 = values(EmojiSymbols.emoji_symbols)
@test isempty(intersect(legacy_emoji, s2))
@test REPL.REPLCompletions.emoji_symbols["\\:thinking_face:"] == "🤔"

@test EmojiSymbols.emoji_name_table["🤔"] == "\\:thinking_face:"
@test sprint(show, MIME"text/plain"(), '🤔') == "\\:thinking_face: '🤔': Unicode U+1F914 (category So: Symbol, other)"
@test sprint(show, MIME"text/plain"(), "🤔") == "\\:thinking_face: \"🤔\""
@test sprint(show, MIME"text/plain"(), '☔') == "\\:umbrella_with_rain_drops: '☔': Unicode U+2614 (category So: Symbol, other)"

@test string('🤔') == string("🤔") == "🤔"

end # module test_emojisymbols_emoji


text:
@reexport module Maze

export MazeEnv

using ReinforcementLearningBase

import Base: *

const Actions = [
    CartesianIndex(0, -1),  # left
    CartesianIndex(0, 1),   # right
    CartesianIndex(-1, 0),  # up
    CartesianIndex(1, 0),   # down
]

mutable struct MazeEnv <: AbstractEnv
    walls::Set{CartesianIndex{2}}
    position::CartesianIndex{2}
    start::CartesianIndex{2}
    goal::CartesianIndex{2}
    NX::Int
    NY::Int
    observation_space::DiscreteSpace
    action_space::DiscreteSpace
    MazeEnv(w, p, s, g, NX, NY) =
        new(w, p, s, g, NX, NY, DiscreteSpace(NX * NY), DiscreteSpace(length(Actions)))
end

RLBase.get_observation_space(env::MazeEnv) = env.observation_space
RLBase.get_action_space(env::MazeEnv) = env.action_space

function MazeEnv()
    walls = Set([
        [CartesianIndex(i, 3) for i = 2:4]
        CartesianIndex(5, 6)
        [CartesianIndex(j, 8) for j = 1:3]
    ])
    start = CartesianIndex(3, 1)
    goal = CartesianIndex(1, 9)
    MazeEnv(walls, start, start, goal, 6, 9)
end

function extend(p::CartesianIndex{2}, n::Int)
    x, y = Tuple(p)
    [CartesianIndex(n * (x - 1) + i, n * (y - 1) + j) for i = 1:n for j = 1:n]
end

function remap(p::CartesianIndex{2}, n::Int)
    x, y = Tuple(p)
    CartesianIndex((x - 1) * n + 1, (y - 1) * n + 1)
end

function *(env::MazeEnv, n::Int)
    walls = Set{CartesianIndex{2}}(ww for w in env.walls for ww in extend(w, n))
    start, position, goal = remap(env.start, n), remap(env.position, n), remap(env.goal, n)
    NX, NY = env.NX * n, env.NY * n
    MazeEnv(walls, position, start, goal, NX, NY)
end

function (env::MazeEnv)(a::Int)
    p = env.position + Actions[a]
    if p == env.goal
        env.position = env.goal
    elseif !(p ∈ env.walls)
        env.position = CartesianIndex(min(max(p[1], 1), env.NX), min(max(p[2], 1), env.NY))
    end
    nothing
end

RLBase.observe(env::MazeEnv) =
    (
        reward = Float64(env.position == env.goal),
        terminal = env.position == env.goal,
        state = (env.position[2] - 1) * env.NX + env.position[1],
    )

function RLBase.reset!(env::MazeEnv)
    env.position = env.start
    nothing
end

end


text:
#!/usr/bin/env julia
using JSON
using Glob
using Suppressor
using Dates
using Statistics
using Random
using LinearAlgebra

runners = filter(x -> strip(x) != "", split(read(ignorestatus(pipeline(
    `ps -U $(ENV["USER"]) -ux`, `grep '[0-9]*:[0-9]*\s*julia\s*.*population.jl'`)), String), "\n"))

if length(runners) == 0
    error("cannot detect self, the programe must be wrong!")
elseif length(runners) > 1
    println(join(runners, "\n"))
    warn("population already running, exiting...")
    exit()
end

if length(ARGS) != 1
    @error "please specify config file"
end

CONFIG_FILE = ARGS[1]
CONFIG = JSON.parse(read(CONFIG_FILE, String))
const TOTAL_STAGE = ceil(Int, CONFIG["total_steps"] / CONFIG["step_interval"])

const populationroot = expanduser("./jobs/population")
const jobdeploy = expanduser("~/jobs/queue")
const jobdone = expanduser("~/jobs/done")
const jobs = expanduser("~/jobs")
const jobroot = joinpath(populationroot, CONFIG["name"])
const jobscript = joinpath(jobroot, "script")
const jobresult = joinpath(jobroot, "result") # results
const jobhistory = joinpath(jobroot, "history")
const logfile = joinpath(jobroot, "population.log")
mkpath.([populationroot, jobroot, jobscript, jobresult, jobhistory])

stage_history_path(id) = joinpath(jobhistory, "stage$(lpad(id, 4, "0")).json")
save_stage(stage) = open(f -> println(f, json(stage, 4)), stage_history_path(stage["id"]), "w")
vec_str(v) = "[$(join(v, ", "))]"
weighted_mean(r) = round(dot(r, normalize([0.95^i for i in length(r):-1:1], 1)), digits=4)

@suppress Base.println(xs...) = open(f -> (println(f, "[$(now())] ", xs...);
    println(stdout, "[$(now())] ", xs...)), logfile, "a")

function next_stage(config)
    # check history folder, which contais info of each stage
    id = 1
    vip = config["vip"]
    weight_dir = ""
    historyfiles = glob(joinpath(jobhistory, "stage????.json"))
    if length(historyfiles) > 0
        id = maximum(parse.(Int, getindex.(basename.(historyfiles), [6:9])))
    end

    if isfile(stage_history_path(id))
        # check if finished
        stage = JSON.parse(read(stage_history_path(id), String))
        @assert(length(stage["population"]) > 0)
        @assert(stage["id"] == id)
        @assert(length(stage["vip"]) == length(config["vip"]))
        if all([p["status"] == "done" for p in stage["population"]])
            rewards = [p["reward"] for p in stage["population"]]
            println("all population done in stage $id, rewards $rewards")
            vip_id = argmax(rewards)
            stage["best_person"] = vip_id
            stage["best_reward"] = stage["population"][vip_id]["reward"]
            stage["next_vip"] = stage["population"][vip_id]["config"]
            stage["next_weight_dir"] = stage["population"][vip_id]["weight_dir"]
            save_stage(stage)

            vip = stage["next_vip"]
            weight_dir = stage["next_weight_dir"]
            println("VIP person $(vip_id) $(vec_str(vip)) selected for its average reward $(maximum(rewards))")
            id += 1
        else
            println("stage $id unfinished, continuing...")
            println("stage $id current status:\n",
                    join(["$(p["runname"]): $(p["status"]), config $(vec_str(p["config"]))" *
                          (p["status"] == "done" ? ", reward $(p["reward"])" : "")
                          for p in stage["population"]], "\n"))
            return stage
        end
    end

    # end?
    if id > TOTAL_STAGE
        println("all $TOTAL_STAGE stages done")
        return nothing
    end

    # need create a new stage
    stage = Dict(
        "id" => id,
        "vip" => vip,
        "max_update" => config["step_interval"] * id,
        "weight_dir" => weight_dir
    )

    # generate population
    population = filter(p -> all(p .<= config["upper_bound"]) && all(p .>= config["lower_bound"]),
                        [vip] .+ config["variations"])
    if length(population) < config["pool"]
        popul2 = filter(p -> all(p .<= config["upper_bound"]) && all(p .>= config["lower_bound"]),
                        [vip] .+ config["additional_variations"])
        population = vcat(population, popul2[randperm(length(popul2))], repeat([vip], config["pool"]))[1:config["pool"]]
    end
    @assert length(population) == config["pool"]

    if length(population) == 0
        println("no population candidate, exiting")
        return nothing
    end

    stage["population"] = Dict.(Pair.("i", 1:length(population)), ["status" => "start"], Pair.("config", population))

    # save
    save_stage(stage)

    println("stage $id history file generated")
    return stage
end

"""
    population status: start | deployed | runover | done
    call this function to go to next state, return the result state
"""
function process_population(stage, i, config)
    id = stage["id"]
    p = stage["population"][i]
    if p["status"] == "done"
        return true
    elseif p["status"] == "start"
        # gen script and deploy
        p["runname"] = "$(config["name"])-stage$(lpad(id, 4, "0"))-person$(lpad(i, 3, "0"))"
        template = read(config["template"], String)
        if length(template) > 2 && template[1] == '"' && template[end] == '"'
            template = strip(template[2:end-1])
        elseif template == ""
            template = "exit 1\necho Empty template!"
        end
        lines = strip.(split(template, "\n", keepempty=false))
        lines[end] *= " --maps '{$(join(["\"$m\": $n" for (m, n) in zip(config["envs"], p["config"])], ", "))}'"
        lines[end] *= " --max_update $(stage["max_update"])"
        lines[end] *= " --run_id $(p["runname"])"
        lines[end] *= " --save_interval $(min(250, ceil(Int, config["step_interval"] / 4)))"
        lines[end] *= " --num_snapshot 1"
        id > 1 && (lines[end] *= " --remote_restore $(stage["weight_dir"])")
        script = join(lines, "\n")

        scriptfilename = joinpath(jobscript, "$(p["runname"]).sh")
        open(f -> println(f, script), scriptfilename, "w")
        cp(scriptfilename, joinpath(jobdeploy, "$(p["runname"]).sh"), force=true)

        p["status"] = "deployed"
    elseif p["status"] == "deployed"
        # check if done, if not, whether in queue or in jobroot, else print a warn
        doneshfile = glob("\\[DONE*\\]$(p["runname"]).sh", jobdone)
        if length(doneshfile) >= 1
            p["status"] = "runover"
        else
            return false
        end
    elseif p["status"] == "runover"
        # summaryise
        # find log files
        logfiles = unique(vcat(glob("\\[DONE*\\]$(p["runname"]).sh.log", jobdone),
                               glob("\\[STOP*\\]$(p["runname"]).sh.log", jobs),
                               glob("\\[?-ORZ*\\]$(p["runname"]).sh.log", jobs),
                               glob("\\[?-ERR*\\]$(p["runname"]).sh.log", jobs)))

        @assert length(logfiles) >= 1
        p["weight_dir"] = strip(split(read(pipeline(`cat $(logfiles[1])`, `grep 'weights are saved at'`, `cut -b 22-`), String), "\n")[1])
        rewards = []
        for f in logfiles
            try
                r = read(pipeline(`cat $f`, `grep $(config["envs"][1] * "_rew_mean")`, `cut -d'|' -f3`), String)
                append!(rewards, parse.(Float32, split(r, "\n", keepempty=false)))
            catch y
                println(y)
            end
        end
        @assert length(rewards) > 0
        p["rewards"] = join(rewards, ",")
        p["reward"] = weighted_mean(rewards)
        for f in logfiles
            for g in [f, f[1:end-4], f[1:end-4] * ".bk"]
                isfile(g) && mv(g, joinpath(jobresult, basename(g)))
            end
        end
        println("stage $id population $i config $(vec_str(p["config"])), weighted mean reward $(p["reward"])",
                "mean reward $(mean(rewards)), std $(std(rewards)), max $(maximum(rewards)), min $(minimum(rewards)), count $(length(rewards))")
        p["status"] = "done"
    else
        println("warn: stage $id population $i in unknown state")
    end

    println("stage $id population $i config $(vec_str(p["config"])) goto state $(p["status"])")
    return false
end

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

function main(config)
    println("whaaa! a new day!")
    println("config file $CONFIG_FILE loaded\n", json(config))
    # get where we are
    stage = next_stage(config)
    while stage != nothing
        println("now on stage $(stage["id"]) / $TOTAL_STAGE, updates from $(stage["max_update"] - config["step_interval"]) to $(stage["max_update"])")
        n = length(stage["population"])
        while true
            n_done = 0
            for i in 1:n
                n_done += process_population(stage, i, config)
                save_stage(stage)
                sleep(0.5)
            end
            n_done == length(stage["population"]) && break
        end
        stage = next_stage(config)
    end
end

main(CONFIG)

max_stars_repo_name: innerlee/runner

max_issues_repo_name: innerlee/runner

max_forks_repo_name: innerlee/runner


text:
## Matlab Author: Rodrigo de Barros Vimieiro
## Translated by Sonia Minseo Kim in 2022

## 2-D Branchless Distance Driven Code

export bdd_2d

# using Plots: plot
using ImagePhantoms: shepp_logan, SheppLoganToft
using LazyGrids: ndgrid
using Dierckx: Spline1D

## Projection
"""
    projection(phantom::AbstractMatrix{<:T}, geo)
Generates a sinogram using the forward projection algorithm,
for a phantom image and a tuple of geometry definitions
"""
function projection(phantom::AbstractMatrix{<:T}, geo ; draw::Bool = false) where {T <: Number}

    DSD = geo.DSD
    DSO = geo.DSO
    pSize = geo.pSize
    dSize = geo.dSize
    nPix = geo.nPix
    nDet = geo.nDet
    theta = geo.theta # vector of projection view angles in radians

    # Detector boundaries
    detX = (-(nDet/2):(nDet/2)) .* dSize
    detY = (-(DSD-DSO)-(dSize/2)) .* ones(nDet+1)

    # Pixel boundaries
    (pixelX,pixelY) = ndgrid((-(nPix/2):(nPix/2)), (-(nPix/2):(nPix/2)))
    pixelX = pixelX' .* pSize
    pixelY = reverse(pixelY' .- pSize/2, dims = 1)

    # Tube
    tubeX = 0
    tubeY = DSO

    # Iso-center
    isoX = geo.isoX
    isoY = geo.isoY

    sinogram = zeros(length(theta),nDet)

    # For each projection
    for proj in 1:length(theta)

        angle = theta[proj]

        # Tube rotation
        rtubeX = ( (tubeX - isoX )*cos(angle) - (tubeY - isoY )*sin(angle) ) + isoX
        rtubeY = ( (tubeX - isoX )*sin(angle) + (tubeY - isoY )*cos(angle) ) + isoY

        # Detector rotation
        rdetX = ( (detX .- isoX ).*cos(angle) - (detY .- isoY ).*sin(angle) ) .+ isoX
        rdetY = ( (detX .- isoX ).*sin(angle) + (detY .- isoY ).*cos(angle) ) .+ isoY

        if (draw == true)
            # drawGeo(rtubeX,rtubeY,rdetX,rdetY)
        end

        # Define angle case & which axis it project boundaries
        # Case 1
        if (((angle>=0)&&(angle<=pi/4))||(angle>=7*pi/4))
            axisXCase = true   # Map on X axis()
            angleCase = 1
            c1=0
            c2=1
        # Case 2
        elseif ((angle>pi/4) && (angle<3*pi/4))
            axisXCase = false   # Map on Y axis()
            angleCase = 2
            c1=0
            c2=1
        # Case 3
        elseif ((angle>=3*pi/4)&&(angle<=5*pi/4))
            axisXCase = true   # Map on X axis()
            angleCase = 3
            c1=0
            c2=-1
        else # Case 4
            axisXCase = false   # Map on Y axis()
            angleCase = 4
            c1=0
            c2=-1
        end

        # Mapping boundaries into a common axis()
        if (axisXCase)
            detm = mapp2x.(rtubeX,rtubeY,rdetX,rdetY)
            pixm = mapp2x.(rtubeX,rtubeY,pixelX,pixelY)
            img = phantom
        else
            detm = mapp2y.(rtubeX,rtubeY,rdetX,rdetY)
            pixm = reverse(mapp2y.(rtubeX,rtubeY,pixelX,pixelY)', dims = 2)
            img = reverse(phantom', dims = 2)
        end

        center_det = floor(Int, (nDet+1)/2+1)

        if (axisXCase)
            # X-Ray pixel intersection calculation
            L = abs(pSize/cos(angle)) # This account for parallel-beam
            # Correction for fan-beam
            L1 = @. sqrt((rtubeX-detm[center_det])^2+(rtubeY-0)^2)/sqrt((rtubeX-detm[1:nDet]')^2+(rtubeY-0)^2)
        else
            # X-Ray pixel intersection calculation
            L = abs(pSize/sin(angle)) # This account for parallel-beam
            # Correction for fan-beam
            L1 = @. sqrt((rtubeX-0)^2+(rtubeY-detm[center_det])^2)/sqrt((rtubeX-0)^2+(rtubeY-detm[1:nDet]')^2)
        end
        L = L ./ L1

        pixIstart = 1
        pixIinc = 1
        if ((angleCase == 1)||(angleCase == 2))
            detIstart = 1
            detIinc = 1
        else
            detIstart = nDet+1
            detIinc = -1
        end

        deltaDetm = detm[detIstart+detIinc] - detm[detIstart] # Mapped detector length()
        deltaPixm = pixm[1,2] - pixm[1,1]   # Mapped pixel length()

        sinoTmp = zeros(1, nDet) # one "row" of sinogram

        # For each image row
        for row in 1:nPix

            rowm = pixm[row,:] # Get first mapped row from image

            detInd = detIstart
            pixInd = pixIstart

            # Find first detector overlap mapped with pixel mapped [Case 1]
            if (detm[detInd]-rowm[pixIstart]<=deltaDetm)
                while detm[detInd]-rowm[pixIstart]<=deltaDetm
                    detInd = detInd + detIinc
                end
            else
            # Find first pixel overlap mapped with detector mapped [Case 2]
                if (detm[detIstart]-rowm[pixInd]>deltaPixm)
                    while detm[detIstart]-rowm[pixInd]>deltaPixm
                        pixInd = pixInd + pixIinc
                    end
                end
            end

            Ppj = integrate1D(img[row,:],deltaPixm)

            if rowm[1] > rowm[2] #if descending
                Pdk = Spline1D(reverse(rowm), reverse(Ppj); k=1, bc="zero")(detm)
                Pdk = reverse(Pdk)
            else
                Pdk = Spline1D(rowm, Ppj; k=1, bc="zero")(detm)
            end

            while detm[detInd+c2]<rowm[end]
                sinoTmp[detInd] = sinoTmp[detInd]+(Pdk[detInd+c2]-Pdk[detInd])/deltaDetm
                detInd = detInd + detIinc
            end
            sinoTmp[detInd] = sinoTmp[detInd]+(Ppj[end]-Pdk[detInd])/deltaDetm

        end # Row loop
        sinogram[proj,:] = sinoTmp .* L

    end # Projection loop

    return sinogram

end # endfunc

## Backprojection
"""
    backprojection(sinogram::AbstractMatrix{<:T}, geo)
Generates a reconstructed image using the back projection algorithm,
for a sinogram and a tuple of geometry definitions
"""
function backprojection(sinogram::AbstractMatrix{<:T},geo ; draw::Bool = false) where {T <: Number}

    DSD = geo.DSD
    DSO = geo.DSO
    pSize = geo.pSize
    dSize = geo.dSize
    nPix = geo.nPix
    nDet = geo.nDet
    theta = geo.theta

    # Detector boundaries
    detX = (-(nDet/2):(nDet/2)) .* dSize
    detY = (-(DSD-DSO)-(dSize/2)) .* ones(nDet+1)

    # Pixel boundaries
    (pixelX,pixelY) = ndgrid(-(nPix/2):(nPix/2), -(nPix/2):(nPix/2))
    pixelX = pixelX' .* pSize
    pixelY = reverse(pixelY' .- pSize/2, dims = 1)

    # Tube
    tubeX = 0
    tubeY = DSO

    # Iso-center
    isoX = geo.isoX
    isoY = geo.isoY

    reconImg = zeros(nPix,nPix)
    reconImgTmp = reconImg

    # For each projection
    for proj in 1:length(theta)

        angle = theta[proj]

        # Tubre rotation
        rtubeX = ( (tubeX - isoX )*cos(angle) - (tubeY - isoY )*sin(angle) ) + isoX
        rtubeY = ( (tubeX - isoX )*sin(angle) + (tubeY - isoY )*cos(angle) ) + isoX

        # Detector rotation
        rdetX = ( (detX .- isoX )*cos(angle) - (detY .- isoY )*sin(angle) ) .+ isoX
        rdetY = ( (detX .- isoX )*sin(angle) + (detY .- isoY )*cos(angle) ) .+ isoX

        if (draw == true)
            # drawGeo(rtubeX,rtubeY,rdetX,rdetY)
        end

        # Define angle case & which axis it project boundaries
        # Case 1
        if (((angle>=0)&&(angle<=pi/4))||(angle>=7*pi/4))
            axisXCase = true   # Map on X axis()
            angleCase = 1
            c1=1
            c2=0
        # Case 2
        elseif ((angle>pi/4)&&(angle<3*pi/4))
            axisXCase = false   # Map on Y axis()
            angleCase = 2
            c1=1
            c2=0
        # Case 3
        elseif (((angle>=3*pi/4)&&(angle<=5*pi/4)))
            axisXCase = true   # Map on X axis()
            angleCase = 3
            c1=0
            c2=1
        else # Case 4
            axisXCase = false   # Map on Y axis()
            angleCase = 4
            c1=0
            c2=1
        end

        # Mapping boundaries into a common axis()
        if (axisXCase)
            detm = mapp2x.(rtubeX,rtubeY,rdetX,rdetY)
            pixm = mapp2x.(rtubeX,rtubeY,pixelX,pixelY)
        else
            detm = mapp2y.(rtubeX,rtubeY,rdetX,rdetY)
            pixm = reverse(mapp2y.(rtubeX,rtubeY,pixelX,pixelY)', dims = 2)
        end

        center_det = floor(Int, (nPix+1)/2+1)

        if (axisXCase)
            # X-Ray pixel intersection calculation
            L = abs(pSize/cos(angle)) # This account for parallel-beam
            # Correction for fan-beam
            L1 = @. sqrt((rtubeX-pixm[center_det])^2+(rtubeY-0)^2)/sqrt((rtubeX-pixm[1:nPix]')^2+(rtubeY-0)^2)
        else
            # X-Ray pixel intersection calculation
            L = abs(pSize/sin(angle)) # This account for parallel-beam
            # Correction for fan-beam
            L1 = @. sqrt((rtubeX-0)^2+(rtubeY-pixm[center_det])^2)/sqrt((rtubeX-0)^2+(rtubeY-pixm[1:nPix]')^2)
        end
        L = L ./ L1

        pixIstart = 1
        pixIinc = 1
        if ((angleCase == 1)||(angleCase == 2))
            detIstart = 1
            detIinc = 1
        else
            detIstart = nDet+1
            detIinc = -1
        end

        deltaDetm = detm[detIstart+detIinc]- detm[detIstart] # Mapped detector length()
        deltaPixm = pixm[1,2]- pixm[1,1]   # Mapped pixel length()

        # For each row
        for row in 1:nPix

            rowm = pixm[row,:] # Get first mapped row from image --> changed to column vec
            reconTmp = zeros(1,nPix)

            detInd = detIstart
            pixInd = pixIstart

            # Find first detector overlap mapped with pixel mapped [Case 1]
            if (detm[detInd]-rowm[pixIstart]<=deltaDetm)
                while ((detm[detInd]-rowm[pixIstart]<=deltaDetm))
                    detInd = detInd + detIinc
                end
            else
            # Find first pixel overlap mapped with detector mapped [Case 2]
                if (detm[detIstart]-rowm[pixInd]>deltaPixm)
                    while (detm[detIstart]-rowm[pixInd]>deltaPixm)
                        pixInd = pixInd + pixIinc
                    end
                end
            end

            detm = vec(detm)
            Ppj = integrate1D(sinogram[proj,:],deltaDetm)

            if detm[1] > detm[2] #if descending
                Pdk = Spline1D(reverse(detm), reverse(vec(Ppj)); k=1, bc="zero")(vec(rowm))
            else
                Pdk = Spline1D(detm, vec(Ppj); k=1, bc="zero")(vec(rowm))
            end

            while (pixInd < nPix+1)
                reconTmp[pixInd] = reconTmp[pixInd]+(Pdk[pixInd+c1]-Pdk[pixInd+c2])/deltaDetm
                pixInd = pixInd + pixIinc
            end

            reconImgTmp[row,:] = reconTmp .* L

        end # Row loop

        if ((angleCase == 1)||(angleCase == 3))
            reconImg = reconImg + reconImgTmp
        else
            reconImg = reconImg + reverse(reconImgTmp, dims = 2)'
        end

    end # Projection loop

    reconImg = reconImg / length(theta)
    return reconImg

end # endfunc

## Integrate function
"""
    integrate1D(p_v::Vector, pixelSize)
Calculates the integral image set,
for a given column vector and a pixel size
"""
function integrate1D(p_v::Vector,pixelSize)

    n_pixel = length(p_v)

    P_x = 0
    Ppj = zeros(n_pixel+1)

    for pj in 1:n_pixel

       P_x += p_v[pj] * pixelSize

       Ppj[pj+1] = P_x

    end

    return Ppj

end

"""
    mapp2y(x1,y1,x2,y2)
Maps detector or pixel boundaries onto y-axis,
for tube and detector rotation angles and detector/pixel boundaries
"""
function mapp2y(x1,y1,x2,y2)
    y = y1-x1*(y1-y2)/(x1-x2)
    return y
end

"""
    mapp2x(x1,y1,x2,y2)
Maps detector or pixel boundaries onto x-axis,
for tube and detector rotation angles and detector/pixel boundaries
"""
function mapp2x(x1,y1,x2,y2)
    x = -(y1)*(x1-x2)/(y1-y2)+x1
    return x
end

## Draw 2D
# Function to draw geometry
#=
function drawGeo(tubeX,tubeY,detX,detY)
    plot(tubeX,tubeY,markershape = :asterisk)
    plot!([detX[1],detX[end]],[detY[1],detY[end]])
end
=#


text:
using Clustering
using ParallelKMeans
using Plots
using BenchmarkTools
using TimerOutputs
using Random
using ProgressMeter

# Create a TimerOutput, this is the main type that keeps track of everything.
const to = TimerOutput()

Random.seed!(2020)
X = rand(60, 1_000_000);
# Timed assingments
global a = Float64[]
global b = Float64[]
global c = Float64[]

p = Progress(9, 10, "Computing clustering...")
@timeit to "Clustering" begin
    for i in 2:10
        @timeit to "$i clusters" push!(a, Clustering.kmeans(X, i, tol=1e-6, maxiter=300).totalcost)
        next!(p)
    end
end

p = Progress(9, 10, "Computing singlethreaded ParallelKMeans...")
@timeit to "PKMeans Singlethread" begin
    for i in 2:10
        @timeit to "$i clusters" push!(b, ParallelKMeans.kmeans(X, i, tol=1e-6, max_iters=300, verbose=false).totalcost)
        next!(p)
    end
end

p = Progress(9, 10, "Computing multithreaded ParallelKMeans...")
@timeit to "PKMeans Multithread" begin
    for i in 2:10
        @timeit to "$i clusters" push!(c, ParallelKMeans.kmeans(X, i, ParallelKMeans.MultiThread(), tol=1e-6, max_iters=300, verbose=false).totalcost)
        next!(p)
    end
end

plot(a, label="Clustering.jl")
plot!(b, label="Single-Thread ParallelKmeans")
plot!(c, label="Multi-Thread ParallelKmeans")

print(to)


text:
## test_separable.jl : test separable kernels for GRF generation

@testset "separable GRFs           " begin

    ## 1d Exponential ##
    cov = SeparableCovarianceFunction(Exponential(0.1))
    pts = range(0,stop = 1,length = 1001)
    grf = GaussianRandomField(cov,KarhunenLoeve(1000),pts)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test ndims(grf.cov) == 1
    @test isa(grf,GaussianRandomField{KarhunenLoeve{1000}})
    @test length(grf.pts[1]) == length(pts)
    @test length(sample(grf)) == length(pts)

    ## 2d Exponential ##
    e1 = Exponential(0.1)
    e2 = Exponential(0.01)
    cov = SeparableCovarianceFunction(e1,e2)
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(500),pts1,pts2)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{500}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## 2d Exponential ##
    e1 = Exponential(0.1)
    e2 = Exponential(0.01)
    cov = SeparableCovarianceFunction(e1,e2)
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(500),pts1,pts2)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{500}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    # matern/exponential 2d
    m1 = Matern(0.2,1.0)
    e1 = Exponential(0.01)
    cov = SeparableCovarianceFunction([m1,e1])
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(250),pts1,pts2,quad=EOLE())
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Matern)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{250}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## 3d separable Matern ##
    m1 = Matern(0.2,1.0)
    m2 = Matern(0.4,3.0)
    m3 = Matern(0.6,4.0)
    cov = SeparableCovarianceFunction(m1,m1,m2)
    pts1 = 0:0.1:1
    pts2 = 0:0.1:1
    pts3 = 0:0.1:1
    grf = GaussianRandomField(cov,KarhunenLoeve(100),pts1,pts2,pts3)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Matern)
    @test isa(grf.cov.cov[2],Matern)
    @test isa(grf.cov.cov[3],Matern)
    @test ndims(grf.cov) == 3
    @test isa(grf,GaussianRandomField{KarhunenLoeve{100}})
    @test length(grf.pts) == 3
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test length(grf.pts[3]) == length(pts3)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)
    @test size(sample(grf),3) == length(pts3)

end


text:
#!/usr/bin/env julia

using ArgParse
using BioStructures
using DataDeps
using PyCall

const shutil = pyimport("shutil")

"Free space in Gb at the given path."
free_space(path) = shutil.disk_usage(path)[3] / 1000_000_000

function check_space(free, min_space, database_name)
    if free < min_space
        throw(ErrorException("""
        Free space: $free Gb
        Please make sure to have at least $min_space Gb of free disk space
        before downloading the $database_name database.
        """))
    end
end

function parse_commandline()
    settings = ArgParseSettings(description = """
      This script set up the needed databases for PhyloSofS in the `--output`
      path. If the `--pdb`, `--uniclust` or `--pdb70` argument is used, the
      script is going to create a symbolic link to the indicated folder
      instead of downloading the database.

      It creates a `databases` folder in `--output` containing three folders:
      `pdb`, `uniclust` and `pdb70`
      """,)

    @add_arg_table settings begin
        "--output", "-o"
        help = "path where the database folder is going to be created."
        default = "."
        "--pdb"
        help = """
        path to an already existing local folder containing the entire pdb
        in mmCIF format.
        """
        default = ""
        "--uniclust"
        help = """
        path to an already existing local folder containing the uniclust
        database from the HH-suite databases.
        """
        default = ""
        "--uniclust_version"
        help = """
        Uniclust30 version to be downloaded: YYYY_MM
        """
        default = "2018_08"
        "--pdb70"
        help = """
        path to an already existing local folder containing the
        pdb70_from_mmcif database from the HH-suite databases.
        """
        default = ""
    end

    return parse_args(settings)
end

function main()
    execution_folder = pwd()
    parsed_args = parse_commandline()

    output_path = joinpath(abspath(parsed_args["output"]), "databases")
    mkpath(output_path)

    pdb = parsed_args["pdb"]
    uniclust = parsed_args["uniclust"]
    pdb70 = parsed_args["pdb70"]

    uniclust_version = parsed_args["uniclust_version"]

    pdb_path = joinpath(output_path, "pdb")
    uniclust_path = joinpath(output_path, "uniclust")
    pdb70_path = joinpath(output_path, "pdb70")

    free = free_space(output_path)

    @info "Setting up Uniclust30 (HH-suite database)"

    if uniclust == ""
        check_space(free, 14, "HH-suite Uniclust30")
        cd(output_path)
        download(
            string(
                "http://wwwuser.gwdg.de/~compbiol/uniclust/2018_08/uniclust30_",
                uniclust_version,
                ".tar.gz",
            ),
            "uniclust.tar.gz",
        )
        unpack("uniclust.tar.gz")
        mv("uniclust30_$uniclust_version", "uniclust")
        cd(execution_folder)
    else
        uniclust = abspath(uniclust)
        mkpath(uniclust_path)
        for file in readdir(uniclust)
            if filesize(file) > 0
                symlink(joinpath(uniclust, file), joinpath(uniclust_path, file))
            end
        end
    end

    cd(uniclust_path)
    to_erase = string("_", uniclust_version)
    for file in readdir()
        if occursin(to_erase, file)
            if filesize(file) > 0
                symlink(file, replace(file, to_erase => ""))
            end
        end
    end
    cd(execution_folder)

    @info "Setting up pdb70 (HH-suite database)"

    if pdb70 == ""
        check_space(free, 42, "HH-suite pdb70_from_mmcif")
        mkpath(pdb70_path)
        cd(pdb70_path)
        download(
            "http://wwwuser.gwdg.de/~compbiol/data/hhsuite/databases/hhsuite_dbs/pdb70_from_mmcif_latest.tar.gz",
            "pdb70.tar.gz",
        )
        unpack("pdb70.tar.gz")
        cd(execution_folder)
    else
        pdb70 = abspath(pdb70)
        mkpath(pdb70_path)
        for file in readdir(pdb70)
            if filesize(file) > 0
                symlink(joinpath(pdb70, file), joinpath(pdb70_path, file))
            end
        end
    end

    @info "Setting up PDB (mmCIF format)"

    if pdb == ""
        check_space(free, 157, "PDB")
        mkpath(pdb_path)
        downloadentirepdb(pdb_dir = pdb_path, file_format = MMCIF)
    else
        pdb = abspath(pdb)
        symlink(pdb, pdb_path)
    end

end

main()


text:
@testset "$(rpad("Tableau coefficients",80))" begin

    # test computation of Gauss-Legendre Runge-Kutta coefficients
    glrk2_tab1 = getCoefficientsGLRK(2)
    glrk2_tab2 = getCoefficientsGLRK(2, high_precision=false)
    @test glrk2_tab1.a ≈ glrk2_tab2.a atol=4eps()
    @test glrk2_tab1.b == glrk2_tab2.b
    @test glrk2_tab1.c ≈ glrk2_tab2.c atol=4eps()

    glrk3_tab1 = getCoefficientsGLRK(3)
    glrk3_tab2 = getCoefficientsGLRK(3, high_precision=false)
    @test glrk3_tab1.a ≈ glrk3_tab2.a atol=4eps()
    @test glrk3_tab1.b == glrk3_tab2.b
    @test glrk3_tab1.c ≈ glrk3_tab2.c atol=4eps()

    # test PGLRK coefficients
    pglrk2 = getCoefficientsPGLRK(2)
    @test typeof(pglrk2) <: CoefficientsPGLRK

end


text:
module FourierBasis

@inline ϕ(k, x; L=1)  = exp(im*2*pi*x*k/L) # L is the length of the interval
@inline ψ(k, h, x, y; Lx=1, Ly=1) = (-1)^(k+h)*ϕ(k, x; L = Lx)*ϕ(h, y; L= Ly)

purge(x, ϵ=2^-10) = abs(x)<ϵ ? 0.0 :  x
purge(x::Complex)= purge(real(x))+im*purge(imag(x))

make_dx(N) = [i/2^N for i in 0:2^N-1]

# converts [0,1,..., Nx,-Nx,...,-1] to [1,...,2*Nx+1]
unidimensional_index(i, Nx) = i>=0 ? i+1 : 2*Nx+2+i
inverse_unidimensional_index(i, Nx) = 1<=i<=Nx+1 ? i-1 : i-2*Nx-2

end


text:
using GeometricAlgebra

@testset "indexing" begin
	sigs = [sig for s ∈ [(1,1,1), (x=1,y=1,z=1)] for sig ∈ [s, OffsetSignature(s,1:3)]]
	@testset "integer indices with $sig" for sig ∈ sigs
		x, y, z = basis(sig)
		@test x[1] == 1
		@test x[2] == 0
		@test (x + 4y)[2] == 4
		@test (1 + 2*x*y)[2,1] == -2
	end

	labelled3d = (x=1,y=1,z=1)
	@testset "symbol indices with $sig" for sig ∈ [labelled3d, OffsetSignature(labelled3d,11:13)]
		x, y, z = basis(sig)
		@test x[:x] == 1
		@test y[:x,:y] == 0
		@test (x + 4y)[:y] == 4
		@test (1 + 2*x*y)[:y,:x] == -2
	end

	@testset "offset indices" begin
		t, x, y, z = basis(OffsetSignature((t=-1, x=1, y=1, z=1),0:3))
		@test t[0] == x[1] == 1
		@test (x + 2y)[2] == 2
		@test (1 + 2*x*y)[2,1] == -2
		@test (t*x*y*z)[0,1,2,3] == 1
	end

end

@testset "@basis" begin
	@basis x y z=-1
	@test x isa Blade{signature(x),1}
	@test x^2 == 1
	@test z^2 == -1

	@basisall x y
	@test xy == -yx
end


text:
# multigrid transfer operators
function intergrid_1d_interp(N, Minv)

    m = size(Minv)[1]
    ic = Int(ceil(m/2.0))
    sten = transpose(Minv) * Float64[(0.5)^(j-1)/factorial(j-1) for j=1:m]

    Ih = zeros(Float64, (N,N))
    Is = I + Ih
    for i=1:m
        k = i - ic
        vec = sten[i] * ones(N - abs(k))
        Ih += diagm(k => vec)
    end

    Iout = zeros(Float64, (2*N,N))
    Iout[1:2:end,:] = Is
    Iout[2:2:end,:] = Ih
    return Iout
end

function transfer_matrices(Nx0, Ny0, Mxinv, Myinv, P0, P1)

    Ix = intergrid_1d_interp(Nx0, Mxinv)
    Iy = intergrid_1d_interp(Ny0, Myinv)
    Ic = kron(sparse(Iy), sparse(Ix))
    Rc = transpose(Ic) / 4.0

    Restr = P0 * Rc * transpose(P1)
    Interp = P1 * Ic * transpose(P0)
    return Restr, Interp
end

# the idea behind this struct is that instances represent a matrix at multiple resolutions. Not only
# can the same matrix at different resolutions be represented by the same variable, but the algebra can be
# simplified in a more elegant way too.
struct MultiResMatrix

    mats :: Array{SparseMatrixCSC{Float64, Int32}, 1}
    lvl_by_row_size :: Dict{Int64, Int64}
    lvl_by_column_size :: Dict{Int64, Int64}

end

# no need for Julia to see this struct as a custom array, but we can overload the multiplication
# function so that multiplying from the left or the right automatically works. Both M*X and X*M will multiply
# X by the correct version of M.
function Base.:*(A::MultiResMatrix, X)
    lvl = A.lvl_by_row_size[size(X)[1]]
    return A.mats[lvl] * X
end
function Base.:*(X, A::MultiResMatrix)
    lvl = A.lvl_by_column_size[size(X)[2]]
    return X * A.mats[lvl]
end

function MultiResMatrix(A::SparseMatrixCSC, grd::Grid)

    As = SparseMatrixCSC[A]
    lvl_by_row_size = Dict(size(A)[2] => 1)
    lvl_by_column_size = Dict(size(A)[1] => 1)

    for l=1:length(grd.restrictions)
        Restr = grd.restrictions[l]
        Interp = grd.interpolations[end-(l-1)]
        Ap = Restr * As[end] * Interp
        append!(As, [Ap])

        lvl_by_row_size[size(Ap)[2]] = l+1
        lvl_by_column_size[size(Ap)[1]] = l+1
    end

    return MultiResMatrix(As, lvl_by_row_size, lvl_by_column_size)
end

function setup_jacobi_smoothing(A::SparseMatrixCSC, w::Float64, grd::Grid)

    A_mr = MultiResMatrix(A, grd)
    Ms = SparseMatrixCSC[]
    Ds = SparseMatrixCSC[]

    for a in A_mr.mats
        D = w*inv(Diagonal(a))
        M = I - D*a
        append!(Ms, [M])
        append!(Ds, [D])
    end

    M_mr = MultiResMatrix(Ms, A_mr.lvl_by_row_size, A_mr.lvl_by_column_size)
    D_mr = MultiResMatrix(Ds, A_mr.lvl_by_row_size, A_mr.lvl_by_column_size)

    return A_mr, M_mr, D_mr
end

function jacobi_smooth(n::Int64, M::MultiResMatrix, D::MultiResMatrix, x0::Vector{Float64}, r::Vector{Float64})

    x = x0[:]
    f = D * r
    for dummy=1:n
        x[:] = M * x + f
    end

    return x
end

function residual(A::MultiResMatrix, x::Vector{Float64}, b::Vector{Float64})
    return b - A * x
end

function vcycle(A::MultiResMatrix, x0::Vector{Float64}, b::Vector{Float64}, M::MultiResMatrix, D::MultiResMatrix, n::Int64, grd::Grid)

    # compute initial residual
    r = b - A*x0

    # relax on coarser grids
    for Restr in grd.restrictions
        x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
        r = Restr*(r - A*x)
    end

    # solve exactly on course grid, then interpolate
    x_crs = grd.interpolations[1] * (A.mats[end] \ r)
    r = grd.interpolations[1]*r - A * (x + x_crs)

    # relax on finer grids
    for Interp in grd.interpolations[2:end]
        x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
        r = Interp*(r - A*x)
    end

    # final smoothing
    x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
    return x0 + x
end

function pure_jacobi(A0::SparseMatrixCSC, x0::Vector{Float64}, b::Vector{Float64}, w::Float64, grd::Grid)

    # setup multiresolution matrices
    A, M, D = setup_jacobi_smoothing(A0, w, grd)

    # compute norm of b, initialize error
    bnorm = norm(b)
    err = 1.0
    x = x0[:]
    println()

    # interate until norm(residual) / bnorm < 10^-8
    n_iters = 0
    max_iters = 1000
    while (err > 1e-8) && (n_iters < max_iters)
        x[:] = jacobi_smooth(100, M, D, x, b) #XXX: This actually doesn't make sense, should just make jacobi_smooth depend on an initial guess.
        r = residual(A, x, b)
        err = norm(r) / bnorm
        n_iters += 1
        println(err)
    end

    return x
end

function multigrid_solve(A0::SparseMatrixCSC, x::Vector{Float64}, b::Vector{Float64}, w::Float64, grd::Grid)

    # setup multiresolution matrices
    A, M, D = setup_jacobi_smoothing(A0, w, grd)

    # compute norm of b, initialize error
    bnorm = norm(b)
    err = 1.0
    println()

    # interate until norm(residual) / bnorm < 10^-8
    n_iters = 0
    max_iters = 1000
    while (err > 1e-8) && (n_iters < max_iters)
        x = vcycle(A, x, b, M, D, 100, grd)
        err = norm(residual(A, x, b)) / bnorm
        n_iters += 1
        println(err)
    end

    return x
end


text:
struct LiftUpdater{T, TV<:AbstractVector{T}, TPD}
    pd::TPD
    ρ::TV
    g::TV
    ng_approx::TV
    n_j::Int
end

function (lu::LiftUpdater{T})() where T
    @unpack ρ, pd = lu
    @unpack x, x1, σ  = pd
    n_i = length(ρ)
    w = mapreduce((x_x1_σ)->((x_x1_σ[1] - x_x1_σ[2])^2 / (x_x1_σ[3]^2 - (x_x1_σ[1] - x_x1_σ[2])^2)), +, zip(x, x1, σ), init=T(0)) / 2
    return mapreduce(i->lu(w, i), or, 1:n_i, init=false)
end
function (lu::LiftUpdater{T})(w, i::Int) where T
    @unpack ρ, g, ng_approx, n_j, pd = lu
    @unpack x, x1, σ = pd
    gi, g_approxi = g[i], -ng_approx[i]
    lift = (gi > g_approxi)
    δi = (gi - g_approxi)/w
    ρi = ρ[i]
    ρ[i] = ifelse(lift, min(10ρi, T(1.1)*(ρi+δi)), ρi)
    return lift
end

struct LiftResetter{T, TV}
    ρ::TV
    ρmin::T
end
function (lr::LiftResetter)(k::Iteration)
    @unpack ρ = lr
    if k.i > 1
        map!(lr, ρ, 1:length(ρ))
    end
    return
end
function (lr::LiftResetter{T})(i::Int) where T
    @unpack ρmin, ρ = lr
    return max(T(0.1)*ρ[i], ρmin)
end