##=============================================================================

##

## Copyright (c) 2019 Marco Colombo

##

## 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/>.

##

##=============================================================================

#' Validate an hsstan object

#'

#' Check that the object has been created by [hsstan()].

#'

#' @param obj An object to be checked.

#'

#' @return

#' Throws an error if the object is not an `hsstan` object.

#'

#' @noRd

validate.hsstan <- function(obj) {

    if (!inherits(obj, "hsstan")) {

        stop("Not an object of class 'hsstan'.")

    }

}

#' Validate the posterior samples

#'

#' Check that the object contains valid posterior samples in the

#' `stanfit` field.

#'

#' @param obj An object of class `hsstan`.

#'

#' @return

#' Throws an error if the object does not contain posterior samples.

#'

#' @noRd

validate.samples <- function(obj) {

    if (!inherits(obj$stanfit, "stanfit")) {

        stop("No valid posterior samples stored in the 'hsstan' object.")

    }

}

#' Validate new data

#'

#' Check that the new data contains all variables used in the model and no

#' missing values, and generate the corresponding model matrix.

#'

#' @param obj Object of class `hsstan`.

#' @param newdata Optional data frame containing the variables used in the

#'        model. If `NULL`, the model matrix used when fitting the model

#'        is returned.

#'

#' @return

#' A design matrix corresponding to the variables used in the model.

#'

#' @noRd

validate.newdata <- function(obj, newdata) {

    if (is.null(newdata))

        newdata <- obj$data

    else if (!inherits(newdata, c("data.frame", "matrix")))

        stop("'newdata' must be a data frame or a matrix.")

    if (nrow(newdata) == 0 || ncol(newdata) == 0)

        stop("'newdata' contains no rows or no columns.")

    ## only check for NAs in the variables used in the model

    vars <- with(obj$model.terms, c(outcome, unpenalized, penalized))

    newdata <- newdata[, colnames(newdata) %in% vars, drop=FALSE]

    if (any(is.na(newdata)))

        stop("'newdata' contains missing values.")

    ## this adds the intercept column back

    ordered.model.matrix(as.data.frame(newdata),

                         obj$model.terms$unpenalized,

                         obj$model.terms$penalized)

}

#' Validate a model formula

#'

#' Check that the formula that specifies a model contains all required elements.

#'

#' @param model Formula to be checked.

#' @param penalized Vector of names for the penalized predictors.

#'

#' @return

#' A list containing the formula representing the covariates model, the name of

#' the outcome variable, the names of the upenalized and penalized predictors.

#'

#' @importFrom stats as.formula terms

#' @noRd

validate.model <- function(model, penalized) {

    if (is.character(model) && length(model) > 1)

        stop("Model formula specified incorrectly.")

    model <- as.formula(model)

    tt <- terms(model)

    if (attr(tt, "response") == 0)

        stop("No outcome variable specified in the model.")

    if (attr(tt, "intercept") == 0)

        stop("Models with no intercept are not supported.")

    if (length(penalized) > 0 && !is.character(penalized))

        stop("'penalized' must be a character vector.")

    if (any(grepl("[:*]", penalized)))

        stop("Interaction terms in penalized predictors are not supported.")

    penalized <- setdiff(unique(trimws(penalized)), "")

    return(list(outcome=as.character(model)[2],

                unpenalized=setdiff(attr(tt, "term.labels"), penalized),

                penalized=penalized))

}

#' Validate the model data

#'

#' Check if the model data can be used with the given model formula and

#' penalized predictors.

#'

#' @param x An object to be checked.

#' @param model Validated model formula.

#'

#' @return

#' A data frame containing the model data. A factor or logical outcome variable

#' is replaced by its numeric equivalent.

#'

#' @noRd

validate.data <- function(x, model) {

    if (!inherits(x, c("data.frame", "matrix")))

        stop("'x' must be a data frame or a matrix.")

    x <- as.data.frame(x)

    validate.variables(x, model$outcome)

    validate.variables(x, c(model$unpenalized, model$penalized))

    x[[model$outcome]] <- validate.outcome(x[[model$outcome]])

    return(x)

}

#' Validate variables

#'

#' Check that the required variables are in the dataset.

#'

#' @param x Data frame containing the variables of interest.

#' @param variables Vector of variable names.

#'

#' @return

#' Throws if variables are not present in the dataset or contain missing values.

#'

#' @noRd

validate.variables <- function(x, variables) {

    ## unpack interaction terms

    variables <- unique(unlist(strsplit(as.character(variables), ":")))

    if (length(variables) == 0)

        stop("No predictors present in the model.")

    var.match <- match(variables, colnames(x))

    if (anyNA(var.match))

        stop(collapse(variables[is.na(var.match)]), " not present in 'x'.")

    if (anyNA(x[, variables]))

        stop("Model variables contain missing values.")

}

#' Validate the outcome variable

#'

#' Check that the outcome variable can be converted to a valid numerical

#' vector.

#'

#' @param y Outcome vector to be checked.

#'

#' @return

#' A numeric vector.

#'

#' @noRd

validate.outcome <- function(y) {

    if (is.factor(y)) {

        if (nlevels(y) != 2)

            stop("A factor outcome variable can only have two levels.")

        y <- as.integer(y) - 1

    }

    if (!(is.numeric(y) || is.logical(y)))

        stop("Outcome variable of invalid type.")

    return(as.numeric(y))

}

#' Validate the family argument

#'

#' Ensure that the family argument has been specified correctly.

#' This is inspired by code in \code{\link{glm}}.

#'

#' @param family Family argument to test.

#' @param y Outcome variable.

#'

#' @return

#' A valid family. The function throws an error if the family argument cannot

#' be used.

#'

#' @importFrom methods is

#' @noRd

validate.family <- function(family, y) {

    if (missing(family))

        stop("Argument of 'family' is missing.")

    if (is.character(family))

        tryCatch(

            family <- get(family, mode="function", envir=parent.frame(2)),

                          error=function(e)

                              stop("'", family, "' is not a valid family.")

        )

    if (is.function(family))

        family <- family()

    if (!is(family, "family"))

        stop("Argument of 'family' is not a valid family.")

    if (!family$family %in% c("gaussian", "binomial"))

        stop("Only 'gaussian' and 'binomial' are supported families.")

    if (family$family == "binomial") {

        if (length(table(y)) != 2)

            stop("Outcome variable must contain two classes with family=binomial.")

        if (!is.factor(y) && any(y < 0 | y > 1))

            stop("Outcome variable must contain 0-1 values with family=binomial.")

    }

    return(family)

}

#' Validate a vector of indices

#'

#' @param x Vector to be checked.

#' @param N Maximum valid index.

#' @param name Name of the vector to report in error messages.

#' @param throw.duplicates Whether the function should throw if the vector

#'        contains duplicate elements (`TRUE` by default).

#'

#' @return

#' Throws an error if the given vector is not an integer vector or contains

#' missing, out of bounds or duplicate indices (if `throw.duplicates` is `TRUE`).

#'

#' @noRd

validate.indices <- function(x, N, name, throw.duplicates=TRUE) {

    if (anyNA(x))

        stop("'", name, "' contains missing values.")

    if (!is.numeric(x) || NCOL(x) > 1 || any(x != as.integer(x)))

        stop("'", name, "' must be an integer vector.")

    if (length(x) < 2)

        stop("'", name, "' must contain at least two elements.")

    if (any(x < 1 | x > N))

        stop("'", name, "' contains out of bounds indices.")

    if (throw.duplicates && any(duplicated(x)))

        stop("'", name, "' contains duplicate indices.")

}

#' Validate the cross-validation folds

#'

#' @param folds Folds to be checked or `NULL`.

#' @param N Number of observations.

#'

#' @return

#' An integer vector with one element per observation indicating the

#' cross-validation fold in which the observation should be withdrawn.

#'

#' @noRd

validate.folds <- function(folds, N) {

    if (is.null(folds))

        return(rep(1, N))

    validate.indices(folds, N, "folds", throw.duplicates=FALSE)

    if (length(folds) != N)

        stop("'folds' should have length ", N, ".")

    K <- length(unique(folds))

    if (!all(1:K %in% folds))

        stop("'folds' must contain all indices up to ", K, ".")

    folds <- as.integer(folds)

}

#' Validate start.from

#'

#' Check that the predictor names provided is a valid subset of the variables

#' used in the model.

#'

#' @param obj An object of class `hsstan`.

#' @param start.from Vector to be checked.

#'

#' @return

#' A list of two elements: the names of the model terms matching `start.from`

#' and a vector of indices corresponding to the names listed in `start.from`.

#' Throws an error if any of the names mentioned does not match those available

#' in the model terms.

#'

#' @noRd

validate.start.from <- function(obj, start.from) {

    unp.terms <- obj$model.terms$unpenalized

    unp.betas <- names(obj$betas$unpenalized)

    mod.terms <- c(unp.terms, obj$model.terms$penalized)

    mod.betas <- c(unp.betas, names(obj$betas$penalized))

    start.from <- setdiff(start.from, "")

    if (is.null(start.from)) {

        if (length(obj$model.terms$penalized) > 0)

            return(list(start.from=unp.terms, idx=seq_along(unp.betas)))

        else

            return(list(start.from=character(0), idx=1))

    }

    if (length(start.from) == 0)

        return(list(start.from=character(0), idx=1))

    if (anyNA(start.from))

        stop("'start.from' contains missing values.")

    var.match <- match(start.from, mod.terms)

    if (anyNA(var.match))

        stop("'start.from' contains ", collapse(start.from[is.na(var.match)]),

             ", which cannot be matched.")

    ## unpack interaction terms so that also main effects are matched

    start.from <- mod.terms[mod.terms %in%

                            c(start.from, unlist(strsplit(start.from, ":")))]

    chosen <- expand.terms(obj$data, start.from)

    ## also consider interaction terms in reverse order

    chosen <- c(chosen, sapply(strsplit(chosen[grep(":", chosen)], ":"),

                               function(z) c(z, paste(rev(z), collapse=":"))))

    return(list(start.from=start.from, idx=which(mod.betas %in% chosen)))

}

#' Validate a positive or non-negative scalar value

#'

#' @param x Value to validate.

#' @param name Variable name to report in case of error.

#' @param int Whether the value has to be an integer (`FALSE` by default).

#'

#' @return

#' Throws an error if the given value is not a positive or non-negative

#' scalar (or integer scalar).

#'

#' @noRd

validate.positive.scalar <- function(x, name, int=FALSE) {

    if (!is.numeric(x) || length(x) != 1 || is.na(x) || x <= 0 ||

        (int && (x > .Machine$integer.max || x != as.integer(x))))

        stop(sprintf("'%s' must be a positive %s.", name,

                     ifelse(int, "integer", "scalar")), call.=FALSE)

}

#' @noRd

validate.nonnegative.scalar <- function(x, name, int=FALSE) {

    if (!is.numeric(x) || length(x) != 1 || is.na(x) || x < 0 ||

        (int && (x > .Machine$integer.max || x != as.integer(x))))

        stop(sprintf("'%s' must be a non-negative %s.", name,

                     ifelse(int, "integer", "scalar")), call.=FALSE)

}

#' Validate adapt.delta

#'

#' Check that an adaptation acceptance probability is valid.

#'

#' @param adapt.delta Value to be checked.

#'

#' @return

#' Throws an error if the given value is not a valid acceptance probability

#' for adaptation.

#'

#' @noRd

validate.adapt.delta <- function(adapt.delta) {

    if (!is.numeric(adapt.delta) || length(adapt.delta) != 1) {

        stop("'adapt.delta' must be a single numerical value.")

    }

    if (adapt.delta < 0.8) {

        stop("'adapt.delta' must be at least 0.8.")

    }

    if (adapt.delta >= 1) {

        stop("'adapt.delta' must be less than 1.")

    }

}

#' Validate a probability

#'

#' Check that a probability value is valid.

#'

#' @param prob Value to be checked.

#'

#' @return

#' Throws an error if the given value is not a valid probability.

#'

#' @noRd

validate.probability <- function(prob) {

    if (length(prob) != 1 || prob <= 0 || prob >= 1)

        stop("'prob' must be a single value between 0 and 1.\n")

}

#' Validate arguments passed to rstan

#'

#' Ensure that the options to be passed to \code{\link[rstan]{sampling}} are

#' valid, as to work around rstan issue #681.

#'

#' @param ... List of arguments to be checked.

#'

#' @return

#' Throws an error if any argument is not valid for \code{\link[rstan]{sampling}}.

#'

#' @noRd

validate.rstan.args <- function(...) {

    valid.args <- c("chains", "cores", "pars", "thin", "init", "check_data",

                    "sample_file", "diagnostic_file", "verbose", "algorithm",

                    "control", "open_progress", "show_messages", "chain_id",

                    "init_r", "test_grad", "append_samples", "refresh",

                    "save_warmup", "enable_random_init", "iter", "warmup")

    dots <- list(...)

    for (arg in names(dots))

        if (!arg %in% valid.args)

           stop("Argument '", arg, "' not recognized.")

}

#' Parameter names

#'

#' Get the parameter names corresponding to the regression coefficients or

#' matching a regular expression.

#'

#' @param obj An object of class `hsstan`.

#' @param pars Regular expression to match a parameter name, or `NULL`

#'        to retrieve the names of all regression coefficients.

#'

#' @return

#' A character vector.

#'

#' @noRd

get.pars <- function(object, pars) {

    if (is.null(pars))

        pars <- grep("^beta_", object$stanfit@model_pars, value=TRUE)

    else {

        if (!is.character(pars))

            stop("'pars' must be a character vector.")

        get.pars <- function(x) grep(x, object$stanfit@sim$fnames_oi, value=TRUE)

        pars <- unlist(lapply(pars, get.pars))

        if (length(pars) == 0)

            stop("No pattern in 'pars' matches parameter names.")

    }

    return(pars)

}

#' Create a design matrix with all unpenalized predictors first

#'

#' This is required as `model.matrix` puts the interaction terms after the

#' penalized predictors, but the Stan models expects all unpenalized terms to

#' appear before the penalized ones.

#'

#' @param x Data frame containing the variables of interest.

#' @param unpenalized Vector of variable names for the unpenalized covariates.

#' @param penalized Vector of variable names for the penalized predictors.

#'

#' @return

#' A design matrix with all unpenalized covariates (including interaction terms)

#' before the penalized predictors.

#'

#' @importFrom stats model.matrix reformulate

#' @noRd

ordered.model.matrix <- function(x, unpenalized, penalized) {

    X <- model.matrix(reformulate(c(unpenalized, penalized)), data=x)

    if (any(grepl("[:*]", unpenalized)))

        X <- X[, c(expand.terms(x, unpenalized), expand.terms(x, penalized)[-1])]

    return(X)

}

#' Expand variable names into formula terms

#'

#' @param x Data frame containing the variables of interest.

#' @param variables Vector of variable names.

#'

#' @return

#' A vector of variable names expanded by factor levels and interaction terms.

#'

#' @importFrom stats model.matrix reformulate

#' @noRd

expand.terms <- function(x, variables) {

    if (length(variables) == 0)

        return(character(0))

    colnames(model.matrix(reformulate(variables), x[1, ]))

}

#' Summarize a vector

#'

#' @param x A numerical vector.

#' @param prob Width of the interval between quantiles.

#'

#' @return

#' The mean, standard deviation and quantiles for the input vector.

#'

#' @noRd

vector.summary <- function(x, prob) {

    lower <- (1 - prob) / 2

    upper <- 1 - lower

    c(mean=mean(x), sd=stats::sd(x), stats::quantile(x, c(lower, upper)))

}

#' Check whether the model fitted is a logistic regression model.

#'

#' @param obj An object of class `hsstan`.

#'

#' @return

#' `TRUE` for logistic regression models, `FALSE` otherwise.

#'

#' @noRd

is.logistic <- function(obj) {

    obj$family$family == "binomial"

}

#' Comma-separated string concatenation

#'

#' Collapse the elements of a character vector into a comma-separated string.

#'

#' @param x Character vector.

#'

#' @return

#' A comma-separated string where each element of the original vector is

#' surrounded by single quotes.

#'

#' @noRd

collapse <- function(x) {

    paste0("'", x, "'", collapse=", ")

}

#' Fast computation of correlations

#'

#' This provides a loopless version of the computation of the correlation

#' coefficient between observed and predicted outcomes.

#'

#' @param y Vector of observed outcome.

#' @param x Matrix with as many columns as the number of elements in `y`,

#'          where each row corresponds to a predicted outcome.

#'

#' @return

#' A vector of correlations with as many elements as the number of rows in `x`.

#'

#' @noRd

fastCor <- function(y, x) {

    yx <- rbind(y, x)

    if (.Machine$sizeof.pointer == 8) {

        yx <- yx - rowMeans(yx)

        yx <- yx / sqrt(rowSums(yx^2))

        corr <- tcrossprod(yx, yx)

    } else {

        corr <- stats::cor(yx)

    }

    return(corr[-1, 1])

}

#' Log of sum of exponentials

#'

#' @noRd

logSumExp <- function(x) {

    xmax <- max(x)

    xmax + log(sum(exp(x - xmax)))

}

#' Log of average of exponentials

#'

#' @noRd

logMeanExp <- function(x) {

    logSumExp(x) - log(length(x))

}