ContingencyTables.jl

# Contingency Tables
# ==================

"""
A `ContingencyTable` is a multidimensional array. It stores the contingency matrix, its
marginal values and total. The type also has an internal and private temporal array and an
alphabet object. It's a parametric type, taking three ordered parameters:

  - `T` : The element type of the multidimensional array.
  - `N` : It's the dimension of the array and should be an `Int`.
  - `A` : This should be a type, subtype of `ResidueAlphabet`, i.e.: `UngappedAlphabet`,
    `GappedAlphabet` or `ReducedAlphabet`.

A `ContingencyTable` can be created from an alphabet if all the parameters are given.
Otherwise, you need to give a type, a number (`Val`) and an alphabet. You can also create a
`ContingencyTable` using a matrix and a alphabet. For example:

```julia
ContingencyTable{Float64,2,UngappedAlphabet}(UngappedAlphabet())
ContingencyTable(Float64, Val{2}, UngappedAlphabet())
ContingencyTable(zeros(Float64, 20, 20), UngappedAlphabet())
```
"""
mutable struct ContingencyTable{T,N,A} <: AbstractArray{T,N}
    alphabet::A
    temporal::Array{T,N}
    table::NamedArray{T,N,Array{T,N},NTuple{N,OrderedDict{String,Int}}}
    marginals::NamedArray{T,2,Array{T,2},NTuple{2,OrderedDict{String,Int}}}
    total::T
end

# Probability and Frequencies
# ----------------------

"""
A `Probabilities` object wraps a `ContingencyTable` storing probabilities. It doesn't
perform any check. If the total isn't one, you must use `normalize` or `normalize!`on the
`ContingencyTable` before wrapping it to make the sum of the probabilities equal to one.
"""
mutable struct Probabilities{T,N,A} <: AbstractArray{T,N}
    table::ContingencyTable{T,N,A}
end

"""
A `Frequencies` object wraps a `ContingencyTable` storing counts/frequencies.
"""
mutable struct Frequencies{T,N,A} <: AbstractArray{T,N}
    table::ContingencyTable{T,N,A}
end

# DEPRECATED: Deprecation warning for Counts type
mutable struct Counts{T,N,A} <: AbstractArray{T,N}
    table::ContingencyTable{T,N,A}
end
function Counts{T,N,A}(table::ContingencyTable{T,N,A}) where {T,N,A}
    Base.depwarn(
        "The `Counts` type is deprecated. Please use `Frequencies` instead.",
        :Counts,
        force = true,
    )
    Frequencies{T,N,A}(table)
end

# Getters

"""
`getcontingencytable` allows to access the wrapped `ContingencyTable` in a `Probabilities`
or `Frequencies` object.
"""
@inline getcontingencytable(p::Probabilities{T,N,A}) where {T,N,A} = p.table
@inline getcontingencytable(n::Frequencies{T,N,A}) where {T,N,A} = n.table

for f in
    (:getalphabet, :gettable, :getmarginals, :gettotal, :gettablearray, :getmarginalsarray)
    @eval $(f)(p::Probabilities{T,N,A}) where {T,N,A} = $(f)(getcontingencytable(p))
    @eval $(f)(n::Frequencies{T,N,A}) where {T,N,A} = $(f)(getcontingencytable(n))
end

# AbstractArray

for f in (:size, :getindex, :setindex!)
    @eval Base.$(f)(p::Probabilities{T,N,A}, args...) where {T,N,A} =
        $(f)(getcontingencytable(p), args...)
    @eval Base.$(f)(n::Frequencies{T,N,A}, args...) where {T,N,A} =
        $(f)(getcontingencytable(n), args...)
end

# Getters
# -------

"""
`getalphabet` allows to access the stored alphabet object.
"""
@inline getalphabet(table::ContingencyTable) = table.alphabet

"""
`gettable` allows to access the table (`NamedArray`).
"""
@inline gettable(table::ContingencyTable) = table.table

"""
`getmarginals` allows to access the array with the marginal values (`NamedArray`).
"""
@inline getmarginals(table::ContingencyTable) = table.marginals

"""
`gettotal` allows to access the stored total value.
"""
@inline gettotal(table::ContingencyTable) = table.total

Base.sum(table::ContingencyTable) = gettotal(table)

"""
`gettablearray` allows to access the table (`Array` without names).
"""
@inline gettablearray(table::ContingencyTable) = getarray(table.table)

"""
`getmarginalsarray` allows to access the array with the marginal values
(`Array` without names).
"""
@inline getmarginalsarray(table::ContingencyTable) = getarray(table.marginals)

# Cartesian (helper functions)
# ----------------------------

"""
`_marginal(1,:A,:i,:value)` generates the expression: `A[i_1, 1] += value`
"""
function _marginal(N::Int, marginal::Symbol, index::Symbol, value::Symbol)
    aexprs = [
        Expr(:escape, Expr(:(+=), Expr(:ref, marginal, Symbol(index, '_', i), i), value)) for i = 1:N
    ]
    Expr(:block, aexprs...)
end

macro _marginal(N, marginal, index, value)
    _marginal(N, marginal, index, value)
end

"""
`_test_index(1, i, continue)` generates the expression: `i_1 >= 22 && continue`
"""
function _test_index(N::Int, index::Symbol, expr::Expr)
    aexprs = Expr[]
    for i = 1:N
        push!(
            aexprs,
            Expr(:escape, Expr(:&&, Expr(:call, :>=, Symbol(index, "_", i), 22), expr)),
        )
    end
    Expr(:block, aexprs...)
end

macro _test_index(N, index, expr)
    _test_index(N, index, expr)
end

# AbstractArray Interface
# -----------------------

Base.size(table::ContingencyTable) = size(table.table)

Base.getindex(table::ContingencyTable, i...) = getindex(table.table, i...)

@generated function Base.getindex(table::ContingencyTable, I::Residue...)
    N = length(I)
    quote
        alphabet = getalphabet(table)
        matrix = getarray(gettable(table))
        # index_1 = alphabet[I[1]]
        # index_2 ...
        @nextract $N index d -> alphabet[I[d]]
        # index_1 >= 22 error("There is a Residue outside the alphabet")
        # index_2 ...
        @_test_index $N index error("There is a Residue outside the alphabet")
        # getindex(matrix, index_1, index_2...
        @nref $N matrix index
    end
end

function Base.setindex!(table::ContingencyTable, value, i...)
    setindex!(table.table, value, i...)
    update_marginals!(table)
end

@generated function Base.setindex!(table::ContingencyTable, value, I::Residue...)
    N = length(I)
    quote
        alphabet = getalphabet(table)
        matrix = getarray(gettable(table))
        @nextract $N index d -> alphabet[I[d]]
        @_test_index $N index error("There is a Residue outside the alphabet")
        @nref($N, matrix, index) = value
        update_marginals!(table)
    end
end

# Similar
# -------

Base.similar(table::ContingencyTable{T,N,A}) where {T,N,A} =
    ContingencyTable(T, Val{N}, table.alphabet)

function Base.similar(table::ContingencyTable{T,N,A}, ::Type{S}) where {T,S,N,A}
    ContingencyTable(S, Val{N}, table.alphabet)
end

# Show
# ====

function Base.show(io::IO, ::MIME"text/plain", table::ContingencyTable{T,N,A}) where {T,N,A}
    println(io, typeof(table), " : ")
    print(io, "\ntable : ")
    show(io, MIME"text/plain"(), gettable(table))
    if length(size(table)) != 1
        print(io, "\n\nmarginals : ")
        show(io, MIME"text/plain"(), getmarginals(table))
    end
    print(io, "\n\ntotal : $(gettotal(table))")
end

function Base.show(
    io::IO,
    ::MIME"text/plain",
    table::Union{Probabilities{T,N,A},Frequencies{T,N,A}},
) where {T,N,A}
    print(io, typeof(table), " wrapping a ")
    show(io, MIME"text/plain"(), getcontingencytable(table))
end

# Creation
# --------

@generated function (::Type{ContingencyTable{T,N,A}})(alphabet::A) where {T,N,A}
    @assert N > 0 "The dimension should be a natural number"
    quote
        n = length(alphabet)
        # residue_names = names(alphabet)
        # namedict = OrderedDict{String,Int}(residue_names[i] => i for i in 1:n)
        namedict = getnamedict(alphabet)
        dimtable = @ntuple $N k -> n # (n, n, ...)
        dimtemporal = @ntuple $N k -> 22 # (22, 22, ...)
        table = NamedArray(
            zeros($T, dimtable),
            @ntuple($N, k -> namedict), # (namedict, namedict, ...)
            @ntuple($N, k -> "Dim_k")
        ) # ("Dim_1", "Dim_2", ...)
        marginals = NamedArray(
            zeros($T, n, $N),
            # OrderedDict{String,Int}("Dim_$i" => i for i in 1:N)
            (namedict, OrderedDict{String,Int}(@ntuple $N k -> "Dim_k" => k)),
            ("Residue", "Dim"),
        )
        ContingencyTable{T,N,A}(alphabet, zeros(T, dimtemporal), table, marginals, zero(T))
    end
end

function ContingencyTable(::Type{T}, ::Type{Val{N}}, alphabet::A) where {T,A,N}
    ContingencyTable{T,N,A}(alphabet)
end

function ContingencyTable(matrix::AbstractArray{T,N}, alphabet::A) where {T,N,A}
    n = length(alphabet)
    @assert size(matrix) == ((n for i = 1:N)...,) "Matrix size doesn't match alphabet length"
    table = ContingencyTable(T, Val{N}, alphabet)
    getarray(table.table)[:] = matrix
    _update_marginals!(table)
    _update_total!(table)
end

# Update!
# -------

# Update the table, marginal and total using temporal

@generated function _update_table!(table::ContingencyTable{T,N,A}) where {T,N,A}
    if A <: ReducedAlphabet
        quote
            temporal = table.temporal::Array{T,N}
            alphabet = getalphabet(table)::A
            freqtable = getarray(table.table)::Array{T,N}
            @inbounds @nloops $N i temporal begin
                @_test_index $N i continue
                @nextract $N a d -> alphabet[i_d]
                @_test_index $N a continue
                @nref($N, freqtable, a) += @nref($N, temporal, i)
            end
            table
        end
    elseif (A === UngappedAlphabet) || (A === GappedAlphabet)
        quote
            temporal = table.temporal::Array{T,N}
            freqtable = getarray(table.table)::Array{T,N}
            @inbounds @nloops $N i freqtable begin
                @nref($N, freqtable, i) += @nref($N, temporal, i)
            end
            table
        end
    end
end

@generated function _update_marginals!(table::ContingencyTable{T,N,A}) where {T,N,A}
    quote
        freqtable = getarray(table.table)::Array{T,N}
        marginal = getarray(table.marginals)::Matrix{T}
        @inbounds @nloops $N i freqtable begin
            value = @nref $N freqtable i
            @_marginal($N, marginal, i, value)
        end
        table
    end
end

function _update_total!(table::ContingencyTable{T,N,A}) where {T,N,A}
    marginals = getarray(table.marginals)::Matrix{T}
    n = size(marginals, 1)
    total = zero(T)
    @inbounds @simd for i = 1:n
        total += marginals[i] # marginals[i,1]
    end
    table.total = total
    table
end

"""
`update_marginals!` updates the marginal and total values using the table.
"""
function update_marginals!(table::ContingencyTable{T,N,A}) where {T,N,A}
    fill!(getarray(table.marginals)::Matrix{T}, zero(T))
    _update_marginals!(table)
    _update_total!(table)
    table
end

function _update!(table::ContingencyTable{T,N,A}) where {T,N,A}
    _update_table!(table)
    update_marginals!(table)
end

function _cleanup_table!(table::ContingencyTable{T,N,A}) where {T,N,A}
    fill!(getarray(table.table)::Array{T,N}, zero(T))
end

function _cleanup_temporal!(table::ContingencyTable{T,N,A}) where {T,N,A}
    fill!(table.temporal, zero(T))
end

"""
`cleanup!` fills the temporal, table and marginals arrays with zeros.
It also sets total to zero.
"""
function cleanup!(table::ContingencyTable{T,N,A}) where {T,N,A}
    _cleanup_temporal!(table)
    _cleanup_table!(table)
    fill!(getarray(table.marginals)::Matrix{T}, zero(T))
    table.total = zero(T)
    table
end

# Fill
# ====

function Base.fill!(table::ContingencyTable{T,N,A}, value::T) where {T,N,A}
    fill!(getarray(table.table)::Array{T,N}, value)
    update_marginals!(table)
end

Base.fill!(table::ContingencyTable{T,N,A}, p::AdditiveSmoothing{T}) where {T,N,A} =
    fill!(table, p.λ)

@inline Base.fill!(table::ContingencyTable{T,N,A}, p::NoPseudocount) where {T,N,A} = table

# Apply pseudocount
# =================

# This is faster than array[:] += value
function _sum!(matrix::NamedArray, value)
    matrix_array = getarray(matrix)
    @inbounds for i in eachindex(matrix_array)
        matrix_array[i] += value
    end
    matrix
end

"""
It adds the `pseudocount` value to the table cells.
"""
function apply_pseudocount!(table::ContingencyTable{T,N,A}, pseudocount::T) where {T,N,A}
    _sum!(table.table, pseudocount)
    update_marginals!(table)
end

function apply_pseudocount!(
    table::ContingencyTable{T,N,A},
    p::AdditiveSmoothing{T},
) where {T,N,A}
    apply_pseudocount!(table, p.λ)
end

@inline apply_pseudocount!(table::ContingencyTable, p::NoPseudocount) = table

# Normalize
# =========

# This is faster than array[:] /= value
function _div!(matrix::NamedArray, value)
    matrix_array = getarray(matrix)
    @inbounds for i in eachindex(matrix_array)
        matrix_array[i] /= value
    end
    matrix
end

"""
`normalize!` makes the sum of the frequencies to be one, in place.
"""
function LinearAlgebra.normalize!(table::ContingencyTable{T,N,A}) where {T,N,A}
    if table.total != one(T)
        _div!(table.table, table.total)
        update_marginals!(table)
    end
    table
end

"""
`normalize` returns another table where the sum of the frequencies is one.
"""
LinearAlgebra.normalize(table::ContingencyTable{T,N,A}) where {T,N,A} =
    normalize!(deepcopy(table))

# Delete Dimensions
# =================

function _list_without_dimensions(len::Int, output_len::Int, dimensions::Int...)
    ndim = length(dimensions)
    @assert (len - ndim) == output_len "$output_len should be = $(len-ndim)"
    index_list = Array{Int}(undef, output_len)
    j = 1
    @inbounds for i = 1:len
        if !(i in dimensions)
            index_list[j] = i
            j += 1
        end
    end
    index_list
end

"""
`delete_dimensions!(out::ContingencyTable, in::ContingencyTable, dimensions::Int...)`

This function fills a ContingencyTable with the counts/probabilities on `in` after the
deletion of `dimensions`. i.e. This is useful for getting Pxy from Pxyz.
"""
function delete_dimensions!(
    output::ContingencyTable{T,S,A},
    input::ContingencyTable{T,N,A},
    dimensions::Int...,
) where {T,N,S,A}
    output_marginals = getarray(output.marginals)
    output_table = getarray(output.table)
    input_marginals = getarray(input.marginals)
    input_table = getarray(input.table)
    output_marginals[:] = input_marginals[:, _list_without_dimensions(N, S, dimensions...)]
    output_table[:] = sum(input_table, dims = dimensions)
    output.total = input.total
    output
end

"""
`delete_dimensions(in::ContingencyTable, dimensions::Int...)`

This function creates a ContingencyTable with the counts/probabilities on `in` after the
deletion of `dimensions`. i.e. This is useful for getting Pxy from Pxyz.
"""
function delete_dimensions(
    input::ContingencyTable{T,N,A},
    dims::Vararg{Int,I},
) where {T,N,A,I}
    delete_dimensions!(ContingencyTable(T, Val{N - I}, input.alphabet), input, dims...)
end

for tp in (:Probabilities, :Frequencies)
    @eval begin
        function delete_dimensions!(
            output::$(tp){T,S,A},
            input::$(tp){T,N,A},
            dimensions::Int...,
        ) where {T,N,S,A}
            delete_dimensions!(
                getcontingencytable(output),
                getcontingencytable(input),
                dimensions...,
            )
            output
        end

        function delete_dimensions(input::$(tp){T,N,A}, dims::Vararg{Int,I}) where {T,N,A,I}
            output = delete_dimensions(getcontingencytable(input), dims...)
            $(tp)(output)
        end
    end
end