#' Win/Loss Ratios with Efficient Influence Functions
#'
#' `eif_win_ratio()` calculates the win/loss proportion, win ratio, and net benefit
#' by adding a correction term to `drf_win_ratio()` using Efficient Influence Functions.
#' @param base A data frame with n rows (individuals) and p columns.
#' It must include a treatment column and one or more outcome columns.
#' All remaining columns are considered covariates for matching.
#' The data frame should not contain missing values. All covariates must be either numeric or factors.
#' @param treatment_name A character vector of length one for the treatment column name, must be in `names(base)`.
#' The treatment column should contain only zeros and ones, with ones indicating the treatment group.
#' @param outcome_names A character vector specifying the names of outcome columns.
#' All names must exist in `names(base)`. Each outcome column should be numeric or binary and
#' binary outcomes must be encoded as 0 and 1.
#' @param win_function By default, this function compares outcomes lexicographically according to the column order in the data frame,
#' treating higher values as better. Once a distinct outcome is found, the function yields 1 (a win) if that outcome is larger
#'  for the treated individual and 0 otherwise; in case of a tie, it yields 0.
#' To change the priority of outcomes, simply reorder the columns of the data frame.
#' Alternatively, `win_function` can be a custom function that takes two arguments, each corresponding to a row of the data frame:
#' `win_function(base[i, ], base[j, ])`
#' where the first argument represents a treated individual and the second a control individual.
#' The function must return a value between 0 and 1, with 1 indicating a win for the treated individual.
#' @param distance A function that is used to compute the nearest neighbor matching on covariates.
#' The distance function will be passed as an argument to \code{\link[MatchIt]{matchit}} (see  \code{\link[MatchIt]{distance}}).
#'
#'  **Default Distance Settings**
#'
#' The default value for the distance parameter is `NULL`. If no distance is provided, the function assigns defaults based on the type of covariates:
#'
#' - **Mixed covariates:** Applies Factor Analysis of Mixed Data transformation (see \code{\link[FactoMineR]{FAMD}}) to the covariates and sets distance to `"euclidean"`.
#' - **Numeric covariates:** Sets distance to `"euclidean"`.
#' - **Categorical covariates:** Applies Multiple Correspondence Analysis (see \code{\link[FactoMineR]{MCA}}) transformation to the covariates and sets distance to `"euclidean"`.
#' @param n_trees_drf Number of trees grown in the distributional random forests.
#' It passed as the argument `num.trees` to \code{\link[drf]{drf}}. The default value is 1000.
#' @param n_of_folds Number of folds used in cross-fitting. The default value is 3.
#' @param seed Numeric value used in `set.seed()` to guarantee reproducible results when sampling the folds. The default value is `NULL`.
#' @param propensity_model Character string specifying the type of model to use.
#'   Must be either "glm" or "probability_forest".
#'   Defaults to "probability_forest". The model is only used if no vector
#'   of propensity scores (`propensity_scores`) is provided.
#' @param n_trees_propensity Number of trees grown when computing the propensity scores using `propensity_model`.
#' @param propensity_scores Numeric vector of propensity score probabilities
#' @param clip_propensity Numeric value used to clip propensity scores,
#'   ensuring they lie within the interval `[clip_propensity, 1 - clip_propensity]`. The default value is zero.
#'
#' @returns A list of six elements, each of which is a list structured as follows:
#'
#' The first element, `win_proportion`, is a list with two components: `value` and `std`,
#' where `std` represents the standard deviation.
#'
#' The second element, `lose_proportion`, is a list with two components: `value` and `std`,
#' where `std` represents the standard deviation.
#'
#' The third element, `win_ratio`, is a list containing three components: `value`, `lower_CI`,
#' and `upper_CI`, with `lower_CI` and `upper_CI` representing the confidence interval bounds.
#'
#' The fourth element, `net_benefit`, is a list containing three components: `value`, `lower_CI`,
#' and `upper_CI`, with `lower_CI` and `upper_CI` representing the confidence interval bounds.
#'
#' The fifth and sixth elements, `win_vector` and `lose_vector`, are numeric vectors containing the plug-in estimator
#' evaluated for each individual, along with the corresponding EIF correction term.
#' @export
#'
#' @examples
#' ## Input Propensity Scores ##
#'
#' set.seed(123)
#' n <- 100
#' base <- data.frame(
#'   X1 = rnorm(n, mean = 5, sd = 2), # numeric covariate 1
#'   X2 = rnorm(n, mean = 10, sd = 3), # numeric covariate 2
#'   X3 = factor(sample(c("A", "B", "C"), n, replace = TRUE)), # factor covariate
#'   Y1 = rnorm(n, mean = 50, sd = 10), # numeric outcome 1
#'   Y2 = rnorm(n, mean = 100, sd = 20), # numeric outcome 2
#'   arm = sample(c(0, 1), n, replace = TRUE) # binary treatment arm
#' )
#' treatment_name <- "arm"
#' outcome_names <- c("Y1", "Y2")
#'
#' ps_model <- stats::glm(arm ~ X1 + X3, # only consider X1 and X3 and no cross-fitting
#'   family = binomial(),
#'   data = base
#' )
#'
#' propensity_scores <- stats::predict(ps_model, type = "response")
#' res <- eif_win_ratio(
#'   base = base,
#'   treatment_name = treatment_name,
#'   outcome_names = outcome_names,
#'   propensity_scores = propensity_scores
#' )
eif_win_ratio <- function(base,
                          treatment_name,
                          outcome_names,
                          win_function = "lexicographic",
                          distance = NULL,
                          n_trees_drf = 1000,
                          n_of_folds = 3,
                          seed = NULL,
                          propensity_model = "probability_forest",
                          n_trees_propensity = 1000,
                          propensity_scores = NULL,
                          clip_propensity = 0.0) {
  # Check Input & Compute 'drf_win_ratio' and 'nn_win_ratio' ----------------

  validate_propensity_input(propensity_model, n_trees_propensity, propensity_scores, clip_propensity)

  # The remaining inputs are checked by the calls to 'nn_win_ratio' and 'drf_win_ratio'
  res_nn <- nn_win_ratio(
    base = base,
    treatment_name = treatment_name,
    outcome_names = outcome_names,
    win_function = win_function,
    distance = distance
  )

  res_drf <- drf_win_ratio(
    base = base,
    treatment_name = treatment_name,
    outcome_names = outcome_names,
    win_function = win_function,
    n_trees_drf = n_trees_drf,
    n_of_folds = n_of_folds,
    seed = seed
  )

  win_vector_nn <- res_nn$win_vector
  lose_vector_nn <- res_nn$lose_vector

  win_vector_drf <- res_drf$win_vector
  lose_vector_drf <- res_drf$lose_vector

  n_of_individuals <- nrow(base)

  # Compute Propensity Scores (if needed) -----------------------------------

  if (is.null(propensity_scores)) {
    propensity_model_str <- sub(
      "^[a-z]",
      toupper(substr(propensity_model, 1, 1)),
      sub("_", " ", propensity_model)
    )
    message(
      sprintf(
        "No 'propensity_scores' was provided. '%s' will be used instead to compute those scores.",
        propensity_model_str
      )
    )
    message("Specify 'propensity_scores' to override this choice.")

    propensity_scores <- compute_propensity_scores(
      base = base,
      treatment_name = treatment_name,
      outcome_names = outcome_names,
      n_of_folds = n_of_folds,
      seed = seed,
      propensity_model = propensity_model,
      n_trees_propensity = n_trees_propensity
    )
  }
  for (i in 1:n_of_individuals) {
    propensity_scores[i] <- min(propensity_scores[i], 1 - clip_propensity)
    propensity_scores[i] <- max(propensity_scores[i], clip_propensity)
  }

  ip_weights <- ifelse(
    base[[treatment_name]] == 1,
    1 / propensity_scores,
    1 / (1 - propensity_scores)
  )
  # Compute Win/Lose vectors ------------------------------------------------

  win_vector <- win_vector_nn + ip_weights * (win_vector_drf - win_vector_nn)
  lose_vector <- lose_vector_nn + ip_weights * (lose_vector_drf - lose_vector_nn)

  # Compute Win/Lose Metrics ------------------------------------------------

  res <- compute_metrics(win_vector, lose_vector)
  res$win_vector <- win_vector
  res$lose_vector <- lose_vector

  return(res)
}


compute_propensity_scores <- function(base,
                                      treatment_name,
                                      outcome_names,
                                      n_of_folds,
                                      seed,
                                      propensity_model,
                                      n_trees_propensity) {
  # Initialize
  n_of_individuals <- nrow(base)
  covariate_names <- setdiff(colnames(base), c(outcome_names, treatment_name))
  propensity_scores <- numeric(n_of_individuals)

  X <- stats::model.matrix(~ . - 1, data = base[, covariate_names, drop = FALSE])
  X <- as.matrix(X)

  Treatment <- base[[treatment_name]]

  # Create Folds
  set.seed(seed = seed)
  folds <- sample(rep(1:n_of_folds, length.out = n_of_individuals))

  for (fold in 1:n_of_folds) {
    test_idx <- which(folds == fold)
    train_idx <- setdiff(1:n_of_individuals, test_idx)

    X_train <- X[train_idx, , drop = FALSE]
    T_train <- Treatment[train_idx]

    if (propensity_model == "probability_forest") {
      T_train <- as.factor(T_train)
    }

    X_test <- X[test_idx, , drop = FALSE]

    if (propensity_model == "probability_forest") {
      regf <- grf::probability_forest(X_train, T_train, num.trees = n_trees_propensity)
      preds <- stats::predict(regf, X_test)$predictions[, 2]
    } else {
      lin_model <- stats::glm(T_train ~ ., data = data.frame(X_train, T_train), family = stats::binomial)
      preds <- stats::predict(lin_model, newdata = data.frame(X_test), type = "response")
    }
    propensity_scores[test_idx] <- preds
  }

  return(propensity_scores)
}
