
abstract type AbstractAggregateTools end

const ElIndices = Vector{Vector{T}} where {T<:Integer}
const VibIndices = Vector{Vector{Vector{Int64}}} where {T<:Integer}
# TODO: be more specific for T
const Indices = Vector{Vector{Vector{T} where T}}
const IndicesMap = Vector{Vector{T}} where {T<:Integer}
const FCfactorsT = Matrix{T} where {T<:AbstractFloat}
const FCproductT = Matrix{T} where {T<:AbstractFloat}

"""
    vibrationalIndices(aggCore)

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

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
"""
function electronicIndices(aggCore::AggregateCore)::ElIndices
    # TODO: add typeofs
    vibInds = Array{Array{Int64,1},1}(undef, 0)
    molCount = length(aggCore.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

"""
    vibrationalIndices(aggCore)

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

"""
function vibrationalIndices(aggCore::AggregateCore)::VibIndices
    # TODO: add typeofs
    NvibMols = getNvib(aggCore)
    molLen = length(aggCore.molecules)
    NvibIndMols = Array{Array{Array{Int64,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{Int64,1},1},1}(undef, 0)
    for molInd in molInds
        aggInd = Array{Array{Int64,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

"""
    getIndices(aggCore)

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

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
"""
function getIndices(aggCore::AggregateCore)::Indices
    vibInds = vibrationalIndices(aggCore)
    vibIndsLen = length(vibInds)
    elInds = electronicIndices(aggCore)
    elIndsLen = length(elInds)
    indices = Array{Array{Array{Int64,1} where Int64,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

"""
getIndicesMap(aggCore)

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 getIndicesMap(aggCore::AggregateCore, aggIndices::Indices)::IndicesMap
    aggIndLen = length(aggIndices)
    indicesMap = Array{Array{Int64,1},1}(undef, 0)
    elLen = length(aggCore.molecules) + 1
    for el_i = 1:elLen
        push!(indicesMap, Array{Int64,1}(undef, 0))
    end

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

"""
    getFranckCondonFactors(aggCore, aggIndices)
    getFranckCondonFactors(aggCore)

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

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
"""
function getFranckCondonFactors(aggCore::AggregateCore, aggIndices::Indices)::FCfactorsT
    aggIndLen = length(aggIndices)
    molLen = length(aggCore.molecules)
    FC = zeros(Float64, (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::Float64
                for mi = 1:molLen
                    mol = aggCore.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(aggCore::AggregateCore)::FCfactorsT =
    (aggIndices = getIndices(aggCore);
    getFranckCondonFactors(aggCore, aggIndices))

# TODO: resurrect sparse version
#=
"""
    getFranckCondonFactorsSparse(agg, aggIndices)
    getFranckCondonFactorsSparse(agg)

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

"""
function getFranckCondonFactorsSparse(
    agg::C,
    aggIndices::Any
) where {C<:AbstractAggregateCore}
    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::C
) where {C<:AbstractAggregateCore} = getFranckCondonFactorsSparse(
    agg::C,
    nothing
)
=#

"""
    getFCProd(agg, FCFact, aggIndices, vibindices)

Get product of Franck-Condon factors. This way the trace over bath will be faster

"""
function getFCproduct(
    aggCore::AggregateCore,
    indices::Indices,
    indicesMap::IndicesMap,
    FCfactors::FCfactorsT,
)::FCproductT
    elLen = length(aggCore.molecules)
    aggIndLen = length(indices)
    vibLen = length(indicesMap[2])
    elLen += 1
    FCproduct = zeros(eltype(FCfactors), aggIndLen, aggIndLen)

    for I = 1:aggIndLen
        elind1, vibind1 = indices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)

        for J = 1:aggIndLen
            elind2, vibind2 = indices[J]
            elOrder2 = OpenQuantumSystems.elIndOrder(elind2)
            #=
            K1 = vibindices[elOrder1][1]
            K2 = vibindices[elOrder1][end]
            L1 = vibindices[elOrder2][1]
            L2 = vibindices[elOrder2][end]
            FCProd[I, J] = sum(FCFact[K1:K2, I] .* FCFact[J, L1:L2])
            =#

            for m = 1:vibLen
                # according to quantarhei, trace_over_vibrations()
                K = indicesMap[1][m]
                L = indicesMap[1][m]
                FCproduct[I, J] += FCfactors[K, I] * FCfactors[J, L]
            end
        end
    end
    return FCproduct
end

struct AggregateTools <: AbstractAggregateTools
    elIndices::ElIndices
    vibIndices::VibIndices
    indices::Indices
    indicesMap::IndicesMap
    FCfactors::FCfactorsT
    FCproduct::FCproductT
    basisSystem::GenericBasis
    basisBath::GenericBasis
    basis::GenericBasis
    bSystemSize::Int64
    bBathSize::Int64
    bSize::Int64
    function AggregateTools(
        elIndices::ElIndices,
        vibIndices::VibIndices,
        indices::Indices,
        indicesMap::IndicesMap,
        FCfactors::FCfactorsT,
        FCproduct::FCproductT,
        basisSystem::GenericBasis,
        basisBath::GenericBasis,
        basis::GenericBasis,
        bSystemSize::Int64,
        bBathSize::Int64,
        bSize::Int64,
    )::AggregateTools
        new(
            elIndices,
            vibIndices,
            indices,
            indicesMap,
            FCfactors,
            FCproduct,
            basisSystem,
            basisBath,
            basis,
            bSystemSize,
            bBathSize,
            bSize,
        )
    end
end

function AggregateTools(aggCore::AggregateCore)::AggregateTools
    elIndices = electronicIndices(aggCore)
    vibIndices = vibrationalIndices(aggCore)
    indices = getIndices(aggCore)
    indicesMap = getIndicesMap(aggCore, indices)
    FCfactors = getFranckCondonFactors(aggCore, indices)
    FCproduct = getFCproduct(aggCore, indices, indicesMap, FCfactors)

    bSystemSize = length(elIndices)
    bBathSize = length(vibIndices)
    bSize = length(indices)

    basisSystem = GenericBasis([bSystemSize])
    basisBath = GenericBasis([bBathSize])
    basis = GenericBasis([bSize])


    return AggregateTools(
        elIndices,
        vibIndices,
        indices,
        indicesMap,
        FCfactors,
        FCproduct,
        basisSystem,
        basisBath,
        basis,
        bSystemSize,
        bBathSize,
        bSize,
    )
end

Base.:(==)(x::AggregateTools, y::AggregateTools) =
    x.elIndices == y.elIndices &&
    x.vibIndices == y.vibIndices &&
    x.indices == y.indices &&
    x.FCfactors == y.FCfactors &&
    x.FCproduct == y.FCproduct &&
    x.bSystemSize == y.bSystemSize &&
    x.bBathSize == y.bBathSize &&
    x.bSize == y.bSize &&
    x.basisSystem == y.basisSystem &&
    x.basisBath == y.basisBath &&
    x.basis == y.basis

"""
    take_el_part(A, a, b, indicesMap)

Take electric part specified by electric indices `a` and `b` from the A (type of Array). 

"""
function take_el_part(A::Array, a, b, indicesMap)
    a1 = indicesMap[a][1]
    a2 = indicesMap[a][end]
    b1 = indicesMap[b][1]
    b2 = indicesMap[b][end]

    return A[a1:a2, b1:b2]
end
