#' Bayes Factors (BF) for model comparison
#'
#' @description This function computes or extracts Bayes factors from fitted
#' models.
#' \cr\cr
#' The `bf_*` function is an alias of the main function.
#' \cr\cr
#' \strong{For more info, see [the Bayes factors vignette](https://easystats.github.io/bayestestR/articles/bayes_factors.html).}
#'
#' @author Mattan S. Ben-Shachar
#'
#' @param ... Fitted models (see details), all fit on the same data, or a single
#'   `BFBayesFactor` object (see 'Details'). Ignored in `as.matrix()`,
#'   `update()`. If the following named arguments are present, they are passed
#'   to [`insight::get_loglikelihood()`] (see details):
#'   - `estimator` (defaults to `"ML"`)
#'   - `check_response`  (defaults to `FALSE`)
#' @param denominator Either an integer indicating which of the models to use as
#'   the denominator, or a model to be used as a denominator. Ignored for
#'   `BFBayesFactor`.
#'
#' @inheritParams hdi
#'
#' @note There is also a [`plot()`-method](https://easystats.github.io/see/articles/bayestestR.html)
#' implemented in the \href{https://easystats.github.io/see/}{\pkg{see}-package}.
#'
#' @details
#' If the passed models are supported by **insight** the DV of all models will
#' be tested for equality (else this is assumed to be true), and the models'
#' terms will be extracted (allowing for follow-up analysis with [bayesfactor_inclusion]).
#'
#' - For `brmsfit` or `stanreg` models, Bayes factors are computed using the \CRANpkg{bridgesampling} package.
#'   - `brmsfit` models must have been fitted with `save_pars = save_pars(all = TRUE)`.
#'   - `stanreg` models must have been fitted with a defined `diagnostic_file`.
#' - For `BFBayesFactor`, `bayesfactor_models()` is a wraparound `BayesFactor::extractBF()`.
#' - For all other model types, Bayes factors are computed using the BIC approximation.
#'   Note that BICs are extracted from using [insight::get_loglikelihood], see documentation
#'   there for options for dealing with transformed responses and REML estimation.
#'
#' ## Additional methods
#' The resulting output is supported by the following methods:
#'
#' - `as.matrix()`: Extract a full matrix of (log-)Bayes factors between all
#'   models (using the transitivity of Bayes factors).
#' - `update()`: subset and/or re-reference the Bayes factors to a different model.
#' - `as.numeric()`: Extract the (possibly log-)Bayes factor values.
#'
#' See examples and [bayesfactor_methods].
#'
#' @inheritSection bayesfactor_methods Prior and posterior considerations
#'
#' @inheritSection bayesfactor_methods Transitivity of Bayes factors
#'
#' @inheritSection bayesfactor_methods Interpreting Bayes Factors
#'
#' @return A data frame containing the models' formulas (reconstructed fixed and
#'   random effects) and their `log(BF)`s  (Use `as.numeric()` to extract the
#'   non-log Bayes factors; see examples), that prints nicely.
#'   \cr\cr
#'   For `as.matrix()` a square matrix of (log) Bayes factors, with rows as
#'   denominators and columns as numerators.
#'
#' @examplesIf require("lme4") && require("BayesFactor") && require("rstanarm") && require("brms")
#' # With lm objects:
#' # ----------------
#' lm1 <- lm(mpg ~ 1, data = mtcars)
#' lm2 <- lm(mpg ~ hp, data = mtcars)
#' lm3 <- lm(mpg ~ hp + drat, data = mtcars)
#' lm4 <- lm(mpg ~ hp * drat, data = mtcars)
#' (BFM <- bayesfactor_models(lm1, lm2, lm3, lm4, denominator = 1))
#' # bayesfactor_models(lm2, lm3, lm4, denominator = lm1) # same result
#' # bayesfactor_models(lm1, lm2, lm3, lm4, denominator = lm1) # same result
#'
#' update(BFM, reference = "bottom")
#' as.matrix(BFM)
#' as.numeric(BFM)
#'
#' lm2b <- lm(sqrt(mpg) ~ hp, data = mtcars)
#' # Set check_response = TRUE for transformed responses
#' bayesfactor_models(lm2b, denominator = lm2, check_response = TRUE)
#'
#' \donttest{
#' # With lmerMod objects:
#' # ---------------------
#' lmer1 <- lme4::lmer(Sepal.Length ~ Petal.Length + (1 | Species), data = iris)
#' lmer2 <- lme4::lmer(Sepal.Length ~ Petal.Length + (Petal.Length | Species), data = iris)
#' lmer3 <- lme4::lmer(
#'   Sepal.Length ~ Petal.Length + (Petal.Length | Species) + (1 | Petal.Width),
#'   data = iris
#' )
#' bayesfactor_models(lmer1, lmer2, lmer3,
#'   denominator = 1,
#'   estimator = "REML"
#' )
#'
#' # rstanarm models
#' # ---------------------
#' # (note that a unique diagnostic_file MUST be specified in order to work)
#' stan_m0 <- suppressWarnings(rstanarm::stan_glm(Sepal.Length ~ 1,
#'   data = iris,
#'   family = gaussian(),
#'   diagnostic_file = file.path(tempdir(), "df0.csv")
#' ))
#' stan_m1 <- suppressWarnings(rstanarm::stan_glm(Sepal.Length ~ Species,
#'   data = iris,
#'   family = gaussian(),
#'   diagnostic_file = file.path(tempdir(), "df1.csv")
#' ))
#' stan_m2 <- suppressWarnings(rstanarm::stan_glm(Sepal.Length ~ Species + Petal.Length,
#'   data = iris,
#'   family = gaussian(),
#'   diagnostic_file = file.path(tempdir(), "df2.csv")
#' ))
#' bayesfactor_models(stan_m1, stan_m2, denominator = stan_m0, verbose = FALSE)
#'
#'
#' # brms models
#' # --------------------
#' # (note the save_pars MUST be set to save_pars(all = TRUE) in order to work)
#' brm1 <- brms::brm(Sepal.Length ~ 1, data = iris, save_pars = save_pars(all = TRUE))
#' brm2 <- brms::brm(Sepal.Length ~ Species, data = iris, save_pars = save_pars(all = TRUE))
#' brm3 <- brms::brm(
#'   Sepal.Length ~ Species + Petal.Length,
#'   data = iris,
#'   save_pars = save_pars(all = TRUE)
#' )
#'
#' bayesfactor_models(brm1, brm2, brm3, denominator = 1, verbose = FALSE)
#'
#'
#' # BayesFactor
#' # ---------------------------
#' data(puzzles)
#' BF <- BayesFactor::anovaBF(RT ~ shape * color + ID,
#'   data = puzzles,
#'   whichRandom = "ID", progress = FALSE
#' )
#' BF
#' bayesfactor_models(BF) # basically the same
#' }
#'
#' @references
#' - Gronau, Q. F., Singmann, H., & Wagenmakers, E. J. (2017). Bridgesampling: An R package for estimating
#'   normalizing constants. arXiv preprint arXiv:1710.08162.
#'
#' - Kass, R. E., and Raftery, A. E. (1995). Bayes Factors. Journal of the American Statistical Association,
#'   90(430), 773-795.
#'
#' - Robert, C. P. (2016). The expected demise of the Bayes factor. Journal of Mathematical Psychology,
#'   72, 33–37.
#'
#' - Wagenmakers, E. J. (2007). A practical solution to the pervasive problems of p values.
#'   Psychonomic bulletin & review, 14(5), 779-804.
#'
#' - Wetzels, R., Matzke, D., Lee, M. D., Rouder, J. N., Iverson, G. J., and Wagenmakers, E.-J. (2011).
#'   Statistical Evidence in Experimental Psychology: An Empirical Comparison Using 855 t Tests.
#'   Perspectives on Psychological Science, 6(3), 291–298. \doi{10.1177/1745691611406923}
#'
#' @seealso [bayesfactor_inclusion()] for testing predictors across Bayesian models.
#' @family Bayes factors
#'
#' @export
bayesfactor_models <- function(..., denominator = 1, verbose = TRUE) {
  UseMethod("bayesfactor_models")
}

#' @rdname bayesfactor_models
#' @export
bf_models <- bayesfactor_models


#' @export
#' @rdname bayesfactor_models
bayesfactor_models.default <- function(..., denominator = 1, verbose = TRUE) {
  # Organize the models and their names
  mods <- list(...)
  denominator <- list(denominator)
  cl <- match.call(expand.dots = FALSE)

  estimator <- mods[["estimator"]]
  check_response <- mods[["check_response"]]
  if (is.null(estimator)) {
    estimator <- "ML"
  }
  if (is.null(check_response)) {
    check_response <- FALSE
  }
  mods[["check_response"]] <- mods[["estimator"]] <- NULL
  cl$...$estimator <- cl$...$check_response <- NULL

  names(mods) <- sapply(cl[["..."]], insight::safe_deparse)
  names(denominator) <- insight::safe_deparse(cl$denominator)

  mods <- .cleanup_BF_models(mods, denominator, cl)
  mforms <- names(mods)
  denominator <- attr(mods, "denominator", exact = TRUE)

  # Get formula / model names
  # supported models
  supported_models <- vapply(mods, insight::is_model_supported, TRUE)
  if (all(supported_models)) {
    temp_forms <- sapply(mods, .find_full_formula)

    has_terms <- sapply(temp_forms, nchar) > 0

    mforms[has_terms] <- temp_forms[has_terms]
    supported_models[!has_terms] <- FALSE
  }

  model_objects <- .safe(do.call(insight::ellipsis_info, c(mods, verbose = FALSE)))
  if (!is.null(model_objects)) {
    were_checked <- inherits(model_objects, "ListModels")

    # Validate response
    if (were_checked && verbose && !isTRUE(attr(model_objects, "same_response"))) {
      insight::format_warning(
        "When comparing models, please note that probably not all models were fit from same data."
      )
    }

    # Get BIC
    if (
      were_checked &&
        estimator == "REML" &&
        any(vapply(mods, insight::is_mixed_model, TRUE)) &&
        !isTRUE(attr(model_objects, "same_fixef")) &&
        verbose
    ) {
      insight::format_warning(paste(
        "Information criteria (like BIC) based on REML fits (i.e. `estimator=\"REML\"`)",
        "are not recommended for comparison between models with different fixed effects.",
        "Concider setting `estimator=\"ML\"`."
      ))
    }
  } else if (verbose) {
    insight::format_alert(
      "Unable to validate that all models were fit with the same data."
    )
  }

  mBIC <- tryCatch(
    sapply(mods, function(m) {
      LL <- insight::get_loglikelihood(
        m,
        estimator = estimator,
        check_response = check_response
      )
      stats::BIC(LL)
    }),
    error = function(...) NULL
  )

  if (is.null(mBIC)) {
    mBIC <- sapply(mods, stats::BIC)
  }

  # Get BF
  mBFs <- bic_to_bf(mBIC, denominator = mBIC[denominator], log = TRUE)

  res <- data.frame(
    Model = mforms,
    log_BF = mBFs,
    stringsAsFactors = FALSE
  )

  .bf_models_output(
    res,
    denominator = denominator,
    bf_method = "BIC approximation",
    unsupported_models = !all(supported_models),
    model_names = names(mods)
  )
}


.bayesfactor_models_stan <- function(mods, denominator = 1, verbose = TRUE) {
  # Warn
  n_samps <- sapply(mods, function(x) {
    alg <- insight::find_algorithm(x)
    if (is.null(alg)) {
      return(NA_integer_)
    }
    if (is.null(alg$iterations)) {
      alg$iterations <- alg$sample
    }
    (alg$iterations - alg$warmup) * alg$chains
  })
  if (verbose) {
    if (any(n_samps < 4e4, na.rm = TRUE)) {
      insight::format_warning(
        "Bayes factors might not be precise.",
        "For precise Bayes factors, sampling at least 40,000 posterior samples is recommended."
      )
    } else if (any(is.na(n_samps))) {
      insight::format_alert(
        "Unable to determine the number of posterior samples.",
        "Bayes factors might not be precise."
      )
    }
  }

  if (inherits(mods[[1]], "blavaan")) {
    res <- .bayesfactor_models_stan_SEM(mods, denominator, verbose)
    bf_method <- "marginal likelihoods (Laplace approximation)"
    unsupported_models <- TRUE
  } else {
    res <- .bayesfactor_models_stan_REG(mods, denominator, verbose)
    bf_method <- "marginal likelihoods (bridgesampling)"
    unsupported_models <- inherits(mods[[1]], c("stanfit", "CmdStanFit"))
  }

  .bf_models_output(
    res,
    denominator = denominator,
    bf_method = bf_method,
    unsupported_models = unsupported_models
  )
}


#' @keywords internal
.bayesfactor_models_stan_REG <- function(mods, denominator, verbose = TRUE) {
  insight::check_if_installed("bridgesampling")

  # Test that all is good:
  resps <- lapply(mods, insight::get_response)
  if (all(!sapply(resps, is.null))) {
    from_same_data_as_den <- sapply(
      resps[-denominator],
      identical,
      y = resps[[denominator]]
    )

    if (!all(from_same_data_as_den)) {
      insight::format_error("Models were not computed from the same data.")
    }
  } else if (verbose) {
    insight::format_alert(
      "Unable to validate that all models were fit with the same data."
    )
  }

  mML <- lapply(mods, .get_marglik, verbose = verbose)

  mBFs <- sapply(mML, function(x) {
    bf <- bridgesampling::bf(x, mML[[denominator]], log = TRUE)
    bf[["bf"]]
  })

  # Get formula
  mforms <- sapply(mods, .find_full_formula)
  mforms[!nzchar(mforms)] <- names(mforms)[!nzchar(mforms)]

  res <- data.frame(
    Model = mforms,
    log_BF = mBFs,
    stringsAsFactors = FALSE
  )
}

.bayesfactor_models_stan_SEM <- function(mods, denominator, verbose = TRUE) {
  utils::capture.output(
    suppressWarnings({
      mBFs <- sapply(mods, function(m) {
        blavaan::blavCompare(m, mods[[denominator]])[["bf"]][1]
      })
    })
  )

  res <- data.frame(
    Model = names(mods),
    log_BF = unname(mBFs),
    stringsAsFactors = FALSE
  )
}


#' @export
bayesfactor_models.stanreg <- function(..., denominator = 1, verbose = TRUE) {
  mods <- list(...)
  if (inherits(mods[[1]], "stanreg")) {
    insight::check_if_installed("rstanarm")
  } else if (inherits(mods[[1]], "brmsfit")) {
    insight::check_if_installed("brms")
  } else if (inherits(mods[[1]], "blavaan")) {
    insight::check_if_installed("blavaan")
  }

  # Organize the models and their names
  denominator <- list(denominator)

  cl <- match.call(expand.dots = FALSE)
  names(mods) <- sapply(cl[["..."]], insight::safe_deparse)
  names(denominator) <- insight::safe_deparse(cl$denominator)

  mods <- .cleanup_BF_models(mods, denominator, cl)
  denominator <- attr(mods, "denominator", exact = TRUE)

  .bayesfactor_models_stan(mods, denominator = denominator, verbose = verbose)
}


#' @export
bayesfactor_models.brmsfit <- bayesfactor_models.stanreg


#' @export
bayesfactor_models.blavaan <- bayesfactor_models.stanreg


#' @export
bayesfactor_models.stanfit <- bayesfactor_models.stanreg


#' @export
bayesfactor_models.CmdStanFit <- bayesfactor_models.stanreg

#' @export
bayesfactor_models.BFBayesFactor <- function(..., verbose = TRUE) {
  models <- c(...)

  insight::check_if_installed("BayesFactor")

  mBFs <- c(0, BayesFactor::extractBF(models, TRUE, TRUE))
  mforms <- sapply(c(models@denominator, models@numerator), function(x) x@shortName)

  if (inherits(models@denominator, "BFlinearModel")) {
    mforms[mforms == "Intercept only"] <- "1"
  } else {
    mforms <- .clean_non_linBF_mods(mforms)
  }

  res <- data.frame(
    Model = unname(mforms),
    log_BF = mBFs,
    stringsAsFactors = FALSE
  )

  .bf_models_output(
    res,
    denominator = 1,
    bf_method = "JZS (BayesFactor)",
    unsupported_models = !inherits(models@denominator, "BFlinearModel")
  )
}

# Helpers -----------------------------------------------------------------

#' @keywords internal
.cleanup_BF_models <- function(mods, denominator, cl) {
  if (length(mods) == 1 && inherits(mods[[1]], "list")) {
    mods <- mods[[1]]
    mod_names <- .safe(sapply(cl[["..."]][[1]][-1], insight::safe_deparse))

    if (!is.null(mod_names) && length(mod_names) == length(mods)) {
      names(mods) <- mod_names
    }
  }

  if (is.numeric(denominator[[1]])) {
    denominator <- denominator[[1]]
  } else {
    denominator_model <- which(names(mods) == names(denominator))

    if (length(denominator_model) == 0) {
      mods <- c(mods, denominator)
      denominator <- length(mods)
    } else {
      denominator <- denominator_model
    }
  }

  attr(mods, "denominator") <- denominator
  mods
}


#' @keywords internal
.bf_models_output <- function(
  res,
  denominator = 1,
  bf_method = "method",
  unsupported_models = FALSE,
  model_names = NULL
) {
  # sanity check - are all BF NA?
  if (!is.null(res$log_BF) && all(is.na(res$log_BF))) {
    insight::format_error(
      "Could not calculate Bayes Factor for these models. You may report this problem at {https://github.com/easystats/bayestestR/issues/}."
    )
  }

  if (is.null(model_names)) {
    model_names <- rownames(res)
  }

  attr(res, "denominator") <- denominator
  attr(res, "BF_method") <- bf_method
  attr(res, "unsupported_models") <- unsupported_models
  attr(res, "model_names") <- model_names
  class(res) <- c(
    "bayestestRBF",
    "bayesfactor_models",
    "see_bayesfactor_models",
    class(res)
  )

  res
}


#' @keywords internal
.find_full_formula <- function(mod) {
  formulas <- insight::find_formula(mod)

  conditional <- random <- NULL
  if (!is.null(formulas$conditional)) {
    conditional <- as.character(formulas$conditional)[3]
  }

  if (!is.null(formulas$random)) {
    if (!is.list(formulas$random)) {
      formulas$random <- list(formulas$random)
    }
    random <- sapply(formulas$random, function(x) {
      paste0("(", as.character(x)[2], ")")
    })
  }
  paste(c(conditional, random), collapse = " + ")
}

#' @keywords internal
.clean_non_linBF_mods <- function(m_names) {
  tryCatch(
    {
      m_txt <- character(length = length(m_names))

      ## Detect types ##
      is_null <- startsWith(m_names, "Null")
      is_rho <- grepl("rho", m_names, fixed = TRUE)
      is_mu <- grepl("mu", m_names, fixed = TRUE)
      is_d <- grepl("d", m_names, fixed = TRUE)
      is_p <- grepl("p", m_names, fixed = TRUE)
      is_range <- grepl("<", m_names, fixed = TRUE)

      ## Range Alts ##
      m_txt[!is_null & is_range] <-
        sub("^[^\\s]*\\s[^\\s]*\\s", "", m_names[!is_null & is_range])

      ## Null models + Not nulls ##
      if (any(is_d & is_p)) {
        is_null <- !startsWith(m_names, "Non")
        temp <- m_names[is_null][1]
        mi <- gregexpr("\\(.*\\)", temp)
        aa <- unlist(regmatches(temp, m = mi), use.names = FALSE)

        m_txt[is_null] <- sub("a=", "a = ", aa, fixed = TRUE)
        m_txt[!is_null & !is_range] <- sub("a=", "a != ", aa, fixed = TRUE)
      } else if (any(is_rho)) {
        m_txt[is_null] <- "rho = 0"
        m_txt[!is_null & !is_range] <- "rho != 0"
        m_txt <- sub("<rho<", " < rho < ", m_txt, fixed = TRUE)
      } else if (any(is_d | is_mu)) {
        m_txt[is_null] <- "d = 0"
        m_txt[!is_null & !is_range] <- "d != 0"
        m_txt <- sub("<d<", " < d < ", m_txt, fixed = TRUE)
      } else if (any(is_p)) {
        temp <- m_names[is_null][1]
        mi <- gregexpr("[0-9|\\.]+", temp)
        pp <- unlist(regmatches(temp, m = mi), use.names = FALSE)

        m_txt[is_null] <- paste0("p = ", pp)
        m_txt[!is_null & !is_range] <- paste0("p != ", pp)
        m_txt <- sub("<p<", " < p < ", m_txt, fixed = TRUE)
      } else {
        stop("!", call. = FALSE)
      }

      ## wrap with () for readability ##
      is_wrapped <- grepl("(", m_txt, fixed = TRUE)
      m_txt[!is_wrapped] <- paste0("(", m_txt[!is_wrapped], ")")

      m_txt
    },
    error = function(e) {
      m_names
    }
  )
}

#' @keywords internal
.get_marglik <- function(mod, verbose, ...) {
  # Add a check here for brmsfit object to avoid unnecessary computation of the ML
  if (inherits(mod, "brmsfit") && "marglik" %in% names(mod$criteria)) {
    return(stats::median(mod$criteria$marglik$logml))
  }

  # Else... Get marginal likelihood
  if (verbose) {
    insight::format_alert(
      "Computation of Marginal Likelihood: estimating marginal likelihood, please wait..."
    )
  }
  # Should probably allow additional arguments such as reps or cores to for bridge_sampler
  bridgesampling::bridge_sampler(mod, silent = TRUE)
}
