# """
#     Chains(chain::RevJumpChain, burnin)
#
# Obtain and MCMCChains object, with summary and diagnostic statistics.
# """
# MCMCChains.Chains(c::RevJumpChain, burnin=1000) = Chains(c.trace, burnin)
#
# function MCMCChains.Chains(trace::DataFrame, burnin=1000)
#     df = select(trace, Not([:wgds]))
#     for col in names(df)
#         df[!,col] .= coalesce.(df[!,col], 0.)
#     end
#     @assert size(df)[1] > burnin "# MCMC generations < $burnin (burnin)"
#     X = reshape(Matrix(df), (size(df)...,1))[burnin+1:end, 2:end, :]
#     return Chains(X, [string(x) for x in names(df)][2:end])
# end

"""
    getwgdtrace(chain)

Summarize all WGD models from an rjMCMC trace. This provdes the data to
evaluate retention rates for WGD models etc. Returns a dict of dicts with
data frames (which is a horrible data structure, I know) structured as
(branch_1 => (1 WGD => trace, 2 WGDs => trace, ...), branch_2 => (), ...).
"""
function getwgdtrace(chain::RevJumpChain)
    tr = Dict{Int64,Dict{Int64,DataFrame}}()
    for (gen, d) in enumerate(chain.trace[!,:wgds])
        for (k,v) in d
            n = length(v)
            if !haskey(tr, k)
                tr[k] = Dict{Int64,DataFrame}()
            end
            if !haskey(tr[k], n)
                tr[k][n] = DataFrame(
                    :gen =>Int64[],
                    [Symbol("q$i")=>Float64[] for i=1:n]...,
                    [Symbol("t$i")=>Float64[] for i=1:n]...)
            end
            sort!(v)
            t = [x[1] for x in v]
            q = [x[2] for x in v]
            push!(tr[k][n], vcat(gen, q, t))
        end
    end
    tr
end

"""
    bayesfactors(trace::DataFrame, model::DLWGD, p::Float64)

Compute Bayes Factors for all branch WGD configurations. Returns a data frame
that is more or less self-explanatory.
"""
function bayesfactors(chain; burnin::Int64=1000)
    @unpack trace, model, prior = chain
    trace_ = trace[burnin+1:end,:]
    df = bayesfactors(trace_, model, prior.Ï€K)
    show_bayesfactors(df)
    df
end

# Prior probabilities for # of WGDs on each branch
function kbranch_prior(k, t, T, prior::Union{UpperBoundedGeometric,DiscreteUniform})
    kmax = prior.b
    f = t/T
    sum([binomial(i, k)*f^k*(1. - f)^(i-k)*pdf(prior, i) for i=k:kmax])
end

# closed form under Geometric prior
function kbranch_prior(k, t, T, prior::Geometric)
    p = prior.p
    a = t/T
    b = 1. - t/T
    q = 1. - p
    p * (a*q)^k * (1. - b*q)^(-k-1)
end

# closed form under Poisson prior
kbranch_prior(k, t, T, prior::Poisson) = pdf(Poisson(prior.Î»*t/T), k)

function bayesfactors(trace::DataFrame, model::DLWGD, p)
    T = treelength(model)
    df = DataFrame()
    for (i,n) in sort(model.nodes)
        if isawgd(n); break; end
        if !("k$i" in names(trace)); continue; end
        fmap = freqmap(trace[!,"k$i"])
        kmax = maximum(keys(fmap))
        te = parentdist(n, nonwgdparent(n.p))
        ps = [haskey(fmap, i) ? fmap[i] : 0. for i in 0:kmax]
        Ï€s = kbranch_prior.(0:kmax, te, T, p)
        cl = Beluga.clade(model, n)
        df_ = DataFrame(:branch=>i,
            :clade=>repeat([cl], kmax+1),
            :k=>0:kmax, :eval1=>"", :eval2=>"",
            :p1=>ps, :p0=>[NaN ; ps[1:end-1]],
            :Ï€1=>Ï€s, :Ï€0=>[NaN ; Ï€s[1:end-1]],
            :t=>te, :tfrac=>te/T)
        df = vcat(df, df_)
    end
    df[!,:p2] = by(df, :branch, :p1=>cumsum0)[!,:p1_cumsum0]
    df[!,:Ï€2] = by(df, :branch, :Ï€1=>cumsum0)[!,:Ï€1_cumsum0]
    # compute BFs
    df[!,:K1] = (df[!,:p1] ./ df[!,:p0]) ./ (df[!,:Ï€1] ./ df[!,:Ï€0])
    df[!,:K2] = (df[!,:p2] ./ df[!,:p0]) ./ (df[!,:Ï€2] ./ df[!,:Ï€0])
    df[!,:log10K1] = log10.(df[!,:K1])
    df[!,:log10K2] = log10.(df[!,:K2])
    for x in [0.5, 1., 2.]
        df[!,:eval1][df[!,:log10K1] .> x] .*= "*"
        df[!,:eval2][df[!,:log10K2] .> x] .*= "*"
    end
    df
end

cumsum0(x) = 1 .- vcat(0., cumsum(x)...)[1:end-1]

# 'pretty' printing of Bayes factor computations
function show_bayesfactors(df::DataFrame)
    for d in groupby(df, :branch)
        @printf "ðŸŒ² %2d: " d[1,:branch]
        println("(", join(string.(d[1,:clade]), ","), ")")
        for (i,r) in enumerate(eachrow(d))
            isnan(r[:p0]) ? continue : nothing
            @printf "[%2d vs. %2d] " i-1 i-2
            @printf "K = (%.2f/%.2f) Ã· (%.2f/%.2f) " r[:p1] r[:p0] r[:Ï€1] r[:Ï€0]
            @printf "= %8.3f [logâ‚�â‚€(K) = %8.3f] %s\n" r[:K1] r[:log10K1] r[:eval1]
        end
        println()
        for (i,r) in enumerate(eachrow(d))
            isnan(r[:p0]) ? continue : nothing
            @printf "[â‰¥%1d vs. %2d] " i-1 i-2
            @printf "K = (%.2f/%.2f) Ã· (%.2f/%.2f) " r[:p2] r[:p0] r[:Ï€2] r[:Ï€0]
            @printf "= %8.3f [logâ‚�â‚€(K) = %8.3f] %s\n" r[:K2] r[:log10K2] r[:eval2]
        end
        println("_"^78)
    end
end

"""
    posteriorE!(chain)

Compute E[Xi|Xparent=1,Î»,Î¼] for the joint posterior; i.e. the expected number
of lineages at node i under the linear birth-death process given that there was
one lineage at the parent of i, for each sample from the posterior. This can
give an idea of gene family expansion/contraction.
"""
function posteriorE!(chain)
    @unpack trace, model = chain
    for i=2:ne(model)+1
        t = parentdist(model[i], nonwgdparent(model[i].p))
        l = trace[!,Symbol("Î»$i")]
        m = trace[!,Symbol("Î¼$i")]
        E = exp.(t*(l .- m))
        V = E .* (E .- 1.) .* (l .+ m) ./ (l .- m)
        trace[!,Symbol("E$i")] .= E
        trace[!,Symbol("V$i")] .= V
    end
end

"""
    posteriorÎ£!(chain)

Sample the covariance matrix of the bivariate process *post hoc* from the
posterior under the Inverse Wishart prior. Based on Lartillot & Poujol 2010.
"""
function posteriorÎ£!(chain)
    @unpack model, prior = chain
    @unpack Î¨ = prior
    chain.trace[!,:varÎ»] .= NaN
    chain.trace[!,:varÎ¼] .= NaN
    chain.trace[!,:cov]  .= NaN
    chain.trace[!,:cor]  .= NaN
    for row in eachrow(chain.trace)
        m = model(row)
        @unpack A, q, n = scattermat(m, prior)
        Î£ = rand(InverseWishart(q + n, Î¨ + A))
        row[:varÎ»] = Î£[1,1]
        row[:varÎ¼] = Î£[2,2]
        row[:cov]  = Î£[1,2]
        row[:cor]  = Î£[1,2]/sqrt(Î£[1,1]*Î£[2,2])
    end
end

"""
    PostPredSim(chain, data::DataFrame, n::Int64)

Perform posterior predictive simulations.
"""
mutable struct PostPredSim
    datastats::DataFrame
    ppstats  ::DataFrame
end

"""
    pppvalues(pps::PostPredSim)

Compute posterior predictive p-values based on the posterior predictive
distribution and the observed sumary statistics (see e.g. Gelman et al. 2013).
"""
function pppvalues(pps::PostPredSim)
    @unpack datastats, ppstats = pps
    nrep = size(ppstats)[1]
    d = Dict{Symbol,Float64}()
    for n in names(ppstats)
        x = datastats[1,n]
        p = length(ppstats[ppstats[!,n] .> x,n])/nrep
        d[n] = p > 0.5 ? 1. - p : p
    end
    df = stack(DataFrame(d))
    df[!,:eval] .= ""
    df[df[!,:value] .< 0.05 ,:eval] .= "*"
    df[df[!,:value] .< 0.01 ,:eval] .= "**"
    df[df[!,:value] .< 0.001,:eval] .= "***"
    df
end

PostPredSim(chain::RevJumpChain, data, n; kwargs...) =
    PostPredSim(chain.trace, chain.model, data, n; kwargs...)

function PostPredSim(trace::DataFrame, model::DLWGD, data::DataFrame, n::Int64;
        burnin=1000, lstats=[mean, std, entr], gstats=lstats)
    N = size(data)[1]
    ngen = size(trace)[1]
    @assert burnin < ngen "burnin ($burnin) < generations ($ngen)"
    d1 = pp_simulate(trace, model, N, n,
        burnin=burnin, lstats=lstats, gstats=gstats)
    d2 = sstats(data, lstats=lstats, gstats=gstats)
    PostPredSim(d2, d1)
end

function pp_simulate(trace::DataFrame, model::DLWGD, N::Int64, n::Int64;
        burnin=1000, lstats=[mean, std, entr], gstats=lstats)
    clades = [clade(model, x) for x in model[1].c]
    trace = trace[burnin+1:end,:]
    # NOTE: could be in parallel!
    sdf = [sstats(rand(model(rand(trace)), N),
        lstats=lstats, gstats=gstats) for i=1:n]
    vcat(sdf...)
end

function sstats(df::DataFrame;lstats=[mean, std, entr], gstats=lstats)
    ldf = describe(df, [Pair(x...) for x in zip(Symbol.(lstats), lstats)]...)
    gdf = flatten(ldf, :variable)
    for s in gstats
        gdf[!,Symbol(s)] .= s(vcat(Matrix(df)...))
    end
    gdf
end

entr(x) = entropy(values(freqmap(x)))

function flatten(df::DataFrame, on::Symbol)
    d = Dict()
    for r in eachrow(df)
        for n in names(r)
            n != on ? d[Symbol("$(r[on])_$n")] = r[n] : continue
        end
    end
    DataFrame(d)
end
