compounds.jl

using StatsFuns.RFunctions: betarand, gammarand
using LogProbs

function categorical_sample(tokens, weights)
    x = rand() * sum(weights)
    cum_weights = zero(eltype(weights))
    for (t, w) in zip(tokens, weights)
        cum_weights += w
        if cum_weights > x
            return t
        end
    end
end

categorical_sample(d::Dict) = categorical_sample(keys(d), values(d))
categorical_sample(v::Vector) = categorical_sample(1:length(v), v)

mutable struct DirCat{T, C}
    counts :: Dict{T, C}
end

DirCat(support, priors) = DirCat(Dict(x => p for (x,p) in zip(support, priors)))
support(dc::DirCat) = keys(dc.counts)

function sample(dc::DirCat)
    weights = [gammarand(c, 1) for c in values(dc.counts)]
    categorical_sample(keys(dc.counts), weights)
end

function logscore(dc::DirCat, obs)
    LogProb(lbeta(sum(values(dc.counts)), 1) - lbeta(dc.counts[obs], 1))
end

function add_obs!(dc::DirCat, obs)
    dc.counts[obs] += 1
end

function rm_obs!(dc::DirCat, obs)
    dc.counts[obs] -= 1
end

mutable struct DirMul{T, C}
    counts :: Dict{T, C}
end

DirMul(support) = DirMul(Dict(x => 1.0 for x in support))
support(dm::DirMul) = keys(dm.counts)

function sample(dm::DirMul, n)
    weights = [gammarand(c, 1) for c in values(dm.counts)]
    d = Dict(k=>0 for k in keys(dm.counts))
    for i in 1:n
        d[categorical_sample(keys(dm.counts), weights)] += 1
    end
    d
end

function logscore(dm::DirMul, obs::Associative)
    n = sum(values(obs))
    LogProb(
        log(n) + lbeta(sum(values(dm.counts)), n) -
        sum(
            log(obs[x]) + lbeta(dm.counts[x], obs[x])
            for x in keys(dm.counts) if obs[x] > 0
        )
    )
end

function add_obs!(dm::DirMul, obs::Associative)
    for x in keys(obs)
        dm.counts[x] += obs[x]
    end
end

function rm_obs!(dm::DirMul, obs::Associative)
    for x in keys(obs)
        dm.counts[x] -= obs[x]
    end
end