#' Win/Loss Ratios with Distributional Random Forests
#'
#' `drf_win_ratio()` Computes the win/loss proportion, win ratio, and net benefit
#' by estimating nuisance parameters with machine learning methods.
#' This estimation is performed via distributional regression, implemented using distributional random forests.
#' @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 n_trees_drf Number of trees grown in the forest.
#' 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`.
#'
#' @returns A list of four 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 nuisance function values evaluated for each individual.
#' @export
#'
#' @examples
#' ## Mixed Covariates ##
#'
#' 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")
#' res <- drf_win_ratio(
#'   base = base,
#'   treatment_name = treatment_name,
#'   outcome_names = outcome_names
#' )
#' ## Categorical Covariates ##
#'
#' set.seed(123)
#' n <- 100
#' base <- data.frame(
#'   X1 = factor(sample(c("A", "B", "C"), n, replace = TRUE)), # factor covariate
#'   X2 = factor(sample(c("A", "B", "C"), n, replace = TRUE)), # factor covariate
#'   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")
#' res <- drf_win_ratio(
#'   base = base,
#'   treatment_name = treatment_name,
#'   outcome_names = outcome_names
#' )
drf_win_ratio <- function(base,
                          treatment_name,
                          outcome_names,
                          win_function = "lexicographic",
                          n_trees_drf = 1000,
                          n_of_folds = 3,
                          seed = NULL) {
  # Check Input & Extract Relevant Variables --------------------------------
  validate_input(
    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
  )

  covariate_names <- setdiff(colnames(base), c(outcome_names, treatment_name))

  if (is.character(win_function) && length(win_function) == 1 && win_function == "lexicographic") {
    win_function <- build_default_win_function(outcome_names)
  }

  lose_function <- function(treated_row, control_row) {
    win_function(control_row, treated_row)
  }

  n_of_individuals <- nrow(base)
  win_vector <- numeric(n_of_individuals)
  lose_vector <- numeric(n_of_individuals)

  # Cross-fitting Loop ------------------------------------------------------

  # 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 <- stats::model.frame(~., data = base[train_idx, covariate_names, drop = FALSE], na.action = stats::na.pass)
    X_train <- stats::model.matrix(~ . - 1, X_train)
    Y_train <- base[train_idx, outcome_names, drop = FALSE]
    T_train <- base[train_idx, treatment_name, drop = FALSE]

    X_control_train <- X_train[T_train == 0, , drop = FALSE]
    Y_control_train <- Y_train[T_train == 0, , drop = FALSE]

    X_treated_train <- X_train[T_train == 1, , drop = FALSE]
    Y_treated_train <- Y_train[T_train == 1, , drop = FALSE]

    X_test <- stats::model.frame(~., data = base[test_idx, covariate_names, drop = FALSE], na.action = stats::na.pass)
    X_test <- stats::model.matrix(~ . - 1, X_test)

    # Train the Distributional Random Forests
    drf_control <- drf::drf(Y = Y_control_train, X = X_control_train, num.trees = n_trees_drf)
    drf_treated <- drf::drf(Y = Y_treated_train, X = X_treated_train, num.trees = n_trees_drf)

    # Compute the DRF weights
    weights_control <- stats::predict(drf_control, X_test, type = "weights")$weights
    weights_treated <- stats::predict(drf_treated, X_test, type = "weights")$weights

    len_test <- length(test_idx)

    for (l in seq_len(len_test)) {
      individual_j <- test_idx[l]
      group_of_j <- base[individual_j, treatment_name, drop = FALSE]
      outcome_of_j <- base[individual_j, outcome_names, drop = FALSE]

      # Select appropriate training outcomes and weights
      if (group_of_j == 1) {
        Y <- Y_control_train
        w <- weights_control[l, ]
        win_args <- function(i) win_function(outcome_of_j, Y[i, , drop = FALSE])
        lose_args <- function(i) lose_function(outcome_of_j, Y[i, , drop = FALSE])
      } else {
        Y <- Y_treated_train
        w <- weights_treated[l, ]
        win_args <- function(i) win_function(Y[i, , drop = FALSE], outcome_of_j)
        lose_args <- function(i) lose_function(Y[i, , drop = FALSE], outcome_of_j)
      }

      # Compute win/lose vectors
      pjwin <- vapply(seq_len(nrow(Y)), win_args, FUN.VALUE = 0)
      pjlose <- vapply(seq_len(nrow(Y)), lose_args, FUN.VALUE = 0)

      win_vector[individual_j] <- sum(w * pjwin)
      lose_vector[individual_j] <- sum(w * pjlose)
    }
  }

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

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

  return(res)
}
