##=============================================================================
##
## Copyright (c) 2019 Marco Colombo and Paul McKeigue
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##
##=============================================================================


#' Summary for the fitted model
#'
#' @param object An object of class `hsstan`.
#' @param pars Vector of parameter names to be extracted. If `NULL`
#'        (default) then this refers to the set of predictors used in the
#'        model.
#' @param prob Width of the posterior intervals (0.95, by default).
#' @param digits Number of decimal places to be reported (2 by default).
#' @param sort Column name used to sort the results according to the absolute
#'        value of the column. If `NULL` (default) or the column name cannot
#'        be found, no sorting occurs.
#' @param decreasing Whether the results should be sorted in decreasing order
#'        when a valid column name is specified in `sort` (`TRUE` by default).
#' @param max.rows Maximum number of rows to be returned. If `NULL`
#'        (default) or 0, all results are returned.
#' @param ... Currently ignored.
#'
#' @return
#' A matrix with summaries from the posterior distribution of the parameters
#' of interest.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' summary(hs.biom)
#'
#' @importMethodsFrom rstan summary
#' @method summary hsstan
#' @export
summary.hsstan <- function(object, pars=NULL, prob=0.95, digits=2,
                           sort=NULL, decreasing=TRUE, max.rows=NULL, ...) {
    validate.samples(object)
    validate.probability(prob)
    pars <- get.pars(object, pars)
    low <- (1 - prob) / 2
    upp <- 1 - low
    summary <- summary(object$stanfit, pars=pars, probs=c(low, upp))$summary
    summary[, "n_eff"] <- round(summary[, "n_eff"], 0)
    if (!is.null(sort) && sort %in% colnames(summary)) {
        ord <- order(abs(summary[, sort]), decreasing=decreasing)
        summary <- summary[ord, ]
    }
    if (!is.null(max.rows)) {
        if (max.rows > 0 && max.rows < nrow(summary))
            summary <- summary[1:max.rows, , drop=FALSE]
    }
    round(summary[, -match("se_mean", colnames(summary)), drop=FALSE],
                digits)
}

#' Print a summary for the fitted model
#'
#' @param x An object of class `hsstan`.
#' @param ... Further arguments to [summary()].
#'
#' @export
print.hsstan <- function(x, ...) {
    print(summary(x, ...))
}

#' Posterior summary
#'
#' Produce a summary of the posterior samples for the quantities of interest.
#'
#' @param x An object containing or representing posterior samples. If `x`
#'        is a matrix, it should have size `S` by `Q`, where `S`
#'        is the number of posterior samples, and `Q` is the number of
#'        quantities of interest.
#' @param prob Width of the posterior intervals (0.95, by default).
#' @param ... Further arguments passed to or from other methods.
#'
#' @return
#' A matrix with columns containing mean, standard deviation and posterior
#' intervals for the given quantities.
#'
#' @seealso
#' [summary()] to produce summaries of `hsstan` objects that include the number
#' of effective samples and the split-Rhat diagnostic.
#'
#' @export
posterior_summary <- function(x, prob, ...) {
    UseMethod("posterior_summary")
}

#' @rdname posterior_summary
#' @export
posterior_summary.default <- function(x, prob=0.95, ...) {
    validate.probability(prob)
    if (NCOL(x) == 1)
        x <- as.matrix(x)
    t(apply(x, 2, vector.summary, prob))
}

#' @rdname posterior_summary
#' @param pars Vector of parameter names to be extracted. If `NULL`
#'        (default) then this refers to the set of predictors used in the
#'        model.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' posterior_summary(hs.biom)
#'
#' @export
posterior_summary.hsstan <- function(x, prob=0.95, pars=NULL, ...) {
    validate.samples(x)
    pars <- get.pars(x, pars)
    posterior_summary(as.matrix(x$stanfit, pars=pars), prob, ...)
}

#' Sampler statistics
#'
#' Report statistics on the parameters used in the sampler, the sampler
#' behaviour and the sampling time.
#'
#' @param object An object of class `hsstan`.
#'
#' @return
#' A matrix with `C + 1` rows, where `C` is the number of Markov
#' chains, reporting average acceptance probability, average stepsize, number
#' of divergent transitions, maximum tree depth, total number of gradient
#' evaluations, warmup and sampling times in seconds.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' sampler.stats(hs.biom)
#'
#' @export
sampler.stats <- function(object) {
    validate.hsstan(object)
    validate.samples(object)
    sp <- rstan::get_sampler_params(object$stanfit, inc_warmup=FALSE)
    accept.stat <- sapply(sp, function(x) mean(x[, "accept_stat__"]))
    stepsize <- sapply(sp, function(x) mean(x[, "stepsize__"]))
    divergences <- sapply(sp, function(x) sum(x[, "divergent__"]))
    treedepth <- sapply(sp, function(x) max(x[, "treedepth__"]))
    gradients <- sapply(sp, function(z) sum(z[, "n_leapfrog__"]))
    et <- round(rstan::get_elapsed_time(object$stanfit), 2)
    res <- cbind(accept.stat, stepsize, divergences, treedepth, gradients, et)
    avg <- colMeans(res)
    tot <- colSums(res)
    round(rbind(res, all=c(avg[1:2], tot[3], max(res[, 4]), tot[5:7])), 4)
}

#' Number of posterior samples
#'
#' Extracts the number of posterior samples stored in a fitted model.
#'
#' @param object An object of class `hsstan`.
#' @param ... Currently ignored.
#'
#' @return
#' The total number of posterior samples across the chains after discarding
#' the warmup iterations.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' nsamples(hs.biom)
#'
#' @importFrom rstantools nsamples
#' @method nsamples hsstan
#' @aliases nsamples
#' @export nsamples
#' @export
nsamples.hsstan <- function(object, ...) {
    validate.samples(object)
    with(object$stanfit@sim, sum(n_save - warmup2))
}