##=============================================================================
##
## Copyright (c) 2017-2019 Marco Colombo and Paul McKeigue
##
## kfold.hsstan() is based on code from https://github.com/stan-dev/rstanarm
## Portions copyright (C) 2015, 2016, 2017 Trustees of Columbia University
##
## 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/>.
##
##=============================================================================


#' Hierarchical shrinkage models
#'
#' Run the No-U-Turn Sampler (NUTS) as implemented in Stan to fit a hierarchical
#' shrinkage model.
#'
#' @param x Data frame containing outcome, covariates and penalized predictors.
#'        Continuous predictors and outcome variable should be standardized
#'        before fitting the models as priors assume them to have mean zero and
#'        unit variance.
#' @param covs.model Formula containing the unpenalized covariates.
#' @param penalized Names of the variables to be used as penalized predictors.
#'        Any variable that is already part of the `covs.model` formula will be
#'        penalized. If `NULL` or an empty vector, a model with only unpenalized
#'        covariates is fitted.
#' @param family Type of model fitted: either `gaussian()` for linear regression
#'        (default) or `binomial()` for logistic regression.
#' @param seed Optional integer defining the seed for the pseudo-random number
#'        generator.
#' @param qr Whether the thin QR decomposition should be used to decorrelate the
#'        predictors (`TRUE` by default). This is silently set to `FALSE` if
#'        there are more predictors than observations.
#' @param adapt.delta Target average proposal acceptance probability for
#'        adaptation, a value between 0.8 and 1 (excluded). If unspecified,
#'        it's set to 0.99 for hierarchical shrinkage models and to 0.95 for
#'        base models.
#' @param iter Total number of iterations in each chain, including warmup
#'        (2000 by default).
#' @param warmup Number of warmup iterations per chain (by default, half the
#'        total number of iterations).
#' @param scale.u Prior scale (standard deviation) for the unpenalized
#'        covariates.
#' @param regularized If `TRUE` (default), the regularized horseshoe prior
#'        is used as opposed to the original horseshoe prior.
#' @param nu Number of degrees of freedom of the half-Student-t prior on the
#'        local shrinkage parameters (by default, 1 if `regularized=TRUE`
#'        and 3 otherwise).
#' @param par.ratio Expected ratio of non-zero to zero coefficients (ignored
#'        if `regularized=FALSE`). The scale of the global shrinkage parameter
#'        corresponds to `par.ratio` divided by the square root of the number of
#'        observations; for linear regression only, it's further multiplied by
#'        the residual standard deviation `sigma`.
#' @param global.df Number of degrees of freedom for the global shrinkage
#'        parameter (ignored if `regularized=FALSE`). Larger values induce more
#'        shrinkage.
#' @param slab.scale Scale of the regularization parameter (ignored if
#'        `regularized=FALSE`).
#' @param slab.df Number of degrees of freedom of the regularization parameter
#'        (ignored if `regularized=FALSE`).
#' @param keep.hs.pars Whether the parameters for the horseshoe prior should be
#'        kept in the `stanfit` object returned (`FALSE` by default).
#' @param ... Further arguments passed to [rstan::sampling()],
#'        such as `chains` (4 by default), `cores` (the value of
#'        `options("mc.cores")` by default), `refresh` (`iter / 10` by default).
#'
#' @return
#' An object of class `hsstan` containing the following fields:
#' \item{stanfit}{an object of class `stanfit` containing the output
#'       produced by Stan, including posterior samples and diagnostic summaries.
#'       It can be manipulated using methods from the **rstan** package.}
#' \item{betas}{posterior means of the unpenalized and penalized regression
#'       parameters.}
#' \item{call}{the matched call.}
#' \item{data}{the dataset used in fitting the model.}
#' \item{model.terms}{a list of names for the outcome variable, the unpenalized
#'       covariates and the penalized predictors.}
#' \item{family}{the `family` object used.}
#' \item{hsstan.settings}{the optional settings used in the model.}
#'
#' @seealso
#' [kfold()] for cross-validating a fitted object.
#'
#' @examples
#' \dontshow{oldopts <- options(mc.cores=2)}
#' data(diabetes)
#'
#' # non-default settings for speed of the example
#' df <- diabetes[1:50, ]
#' hs.biom <- hsstan(df, Y ~ age + sex, penalized=colnames(df)[5:10],
#'                   chains=2, iter=250)
#' \dontshow{options(oldopts)}
#'
#' @importFrom stats gaussian
#' @export
hsstan <- function(x, covs.model, penalized=NULL, family=gaussian,
                   iter=2000, warmup=floor(iter / 2),
                   scale.u=2, regularized=TRUE, nu=ifelse(regularized, 1, 3),
                   par.ratio=0.05, global.df=1, slab.scale=2, slab.df=4,
                   qr=TRUE, seed=123, adapt.delta=NULL,
                   keep.hs.pars=FALSE, ...) {

    model.terms <- validate.model(covs.model, penalized)
    x <- validate.data(x, model.terms)
    y <- x[[model.terms$outcome]]
    family <- validate.family(family, y)
    regularized <- as.integer(regularized)
    validate.positive.scalar(iter, "iter", int=TRUE)
    validate.positive.scalar(warmup, "warmup", int=TRUE)
    if (iter <= warmup)
        stop("'warmup' must be smaller than 'iter'.", call.=FALSE)

    ## stop if options to be passed to rstan::sampling are not valid, as
    ## to work around rstan issue #681
    validate.rstan.args(...)

    ## parameter names not to include by default in the stanfit object
    hs.pars <- c("lambda", "tau", "z", "c2")
    if (keep.hs.pars)
        hs.pars <- NA

    ## retrieve the call and its actual argument values
    call <- match.call(expand.dots=TRUE)
    args <- c(as.list(environment()), list(...))
    for (nm in names(call)[-c(1:2)]) # exclude "" and "x"
        call[[nm]] <- args[[nm]]

    ## choose the model to be fitted
    model <- ifelse(length(penalized) == 0, "base0", "hs")
    if (family$family == "binomial") model <- paste0(model, "_logit")

    ## set or check adapt.delta
    if (is.null(adapt.delta)) {
        adapt.delta <- ifelse(grepl("hs", model), 0.99, 0.95)
    } else {
        validate.adapt.delta(adapt.delta)
    }

    ## create the design matrix
    X <- ordered.model.matrix(x, model.terms$unpenalized, model.terms$penalized)
    N <- nrow(X)
    P <- ncol(X)

    ## number of penalized and unpenalized columns in the design matrix
    K <- length(expand.terms(x, model.terms$penalized)[-1])
    U <- P - K

    ## thin QR decomposition
    if (P > N) qr <- FALSE
    if (qr) {
        qr.dec <- qr(X)
        Q.qr <- qr.Q(qr.dec)
        R.inv <- qr.solve(qr.dec, Q.qr) * sqrt(N - 1)
        Q.qr <- Q.qr * sqrt(N - 1)
    }

    ## global scale for regularized horseshoe prior
    global.scale <- if (regularized) par.ratio / sqrt(N) else 1

    ## core block
    {
        ## parameters not used by a model are ignored
        data.input <- list(X=if (qr) Q.qr else X, y=y, N=N,
                           P=P, U=U, scale_u=scale.u,
                           regularized=regularized, nu=nu,
                           global_scale=global.scale, global_df=global.df,
                           slab_scale=slab.scale, slab_df=slab.df)

        ## run the stan model
        samples <- rstan::sampling(stanmodels[[model]], data=data.input,
                                   iter=iter, warmup=warmup, seed=seed, ...,
                                   pars=hs.pars, include=keep.hs.pars,
                                   control=list(adapt_delta=adapt.delta))
        if (is.na(nrow(samples)))
            stop("rstan::sampling failed, see error message above.", call.=FALSE)

        ## assign proper names
        par.idx <- grep("^beta_[up]", names(samples))
        stopifnot(length(par.idx) == ncol(X))
        names(samples)[par.idx] <- colnames(X)

        if (qr) {
            pars <- grep("beta_", samples@sim$pars_oi, value=TRUE)
            stopifnot(pars[1] == "beta_u")
            beta.tilde <- rstan::extract(samples, pars=pars,
                                         inc_warmup=TRUE, permuted=FALSE)
            B <- apply(beta.tilde, 1:2, FUN=function(z) R.inv %*% z)
            chains <- ncol(beta.tilde)
            for (chain in 1:chains) {
                for (p in 1:P)
                    samples@sim$samples[[chain]][[par.idx[p]]] <- B[p, , chain]
            }
        }

        ## store the hierarchical shrinkage settings
        opts <- list(adapt.delta=adapt.delta, qr=qr, seed=seed, scale.u=scale.u)
        if (K > 0)
            opts <- c(opts, regularized=regularized, nu=nu, par.ratio=par.ratio,
                      global.scale=global.scale, global.df=global.df,
                      slab.scale=slab.scale, slab.df=slab.df)

        ## compute the posterior means of the regression coefficients
        betas <- list(unpenalized=colMeans(as.matrix(samples, pars="beta_u")),
                      penalized=tryCatch(colMeans(as.matrix(samples,
                                                            pars="beta_p")),
                                         error=function(e) NULL))
        obj <- list(stanfit=samples, betas=betas, call=call, data=x,
                    model.terms=model.terms, family=family, hsstan.settings=opts)
        class(obj) <- "hsstan"
    }

    return(obj)
}

#' K-fold cross-validation
#'
#' Perform K-fold cross-validation using the same settings used when fitting
#' the model on the whole data.
#'
#' @param x An object of class `hsstan`.
#' @param folds Integer vector with one element per observation indicating the
#'        cross-validation fold in which the observation should be withdrawn.
#' @param chains Number of Markov chains to run. By default this is set to 1,
#'        independently of the number of chains used for `x`.
#' @param store.fits Whether the fitted models for each fold should be stored
#'        in the returned object (`TRUE` by default).
#' @param cores Number of cores to use for parallelization (the value of
#'        `options("mc.cores")` by default). The cross-validation folds will
#'        be distributed to the available cores, and the Markov chains for each
#'        model will be run sequentially.
#' @param ... Further arguments passed to [rstan::sampling()].
#'
#' @return
#' An object with classes `kfold` and `loo` that has a similar structure as the
#' objects returned by [loo()] and [waic()] and is compatible with the
#' \code{\link[loo]{loo_compare}} function for
#' comparing models. The object contains the following fields:
#' \item{estimates}{a matrix containing point estimates and standard errors of
#'       the expected log pointwise predictive density ("elpd_kfold"),
#'       the effective number of parameters ("p_kfold", always `NA`) and the
#'       K-fold information criterion "kfoldic" (which is `-2 * elpd_kfold`,
#'       i.e., converted to the deviance scale).}
#' \item{pointwise}{a matrix containing the pointwise contributions of
#'       "elpd_kfold", "p_kfold" and "kfoldic".}
#' \item{fits}{a matrix with two columns and number of rows equal to the number
#'       of cross-validation folds. Column `fit` contains the fitted
#'       `hsstan` objects for each fold, and column `test.idx` contains
#'       the indices of the withdrawn observations for each fold. This is not
#'       present if `store.fits=FALSE`.}
#' \item{data}{the dataset used in fitting the model (before withdrawing
#'       observations). This is not present if `store.fits=FALSE`.}
#'
#' @examples
#' \donttest{
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' # only 2 folds for speed of example
#' folds <- rep(1:2, length.out=length(df$Y))
#' cv.biom <- kfold(hs.biom, folds=folds, cores=2)
#' }
#'
#' @importFrom loo kfold
#' @method kfold hsstan
#' @aliases kfold
#' @export kfold
#' @export
kfold.hsstan <- function(x, folds, chains=1, store.fits=TRUE,
                         cores=getOption("mc.cores", 1), ...) {
    data <- x$data
    N <- nrow(data)
    folds <- validate.folds(folds, N)
    num.folds <- max(folds)
    validate.positive.scalar(chains, "chains", int=TRUE)
    validate.rstan.args(...)

    ## collect the list of calls to be evaluated in parallel
    calls <- list()
    for (fold in 1:num.folds) {
        test.idx <- which(folds == fold)
        fit.call <- stats::update(object=x, x=data[-test.idx, , drop=FALSE],
                                  chains=chains, cores=1, refresh=0,
                                  open_progress=FALSE, evaluate=FALSE, ...)
        fit.call$x <- eval(fit.call$x)
        calls[[fold]] <- fit.call
    }

    ## evaluate the models
    message("Fitting ", num.folds, " models using ",
            min(cores, num.folds), " cores")
    par.fun <- function(fold) {
        fit <- eval(calls[[fold]])

        ## log pointwise predictive densities (pointwise test log-likelihood)
        lppd <- log_lik(fit, newdata=data[which(folds == fold), , drop=FALSE])
        return(list(lppd=lppd, fit=if (store.fits) fit else NULL))
    }
    if (.Platform$OS.type != "windows") {
        cv <- parallel::mclapply(X=1:num.folds, mc.cores=cores,
                                 mc.preschedule=FALSE, FUN=par.fun)
    } else { # windows
        cl <- parallel::makePSOCKcluster(cores)
        on.exit(parallel::stopCluster(cl))
        cv <- parallel::parLapply(X=1:num.folds, cl=cl, fun=par.fun)
    }

    ## expected log predictive densities
    elpds.unord <- unlist(lapply(cv, function(z) apply(z$lppd, 2, logMeanExp)))
    obs.idx <- unlist(lapply(1:num.folds, function(z) which(folds == z)))
    elpds <- elpds.unord[obs.idx]

    pointwise <- cbind(elpd_kfold=elpds, p_kfold=NA, kfoldic=-2 * elpds)
    estimates <- colSums(pointwise)
    se.est <- sqrt(N * apply(pointwise, 2, stats::var))
    out <- list(estimates=cbind(Estimate=estimates, SE=se.est),
                pointwise=pointwise)
    rownames(out$estimates) <- colnames(pointwise)
    if (store.fits) {
        fits <- array(list(), c(num.folds, 2), list(NULL, c("fit", "test.idx")))
        for (fold in 1:num.folds)
            fits[fold, ] <- list(fit=cv[[fold]][["fit"]],
                                 test.idx=which(folds == fold))
        out$fits <- fits
        out$data <- data
    }
    attr(out, "K") <- num.folds
    class(out) <- c("kfold", "loo")
    return(out)
}