##=============================================================================
##
## Copyright (c) 2017-2021 Marco Colombo and Paul McKeigue
##
## Parts of the code are based on https://github.com/jpiironen/rstan-varsel
## Portions copyright (c) 2015-2017 Juho Piironen
##
## 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/>.
##
##=============================================================================


#' Compute projections of full predictors on to subspace of predictors
#'
#' @param x Design matrix.
#' @param fit Matrix of fitted values for the full model.
#' @param sigma2 Residual variance (1 for logistic regression).
#' @param indproj Vector of indices of the columns of `x` that form the
#'        projection subspace.
#' @param is.logistic Set to `TRUE` for a logistic regression model.
#'
#' @importFrom stats binomial glm.fit quasibinomial
#' @keywords internal
lm_proj <- function(x, fit, sigma2, indproj, is.logistic) {

    P <- ncol(x)
    S <- ncol(fit)

    ## pick the Q columns of x that form the projection subspace
    xp <- x[, indproj, drop=FALSE] # matrix of dimension N x Q

    ## logistic regression model
    if (is.logistic) {

        ## compute the projection for each sample
        par.fun <- function(s) glm.fit(xp, fit[, s],
                                       family=quasibinomial())$coefficients
        if (.Platform$OS.type != "windows") {
            wp <- parallel::mclapply(X=1:S, mc.preschedule=TRUE, FUN=par.fun)
        } else { # windows
            cl <- parallel::makePSOCKcluster(getOption("mc.cores", 1))
            on.exit(parallel::stopCluster(cl))
            wp <- parallel::parLapply(X=1:S, cl=cl, fun=par.fun)
        }
        wp <- matrix(unlist(wp, use.names=FALSE), ncol=S)

        ## estimate the KL divergence between full and projected model
        fitp <- binomial()$linkinv(multiplyAB(xp, wp))
        kl <- 0.5 * mean(colMeans(binomial()$dev.resids(fit, fitp, 1)))
        sigma2p <- 1
    }

    ## linear regression model
    else {
        ## solve the projection equations
        wp <- solve(crossprod(xp, xp), multiplyAtB(xp, fit)) # Q x S matrix

        ## fit of the projected model
        fitp <- multiplyAB(xp, wp)

        ## estimate the KL divergence between full and projected model
        sigma2p <- sigma2 + colMeans((fit - fitp)^2)
        kl <- mean(0.5 * log(sigma2p / sigma2))
    }

    ## reshape wp so that it has same dimension (P x S) as t(x), and zeros for
    ## those variables that are not included in the projected model
    wp.all <- matrix(0, P, S)
    wp.all[indproj, ] <- wp
    return(list(w=wp.all, sigma2=sigma2p, fit=fitp, kl=kl))
}

#' Next variable to enter the current submodel
#'
#' Return the index of the variable that should be added to the current model
#' according to the smallest KL-divergence (linear regression) or the largest
#' score test (logistic regression).
#'
#' @param x Design matrix.
#' @param sigma2 Residual variance (1 for logistic regression).
#' @param fit Matrix of fitted values for the full model.
#' @param fitp Matrix of fitted values for the projected model.
#' @param chosen Vector of indices of the columns of `x` in the current submodel.
#' @param is.logistic Set to `TRUE` for a logistic regression model.
#'
#' @keywords internal
choose.next <- function(x, sigma2, fit, fitp, chosen, is.logistic) {
    notchosen <- setdiff(1:ncol(x), chosen)
    par.fun <- function(idx) {
        lm_proj(x, fit, sigma2, sort(c(chosen, notchosen[idx])), FALSE)$kl
    }
    if (!is.logistic) {
        if (.Platform$OS.type != "windows") {
            kl <- parallel::mclapply(X=1:length(notchosen), mc.preschedule=TRUE,
                                     FUN=par.fun)
        } else { # windows
            cl <- parallel::makePSOCKcluster(getOption("mc.cores", 1))
            on.exit(parallel::stopCluster(cl))
            kl <- parallel::parLapply(X=1:length(notchosen), cl=cl, fun=par.fun)
        }
        idx.selected <- which.min(unlist(kl))
        return(notchosen[idx.selected])
    }

    ## score test
    yminusexp <- fit - fitp
    dinv.link <- fitp * (1 - fitp)
    U <- colMeans(multiplyAtB(yminusexp, x[, notchosen]))
    V <- colMeans(multiplyAtB(dinv.link, x[, notchosen]^2))
    idx.selected <- which.max(U^2 / V)
    return(notchosen[idx.selected])
}

#' Compute the fit of a submodel
#'
#' @param x Design matrix.
#' @param sigma2 Residual variance (1 for logistic regression).
#' @param mu Matrix of fitted values for the full model.
#' @param chosen Vector of indices of the columns of `x` in the current submodel.
#' @param xt Design matrix for test data.
#' @param yt Outcome variable for test data.
#' @param logistic Set to `TRUE` for a logistic regression model.
#'
#' @keywords internal
fit.submodel <- function(x, sigma2, mu, chosen, xt, yt, logistic) {
    ## projected parameters
    submodel <- lm_proj(x, mu, sigma2, chosen, logistic)
    eta <- multiplyAB(xt, submodel$w)
    return(c(submodel, elpd=elpd(yt, eta, submodel$sigma2, logistic)))
}

#' Expected log predictive density
#'
#' @param yt Outcome variable.
#' @param eta Matrix of linear predictors, with as many rows as the number of
#'        observations.
#' @param sigma2 Residual variance (1 for logistic regression).
#' @param logistic Set to `TRUE` for a logistic regression model.
#'
#' @noRd
elpd <- function(yt, eta, sigma2, logistic) {
    if (logistic)
        pd <- stats::dbinom(yt, 1, binomial()$linkinv(eta), log=TRUE)
    else {
        sigma <- sqrt(sigma2)
        pd <- t(cbind(sapply(1:nrow(eta),
                             function(z) stats::dnorm(yt[z], eta[z, ],
                                                      sigma, log=TRUE))))
    }
    sum(rowMeans(pd))
}

#' Forward selection minimizing KL-divergence in projection
#'
#' @param obj Object of class `hsstan`.
#' @param max.iters Maximum number of iterations (number of predictors selected)
#'        after which the selection procedure should stop.
#' @param start.from Vector of variable names to be used in the starting
#'        submodel. If `NULL` (default), selection starts from the set of
#'        unpenalized covariates if the model contains penalized predictors,
#'        otherwise selection starts from the intercept-only model.
#' @param out.csv If not `NULL`, the name of a CSV file to save the
#'        output to.
#'
#' @return
#' A data frame of class `projsel` where each row corresponds to a
#' forward-selected submodel that contains all variables listed up to that row.
#' Attribute `start.from` reports the predictors in the initial model.
#' The data frame contains the following columns:
#' \item{var}{names of the variables selected.}
#' \item{kl}{KL-divergence from the full model to the submodel.}
#' \item{rel.kl.null}{relative explanatory power of predictors starting from the
#'       intercept-only model.}
#' \item{rel.kl}{relative explanatory power of predictors starting from the
#'       initial submodel.}
#' \item{elpd}{the expected log predictive density of the submodels.}
#' \item{delta.elpd}{the difference in elpd from the full model.}
#'
#' @examples
#' \donttest{
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' \dontshow{oldopts <- options(mc.cores=2)}
#' # continued from ?hsstan
#' sel <- projsel(hs.biom, max.iters=3)
#' plot(sel)
#' \dontshow{options(oldopts)}
#' }
#'
#' @importFrom utils write.csv
#' @export
projsel <- function(obj, max.iters=30, start.from=NULL,
                    out.csv=NULL) {
    validate.hsstan(obj)
    validate.samples(obj)
    validate.nonnegative.scalar(max.iters, "max.iters", int=TRUE)

    x <- xt <- validate.newdata(obj, obj$data)
    yt <- obj$data[[obj$model.terms$outcome]]
    is.logistic <- is.logistic(obj)
    sigma2 <- if (is.logistic) 1 else as.matrix(obj$stanfit, pars="sigma")^2

    ## set of variables in the initial submodel
    vsf <- validate.start.from(obj, start.from)
    start.from <- vsf$start.from
    chosen <- start.idx <- vsf$idx

    ## number of model parameters (including intercept)
    P <- sum(sapply(obj$betas, length))

    ## number of iterations to run
    K <- min(P - length(start.idx), max.iters)

    message(sprintf("%58s  %8s %11s", "Model", "KL", "ELPD"))
    report.iter <- function(msg, kl, elpd)
        message(sprintf("%58s  %8.5f  %8.5f", substr(msg, 1, 55), kl, elpd))

    ## fitted values for the full model (N x S)
    fit <- t(posterior_linpred(obj, transform=TRUE))

    ## intercept only model
    sub <- fit.submodel(x, sigma2, fit, 1, xt, yt, is.logistic)
    kl.elpd <- cbind(sub$kl, sub$elpd)
    report.iter("Intercept only", sub$kl, sub$elpd)

    ## initial submodel with the set of chosen covariates
    if (length(start.idx) > 1) {
        sub <- fit.submodel(x, sigma2, fit, chosen, xt, yt, is.logistic)
        kl.elpd <- rbind(kl.elpd, c(sub$kl, sub$elpd))
        report.iter("Initial submodel", sub$kl, sub$elpd)
    }

    ## add variables one at a time
    for (iter in seq_len(K)) {
        sel.idx <- choose.next(x, sigma2, fit, sub$fit, chosen, is.logistic)
        chosen <- c(chosen, sel.idx)

        ## evaluate current submodel according to projected parameters
        sub <- fit.submodel(x, sigma2, fit, chosen, xt, yt, is.logistic)
        kl.elpd <- rbind(kl.elpd, c(sub$kl, sub$elpd))
        report.iter(colnames(x)[sel.idx], sub$kl, sub$elpd)

        ## check if the current submodel is fully saturated
        if (sub$kl < 1e-8 && iter < K) {
            message("Fully saturated model reached, selection terminated.")
            break()
        }
    }

    ## evaluate the full model if selection stopped before reaching it
    full.elpd <- if (length(chosen) == P)
        sub$elpd else elpd(yt, t(posterior_linpred(obj)), sigma2, is.logistic)

    kl <- kl.elpd[, 1]
    rel.kl <- c(NA, if (length(kl) > 1) 1 - kl[-1] / kl[2])
    if (length(rel.kl) == 2 && is.nan(rel.kl[2]))
        rel.kl[2] <- 1

    res <- data.frame(var=c("Intercept only",
                            if (length(start.idx) > 1) "Initial submodel",
                            colnames(x)[setdiff(chosen, start.idx)]),
                      kl=kl,
                      rel.kl.null=1 - kl / kl[1],
                      rel.kl=rel.kl,
                      elpd=kl.elpd[, 2],
                      delta.elpd=kl.elpd[, 2] - full.elpd,
                      stringsAsFactors=FALSE, row.names=NULL)
    attr(res, "start.from") <- start.from
    if (!is.null(out.csv))
        write.csv(file=out.csv, res, row.names=FALSE)

    class(res) <- c("projsel", "data.frame")
    return(res)
}

#' Plot of relative explanatory power of predictors
#'
#' The function plots the relative explanatory power of each predictor in order
#' of selection. The relative explanatory power of predictors is computed
#' according to the KL divergence from the full model to each submodel, scaled
#' in such a way that the baseline model (either the null model or the model
#' containing only unpenalized covariates) is at 0, while the full model is at 1.
#'
#' @param x A data frame created by [projsel()].
#' @param title Title of the plot. If `NULL`, no title is displayed.
#' @param max.points Maximum number of predictors to be plotted. If `NULL`
#'        (default) or 0, all points are plotted.
#' @param max.labels Maximum number of predictors to be labelled. If `NULL`
#'        (default), all predictor labels present in `x` are displayed, which
#'        may result in overprinting.
#' @param from.covariates Whether the plotting should start from the unpenalized
#'        covariates (`TRUE` by default). If set to `FALSE`, the plot includes a
#'        point for the null (intercept-only) model.
#' @param font.size Size of the textual elements (labels and axes).
#' @param hadj,vadj Horizontal and vertical adjustment for the labels.
#' @param ... Currently ignored.
#'
#' @return
#' A **ggplot2** object showing the relative incremental contribution of each
#' predictor starting from the initial set of unpenalized covariates.
#'
#' @import ggplot2
#' @method plot projsel
#' @export
plot.projsel <- function(x, title=NULL, max.points=NULL, max.labels=NULL,
                         from.covariates=TRUE,
                         font.size=12, hadj=0.05, vadj=0, ...) {

    ## prepare the set of points to plot
    sel <- x
    if (!is.null(max.points) && max.points > 0)
        sel <- utils::head(sel, n=max.points + 2)
    if (!is.null(max.labels))
        sel$var[-c(1:(max.labels + 2))] <- ""
    if (from.covariates)
        sel <- sel[-1, ]
    labs <- sel$var

    ## convert from points to millimetres
    geom.text.size <- font.size * 25.4 / 72

    x <- seq(nrow(sel)) - 1
    text_idx <- x < mean(x) | x - floor(x / 2) * 2 == 1
    rel <- if (from.covariates) sel$rel.kl else sel$rel.kl.null
    p <- ggplot(data=sel, aes(x=x, y=rel, label=labs)) +
      coord_cartesian(ylim=range(c(0, 1))) +
      geom_line() + geom_point(size=geom.text.size / 3) +
      geom_text(aes(x=x + ifelse(text_idx, hadj, -hadj),
                    y=rel + ifelse(text_idx, -vadj, vadj)),
                size=geom.text.size,
                hjust=ifelse(text_idx, "left", "right")) +
      xlab("Number of biomarkers") +
      ylab("Relative explanatory power") +
      theme(text=element_text(size=font.size))

    if (!is.null(title))
        p <- p + ggtitle(title)

    return(p)
}