R/match_gps.R

In vecmatch: Generalized Propensity Score Estimation and Matching for Multiple Groups

#' @title Match the Data Based on Generalized Propensity Scores
#'
#' @description The `match_gps()` function performs sample matching based on
#' generalized propensity scores (GPS). It applies a k-means clustering
#' step to the GPS to partition the data into clusters and then matches
#' all treatment groups within each cluster. This yields balanced
#' comparisons across treatment levels while preserving the GPS
#' structure. To our knowledge, `vecmatch` provides the first
#' implementation in R of the vector-matching algorithm described by
#' Lopez and Gutman (2017).
#'
#' @param csmatrix An object of class `gps` and/or `csr` representing a data
#'   frame of generalized propensity scores. The first column must be the
#'   treatment variable, with additional attributes describing the calculation
#'   of the common support region and the estimation of generalized propensity
#'   scores. It is crucial that the common support region was calculated using
#'   the `csregion()` function to ensure compatibility.
#' @param method A single string specifying the matching method to use. The
#'   default is `"nnm"`, which applies the k-nearest neighbors matching
#'   algorithm. See the Details section for a full list of available methods.
#' @param caliper A numeric value specifying the caliper width, which defines
#'   the allowable range within which observations can be matched. It is
#'   expressed as a percentage of the standard deviation of the
#'   logit-transformed generalized propensity scores. To perform matching
#'   without a caliper, set this parameter to a very large value. For exact
#'   matching, set `caliper = 0` and enable the `exact` option by setting it to
#'   `TRUE`.
#' @param reference A single string specifying the exact level of the treatment
#'   variable to be used as the reference in the matching process. All other
#'   treatment levels will be matched to this reference level. Ideally, this
#'   should be the control level. If no natural control is present, avoid
#'   selecting a level with extremely low or high covariate or propensity score
#'   values. Instead, choose a level with covariate or propensity score
#'   distributions that are centrally positioned among all treatment groups to
#'   maximize the number of matches.
#' @param ratio A scalar for the number of matches which should be found for
#'   each control observation. The default is one-to-one matching. Only
#'   available for the methods `"nnm"` and `"pairopt"`.
#' @param replace Logical value indicating whether matching should be done with
#'   replacement. If `FALSE`, the order of matches generally matters. Matches
#'   are found in the same order as the data is sorted. Specifically, the
#'   matches for the first observation will be found first, followed by those
#'   for the second observation, and so on. Matching without replacement is
#'   generally not recommended as it tends to increase bias. However, in cases
#'   where the dataset is large and there are many potential matches, setting
#'   `replace = FALSE` often results in a substantial speedup with negligible or
#'   no bias. Only available for the method `"nnm"`
#' @param order A string specifying the order in which logit-transformed GPS
#'   values are sorted before matching. The available options are:
#'  * `"desc"` – sorts GPS values from highest to lowest (default).
#'  * `"asc"` – sorts GPS values from lowest to highest.
#'  * `"original"` – preserves the original order of GPS values.
#'  * `"random"` – randomly shuffles GPS values. To generate different random
#'  orders, set a seed using [set.seed()].
#' @param ties A logical flag indicating how tied matches should be handled.
#'   Available only for the `"nnm"` method, with a default value of `FALSE` (all
#'   tied matches are included in the final dataset, but only unique
#'   observations are retained). For more details, see the `ties` argument in
#'   [Matching::Matchby()].
#' @param min_controls The minimum number of treatment observations that should
#'   be matched to each control observation. Available only for the `"fullopt"`
#'   method. For more details, see the `min.controls` argument in
#'   [optmatch::fullmatch()].
#' @param max_controls The maximum number of treatment observations that can be
#'   matched to each control observation. Available only for the `"fullopt"`
#'   method. For more details, see the `max.controls` argument in
#'   [optmatch::fullmatch()].
#' @param kmeans_args A list of arguments to pass to [stats::kmeans]. These
#'   arguments must be provided inside a `list()` in the paired `name = value`
#'   format.
#' @param kmeans_cluster An integer specifying the number of clusters to pass to
#'   [stats::kmeans].
#' @param verbose_output a logical flag. If `TRUE` a more verbose version of the
#'   function is run and the output is printed out to the console.
#' @param ... Additional arguments to be passed to the matching
#'   function.
#'
#' @details Propensity score matching can be performed using various matching
#'   algorithms. Lopez and Gutman (2017) do not explicitly specify the matching
#'   algorithm used, but it is assumed they applied the commonly used k-nearest
#'   neighbors matching algorithm, implemented as `method = "nnm"`. However,
#'   this algorithm can sometimes be challenging to use, especially when
#'   treatment and control groups have unequal sizes. When `replace = FALSE`,
#'   the number of matches is strictly limited by the smaller group, and even
#'   with `replace = TRUE`, the results may not always be satisfactory. To
#'   address these limitations, we have implemented an additional matching
#'   algorithm to maximize the number of matched observations within a dataset.
#'
#'   The available matching methods are:
#'
#'   * `"nnm"` – classic k-nearest neighbors matching, implemented using
#'   [Matching::Matchby()]. The tunable parameters in `match_gps()` are
#'   `caliper`, `ratio`, `replace`, `order`, and `ties`. Additional arguments
#'   can be passed to [Matching::Matchby()] via the `...` argument.
#'   * `"fullopt"` – optimal full matching algorithm, implemented with
#'   [optmatch::fullmatch()]. This method calculates a discrepancy matrix to
#'   identify all possible matches, often optimizing the percentage of matched
#'   observations. The available tuning parameters are `caliper`,
#'   `min_controls`, and `max_controls`.
#'   * `"pairmatch"` – optimal 1:1 and 1:k matching algorithm, implemented using
#'   [optmatch::pairmatch()], which is actually a wrapper around
#'   [optmatch::fullmatch()]. Like `"fullopt"`, this method calculates a
#'   discrepancy matrix and finds matches that minimize its sum. The available
#'   tuning parameters are `caliper` and `ratio`.
#'
#'   Each method only accepts the tuning parameters listed above. Arguments
#'   that do not apply to the selected `method` are dropped, and a warning
#'   naming them is emitted only when they were explicitly supplied. Tuning
#'   parameters that are left at their default are filled in silently, so a
#'   call relying on the defaults runs without any warnings.
#'
#' @return A `data.frame` similar to the one provided as the `data` argument in
#'   the [estimate_gps()] function, containing the same columns but only the
#'   observations for which a match was found. The returned object includes two
#'   attributes, accessible with the `attr()` function:
#' * `original_data`: A `data.frame` with the original data returned by the
#'   [csregion()] or [estimate_gps()] function, after the estimation of the csr
#'   and filtering out observations not within the csr.
#'
#' * `matching_filter`: A logical vector indicating which rows from
#'   `original_data` were included in the final matched dataset.
#'
#' @references Michael J. Lopez, Roee Gutman "Estimation of Causal Effects with
#' Multiple Treatments: A Review and New Ideas," Statistical Science, Statist.
#' Sci. 32(3), 432-454, (August 2017)
#'
#' @seealso [estimate_gps()] for the calculation of generalized propensity
#'   scores; [MatchIt::matchit()], [optmatch::fullmatch()] and
#'   [optmatch::pairmatch()] for the documentation of the matching functions;
#'   [stats::kmeans()] for the documentation of the k-Means algorithm.
#'
#' @examples
#'
#' # Step 1.) Estimation of the generalized propensity scores
#' gp_scores <- estimate_gps(
#'   formula(status ~ age * sex),
#'   data = cancer,
#'   method = "multinom",
#'   reference = "control",
#'   verbose_output = TRUE
#' )
#'
#' # Step 2.) Defining the common support region
#' gps_csr <- csregion(gp_scores)
#'
#' # Example 1: optimal full matching ("fullopt") with explicit tuning params
#' matched_cancer_full <- match_gps(
#'   gps_csr,
#'   caliper = 0.25,
#'   reference = "control",
#'   method = "fullopt",
#'   kmeans_cluster = 2,
#'   kmeans_args = list(
#'     iter.max = 200,
#'     algorithm = "Forgy"
#'   ),
#'   min_controls = 0,
#'   max_controls = Inf,
#'   order = "desc",
#'   verbose_output = TRUE
#' )
#'
#' # Example 2: nearest neighbour matching ("nnm") demonstrating ratio /
#' # replace / order / ties
#' matched_cancer_nnm <- match_gps(
#'   gps_csr,
#'   caliper = 0.25,
#'   reference = "control",
#'   method = "nnm",
#'   ratio = 2,
#'   replace = TRUE,
#'   order = "asc",
#'   ties = TRUE,
#'   verbose_output = FALSE
#' )
#' @export
match_gps <- function(csmatrix = NULL,
                      method = "nnm",
                      caliper = 0.2,
                      reference = NULL,
                      ratio = NULL,
                      replace = NULL,
                      order = NULL,
                      ties = NULL,
                      min_controls = NULL,
                      max_controls = NULL,
                      kmeans_args = NULL,
                      kmeans_cluster = 5,
                      verbose_output = FALSE,
                      ...) {
  ####################### INPUT CHECKING #######################################
  ########################### AND ##############################################
  ####################### DATA PROCESSING ######################################

  # capture arguments
  all_args <- as.list(environment())
  ellipsis_args <- list(...)
  kmeans_args <- kmeans_args %||% list()
  args <- list()

  # define the number of unique treatment groups
  n_treat <- length(unique(csmatrix[, 1]))

  # replace the method to "nnm_2t" - uses Matching::Match()
  # [Matching:Matchby() can not be used, as it needs the "by" argument;
  # it cannot be bypassed by passing a vector of 1s]
  if (method == "nnm" & n_treat == 2) {
    method <- "nnm_2t"
  }

  # check and process the csmatrix
  .chk_cond(
    is.null(csmatrix) || missing(csmatrix),
    "The argument `csmatrix` is missing with no default."
  )

  .chk_cond(
    any(c("gps", "data.frame") %nin% class(csmatrix)),
    "The argument `csmatrix` has to be of classes `gps` and
            `data.frame`."
  )

  .chk_cond(
    n_treat < 2,
    "The number of unique treatments in the `csmatrix` must be higher
            or equal to 2."
  )

  # Perform the logit transformation and combine with treatment
  logit_matrix <- cbind(treatment = csmatrix[, 1], logit(csmatrix[, -1]))

  # defining the main data.frame
  data_name <- .match_methods[[method]]$data_name
  args[[data_name]] <- logit_matrix

  # Reset the rows counter on the logit_matrix
  rownames(args[[data_name]]) <- NULL

  # Control the method argument
  .chk_cond(
    method %nin% names(.match_methods),
    sprintf(
      "The argument `method` is only allowed to have following
             values: %s",
      word_list(names(.match_methods), quotes = TRUE)
    )
  )

  # Process the variable args for different methods
  varargs <- c(
    "caliper", "ratio", "replace", "order",
    "ties", "min_controls", "max_controls"
  )

  # define which varags are possible for given methods (both named args and ...)
  allowed_varargs <- unlist(
    lapply(
      .match_methods[[method]]$args_check_fun,
      function(x) {
        forms <- formals(x)
        forms[unlist(lapply(forms, is.null))] <- NA
        names(forms)
      }
    )
  )

  allowed_varargs <- c(
    .match_methods[[method]]$allowed_args,
    unique(allowed_varargs)
  )

  # check which varargs were specified (both names and ...)
  specified_varargs <- Filter(Negate(is.null), all_args)
  varargs_filter <- names(specified_varargs) %in% varargs
  specified_varargs <- c(
    specified_varargs[varargs_filter],
    ellipsis_args
  )

  # filter out only allowed args from the specified ones
  checked_varargs_filter <- names(specified_varargs) %in% allowed_varargs

  # if not allowed varargs were specified:
  .chk_cond(
    !all(checked_varargs_filter),
    error = FALSE,
    sprintf(
      "Following specified arguments are not
                    allowed for the method %s and will be ignored: %s",
      add_quotes(method),
      word_list(
        names(specified_varargs[!checked_varargs_filter])
      )
    )
  )

  # reference processing
  csmatrix$treatment <- as.factor(csmatrix$treatment)
  ref_list <- .process_ref(csmatrix$treatment,
    ordinal_treat = NULL,
    reference = reference
  )

  if (.match_methods[[method]]$treat_var) {
    args[["Tr"]] <- ref_list[["data.relevel"]]
  }

  reference <- ref_list[["reference"]]

  # kmeans_args, check list, process later
  if (!is.null(kmeans_args)) {
    chk::chk_list(kmeans_args)
  }

  # process combos
  if ("combos" %nin% names(ellipsis_args)) {
    # generate all possible matches with reference on the left
    combos <- expand.grid(group1 = reference, group2 = colnames(csmatrix)[-1])
    combos <- combos[combos[, 2] != reference, ]
  } else {
    # extracting combos
    combos <- ellipsis_args[["combos"]]

    # check if data frame
    .check_df(combos, data_name = "combos")

    # check if two cols
    .chk_cond(
      ncol(combos) != 2,
      "The `combos` data frame must have exactly 2 columns."
    )

    # check if all values in colnames(csmatrix)
    vectorized_combos <- as.character(c(combos[, 1], combos[, 2]))
    .chk_cond(
      any(vectorized_combos %nin% colnames(csmatrix)[-1]),
      "All values in the `combos` table must match the names of the
              unique levels of the treatment variable or the column names of
              the `csmatrix`."
    )

    # check if no repeats (e.g. 1, 1)
    combos_equal <- combos[, 1] == combos[, 2]
    .chk_cond(
      any(combos_equal),
      sprintf(
        "You can not match a group to itself, rows: %s",
        word_list(which(combos_equal))
      )
    )

    # check if all combinations unique
    # Sort each row and convert it into a string for comparison
    row_sorted <- apply(combos, 1, function(row) {
      paste(sort(row),
        collapse = ","
      )
    })

    # Check for duplicates
    duplicates <- duplicated(row_sorted)

    .chk_cond(
      any(duplicates),
      sprintf(
        "You can not check the same combination twice, rows: %s",
        word_list(which(duplicates))
      )
    )
  }

  # define combos
  colnames(combos) <- c("group1", "group2")

  combos[] <- lapply(combos, as.character)
  matches_n <- nrow(combos)

  ## varargs processing section ================================================
  # ratio
  if ("ratio" %in% allowed_varargs) {
    ratio <- ratio %||% 1

    .chk_vararg_length(ratio, "ratio",
      type_n = "integer",
      matches_n = matches_n
    )
    .check_integer(ratio, x_name = "ratio")

    args[["M"]] <- .vectorize(ratio, matches_n)
  }

  # replace
  if ("replace" %in% allowed_varargs) {
    replace <- replace %||% FALSE

    .chk_vararg_length(replace, "replace",
      type_n = "logical flag",
      matches_n = matches_n
    )

    .chk_cond(
      !is.logical(replace) || anyNA(replace),
      "All values in the `replace` argument have to be logical flags."
    )

    args[["replace"]] <- .vectorize(replace, matches_n)
  }

  # process caliper
  if ("caliper" %in% allowed_varargs) {
    caliper <- caliper %||% 0.25

    .chk_vararg_length(caliper, "caliper",
      TRUE,
      type_n = "numeric",
      matches_n
    )

    .chk_cond(
      any(caliper <= 0),
      "The `caliper` argument has to be a positive number."
    )

    caliper <- caliper * stats::sd(as.matrix(args[[data_name]][, reference]))

    args[["caliper"]] <- .vectorize(caliper, matches_n)
  }

  # define the matching formula based on combos
  if (.match_methods[[method]]$formula_necessary) {
    matching_formulas <- apply(
      combos,
      1,
      function(row) {
        paste0("treatment", " ~ ", row[1])
      }
    )
    args[["x"]] <- unlist(matching_formulas)
  }

  # processing the ties argument
  if ("ties" %in% allowed_varargs) {
    ties <- ties %||% TRUE

    .chk_vararg_length(ties, "ties",
      type_n = "logical flag",
      matches_n = matches_n
    )

    .chk_cond(
      !is.logical(ties) || anyNA(ties),
      "All values in the `ties` argument have to be logical flags."
    )

    args[["ties"]] <- .vectorize(ties, matches_n)
  }

  # processing the order argument
  if ("order" %in% allowed_varargs) {
    order <- order %||% "desc"

    .chk_vararg_length(order, "order",
      type_n = "text string",
      matches_n = matches_n
    )

    allowed_orders <- c("desc", "asc", "original", "random")

    .chk_cond(
      order %nin% allowed_orders,
      sprintf(
        "The `order` argument has to be one of the following values: %s",
        word_list(allowed_orders, quotes = TRUE)
      )
    )

    sort_before_matching <- .vectorize(order, matches_n)
  }

  # processing the min_controls
  if ("min_controls" %in% allowed_varargs) {
    min_controls <- min_controls %||% 0

    .chk_vararg_length(min_controls, "min_controls",
      TRUE,
      type_n = "numeric",
      matches_n
    )

    .chk_cond(
      any(min_controls < 0),
      "The `min_controls` argument has to be a non-negative finite number."
    )

    args[["min.controls"]] <- .vectorize(min_controls, matches_n)
  }

  # processing the max_controls
  if ("max_controls" %in% allowed_varargs) {
    max_controls <- max_controls %||% Inf

    .chk_vararg_length(max_controls, "max_controls",
      TRUE,
      type_n = "numeric",
      matches_n
    )

    .chk_cond(
      any(max_controls <= 0),
      "The `max_controls` argument has to be a non-negative number."
    )

    args[["max.controls"]] <- .vectorize(max_controls, matches_n)
  }

  # process kmeans_cluster
  .chk_cond(
    is.null(kmeans_cluster),
    "The `kmeans_cluster` argument can not be NULL."
  )

  .chk_vararg_length(kmeans_cluster, "kmeans_cluster",
    TRUE,
    type_n = "integer",
    matches_n
  )

  .check_integer(kmeans_cluster, x_name = "kmeans_cluster")

  .chk_cond(
    any(kmeans_cluster < 1),
    "The `kmeans_cluster` argument must be an integer
            greater than or equal 1."
  )

  kmeans_args[["centers"]] <- .vectorize(kmeans_cluster, matches_n)

  ## deal with algorithm argument
  if (is.null(kmeans_args[["algorithm"]])) {
    kmeans_args["algorithm"] <- "Hartigan-Wong"
  }

  ## processing the kmeans arglist
  kmeans_args <- match_add_args(
    arglist = kmeans_args,
    funlist = stats::kmeans
  )

  ## vectorize the args
  kmeans_args <- lapply(kmeans_args, .vectorize, matches_n)

  ## change ratio to controls in "pairopt"
  if (method == "pairopt") names(args)[names(args) == "M"] <- "controls"

  ## processing the matching arglist
  args <- match_add_args(
    arglist = args,
    funlist = .match_methods[[method]]$args_check_fun
  )

  if (method == "nnm" || method == "nnm_2t") {
    args <- args[names(args) %nin% c("Z", "V", "tolerance")]
  }

  args <- Filter(function(x) !all(is.na(x)), args)

  ## vectorize args
  args <- lapply(args, .vectorize, matches_n)

  ######################## MATCHING ############################################
  match_results <- list()

  for (i in seq_len(matches_n)) {
    ## select only treatment which in current combos loop
    obs_filter <- args[[data_name]][, "treatment"] %in% as.vector(combos[i, ])

    # selecting elements from arguments list based on current iteration
    kmeans_args_loop <- lapply(kmeans_args, function(x) x[[i]])

    # doing the same for args, but without data args
    args_length <- lengths(args)
    args_length[names(args_length) %in% c("Tr", "data", "X")] <- 0

    args_loop <- mapply(
      function(x, y) {
        if (x == matches_n) {
          y[[i]]
        } else {
          y
        }
      }, args_length,
      args,
      SIMPLIFY = FALSE
    )

    # selecting columns for kmeans clustering
    cols_kmeans <- colnames(args[[data_name]])[-1]
    cols_kmeans <- cols_kmeans[cols_kmeans %nin% as.vector(combos[i, ])]

    # selecting from df
    kmeans_args_loop[["x"]] <- args[[data_name]][, cols_kmeans]


    # when only 2 levels, then cols_kmeans is empty and kmeans throws an error
    if (n_treat > 2) {
      # fitting kmeans clusters
      tryCatch(
        {
          verbosely(
            withr::with_preserve_seed(
              k_res <- do.call(
                stats::kmeans,
                kmeans_args_loop
              )
            ),
            verbose = verbose_output
          )
        },
        error = function(e) {
          chk::abort_chk(strwrap(sprintf(
            "There was a problem fitting the kmeans clustering with
        `stats::kmeans()`.\n Error message: (from `stats::kmeans()`) %s",
            conditionMessage(e)
          ), prefix = " ", initial = ""), tidy = FALSE)
        }
      )

      # adding the clusters to matching arguments
      kmeans_strata <- as.factor(k_res$cluster[obs_filter])
    }

    # selecting columns for matching --> we need only one column of gps!
    cols_matching <- colnames(args[[data_name]]) %in% c(
      "treatment",
      combos[i, 1]
    )

    # arg processing for Matching::Matchby

    if (method == "nnm" || method == "nnm_2t") {
      # filtering the data frame for matching; only the gps column of the
      # first group of the current combination is used as the covariate
      args_loop[[data_name]] <- args_loop[[data_name]][
        obs_filter,
        combos[i, 1],
        drop = FALSE
      ]

      # defining data order
      order_data <- .ordering_func(args_loop[[data_name]][, combos[i, 1]],
        order = sort_before_matching[i]
      )
      order_original <- seq_len(nrow(args_loop[[data_name]]))
      order_original <- order_original[order_data]

      # reordering the data
      args_loop[[data_name]] <- args_loop[[data_name]][order_data, ,
        drop = FALSE
      ]

      # adding the clusters to matching arguments
      if (n_treat > 2) {
        args_loop[["by"]] <- kmeans_strata
        args_loop[["by"]] <- args_loop[["by"]][order_data]
      } else {
        # if 2 treatments - remove "by"
        args_loop[["by"]] <- NULL
      }

      # defining Tr args for matching function
      args_loop[["Tr"]] <- args_loop[["Tr"]][obs_filter]
      args_loop[["Tr"]] <- ifelse(args_loop[["Tr"]] == combos[i, 1],
        1,
        0
      )
      args_loop[["Tr"]] <- args_loop[["Tr"]][order_data]
    } else if (method %in% c("fullopt", "pairopt")) {
      # selecting observations and adding "by" argument from kmeans
      args_loop[[data_name]] <- args[[data_name]][obs_filter, cols_matching]

      # if more than 2 treatments
      if (n_treat > 2) {
        args_loop[[data_name]] <- cbind(args_loop[[data_name]],
          kmeans_strata = kmeans_strata
        )

        # converting the formula to add kmeans stratas
        args_loop[["x"]] <- paste0(
          args_loop[["x"]],
          " + optmatch::strata(kmeans_strata)"
        )
      }

      args_loop[[data_name]] <- as.data.frame(args_loop[[data_name]])

      # defining the order
      order_data <- .ordering_func(args_loop[[data_name]][, combos[i, 1]],
        order = order
      )
      order_original <- seq_len(nrow(args_loop[[data_name]]))
      order_original <- order_original[order_data]

      # reordering the data
      args_loop[[data_name]] <- args_loop[[data_name]][order_data, ]

      # converting treatment to binary
      args_loop[[data_name]][, "treatment"] <- ifelse(
        args_loop[[data_name]][, "treatment"] == combos[i, 1],
        0,
        1
      )

      # converting x to valid formula
      args_loop[["x"]] <- stats::as.formula(args_loop[["x"]])

      # performing the prematching =============================================
      tryCatch(
        {
          verbosely(
            withr::with_preserve_seed(
              prematched <- do.call(
                .match_methods[[method]]$args_check_fun[[1]],
                args_loop
              )
            ),
            verbose = verbose_output
          )
        },
        error = function(e) {
          chk::abort_chk(strwrap(sprintf(
            "There was a problem with matching the samples using
        `%s`.\n Error message:
        (from `%s`) %s",
            deparse(substitute(.match_methods[[method]]$args_check_fun[[1]])),
            deparse(substitute(.match_methods[[method]]$args_check_fun[[1]])),
            conditionMessage(e)
          ), prefix = " ", initial = ""), tidy = FALSE)
        }
      )

      # replacing args in the args_loop list for the subsequent matching
      args_loop[["x"]] <- prematched
    }

    # performing the matching ==================================================
    tryCatch(
      {
        verbosely(
          withr::with_preserve_seed(
            matched <- do.call(
              .match_methods[[method]]$matching_fun,
              args_loop
            )
          ),
          verbose = verbose_output
        )
      },
      error = function(e) {
        chk::abort_chk(strwrap(sprintf(
          "There was a problem with matching the samples using
        `%s`.\n Error message:
        (from `%s`) %s",
          deparse(substitute(.match_methods[[method]]$matching_fun)),
          deparse(substitute(.match_methods[[method]]$matching_fun)),
          conditionMessage(e)
        ), prefix = " ", initial = ""), tidy = FALSE)
      }
    )

    ## when no matches found, return an error
    .chk_cond(
      all(is.na(matched)),
      "No matches found!"
    )

    ## ids of the matched samples
    if (method == "nnm" || method == "nnm_2t") {
      ids_filtered <- as.numeric(rownames(args_loop[[data_name]]))


      ids_matched <- data.frame(
        control = ids_filtered[matched$index.treated],
        treatment = ids_filtered[matched$index.control]
      )

      colnames(ids_matched) <- paste0("level_", combos[i, 1:2])
    } else if (method %in% c("fullopt", "pairopt")) {
      ids_matched <- stats::na.omit(matched)
    }

    match_results <- append(
      match_results,
      list(ids_matched)
    )
  }

  ## OUTPUT section ############################################################
  if (method == "nnm" || method == "nnm_2t") {
    # defining the name of control column
    control_name <- paste0("level_", reference)

    # extracting the control columns
    matched_ids <- lapply(match_results, function(x) {
      sort(unique(x[, control_name, drop = TRUE]))
    })

    common_controls <- Reduce(intersect, matched_ids)

    # extract the respective matches
    extracted_matches <- lapply(match_results, function(x) {
      filter_vector <- x[, control_name] %in% common_controls
      unique(x[filter_vector, 2, drop = TRUE])
    })

    extracted_matches <- append(extracted_matches, list(common_controls))

    # merging all extracted ids
    all_extracted_ids <- Reduce(c, extracted_matches)
  } else if (method %in% c("fullopt", "pairopt")) {
    # extract only the ids of matched samples for common control identification
    matched_ids <- lapply(match_results, function(x) as.numeric(names(x)))

    # define the numbers of controls in logit_matrix
    control_range <- which(args[[data_name]]$treatment == reference)

    # extract the common controls
    matched_ids <- append(matched_ids, list(control_range))
    common_controls <- Reduce(intersect, matched_ids)

    # extract matched groups of controls
    common_controls_groups <- lapply(
      match_results,
      function(x) {
        filter_ids <- names(x) %in% common_controls
        x[filter_ids]
      }
    )

    # remove controls from match_results
    match_results <- lapply(
      match_results,
      function(x) {
        filter_ids <- names(x) %in% control_range
        x[!filter_ids]
      }
    )

    # extract matches with same groups
    extracted_matches <- mapply(function(x, y) y[y %in% x],
      common_controls_groups,
      match_results,
      SIMPLIFY = FALSE
    )

    extracted_matches <- lapply(
      extracted_matches,
      function(x) as.numeric(names(x))
    )

    extracted_matches <- append(extracted_matches, list(common_controls))

    # merging all extracted ids
    all_extracted_ids <- Reduce(c, extracted_matches)
  }

  # returning logical vector: which rows of the original csmatrix are matched
  matched_filter <- seq_len(nrow(csmatrix)) %in% all_extracted_ids

  # figure out which "original" data we should use as the base
  # - if csmatrix is a csr object: original_data = csr_data
  # (i.e. data within CSR)
  # - otherwise: original_data = original_data from estimate_gps()
  if (inherits(csmatrix, "csr")) {
    original_data <- attr(csmatrix, "csr_data")
  } else {
    original_data <- attr(csmatrix, "original_data")
  }

  # subset to matched rows
  matched_data <- as.data.frame(original_data[matched_filter, , drop = FALSE])

  ## ---- extract treatment variable name from function call -------------
  # function_call on csmatrix
  est_call <- attr(csmatrix, "function_call")
  treatment_var <- NULL

  if (!is.null(est_call)) {
    # try to find the formula in the call
    form <- est_call[["formula"]]

    if (inherits(form, "formula") && length(form) >= 3L) {
      # LHS is the treatment variable
      # e.g. status ~ age * sex  -> "status"
      treatment_var <- as.character(form[[2L]])
    }
  }

  # final fallback: if something went wrong, use first column name of
  # original_data
  if (is.null(treatment_var) || !treatment_var %in% names(original_data)) {
    treatment_var <- names(original_data)[1L]
  }

  # assemble result object: attributes + class in ONE place
  results <- structure(
    matched_data,
    original_data   = original_data,
    matching_filter = matched_filter,
    function_call   = match.call(),
    treatment_var   = treatment_var,
    class           = c("matched", "data.frame")
  )

  return(results)
}

#' @export
str.matched <- function(object, ...) {
  original_data <- attr(object, "original_data")
  matching_filter <- attr(object, "matching_filter")

  n_after <- nrow(object)
  n_before <- if (!is.null(original_data)) nrow(original_data) else NA_integer_

  cat(sprintf(
    "Matched data.frame: %d rows (from %s rows in original_data)\n",
    n_after,
    if (!is.na(n_before)) n_before else "NA"
  ))

  if (!is.null(matching_filter) && !is.na(n_before) &&
    length(matching_filter) == n_before) {
    cat(sprintf(
      "Proportion of original units kept: %.1f%%\n",
      100 * mean(matching_filter)
    ))
  }

  # show underlying structure
  utils::str(unclass(object), ...)
  invisible(object)
}

# internal helper: print matched data.frame with a header
.print_matched_core <- function(x, ...) {
  n_after <- nrow(x)

  treatment <- x[["treatment"]]
  trt_levels <- NULL
  k <- NA_integer_

  original_data <- attr(x, "original_data")
  matching_filter <- attr(x, "matching_filter")

  n_before <- if (!is.null(original_data)) nrow(original_data) else NA_integer_
  prop_kept <- if (!is.null(matching_filter) &&
    !is.na(n_before) &&
    length(matching_filter) == n_before) {
    sprintf("%.1f%%", 100 * mean(matching_filter))
  } else {
    NA_character_
  }

  cli::cli_ul()
  cli::cli_li("Number of units after matching: {n_after}")
  if (!is.na(n_before)) {
    cli::cli_li("Number of units in original_data: {n_before}")
  }
  if (!is.na(prop_kept)) {
    cli::cli_li("Proportion of original units kept: {prop_kept}")
  }

  if (!is.null(treatment)) {
    cli::cli_li("Treatment column: {.field treatment}")
    cli::cli_li("Number of treatment levels: {k}")
    cli::cli_li("Treatment levels: {paste(trt_levels, collapse = ', ')}")
  } else {
    cli::cli_li("Treatment column not found in the matched data.")
  }

  cli::cli_li("Underlying classes: {.field {paste(class(x), collapse = ', ')}}")
  cli::cli_end()

  cli::cli_text("")

  # print underlying data.frame directly
  base::print.data.frame(x, ...)

  invisible(x)
}

#' @export
print.matched <- function(x, ...) {
  cli::cli_text("{.strong matched object} (GPS-based matched dataset)")
  .print_matched_core(x, ...)
}

#' @export
summary.matched <- function(object, digits = 1, ...) {
  original_data <- attr(object, "original_data")
  matching_filter <- attr(object, "matching_filter")
  call <- attr(object, "function_call")

  n_total <- NROW(original_data)
  n_matched <- NROW(object)
  n_unmatched <- n_total - n_matched
  overall_retained <- if (n_total > 0) 100 * n_matched / n_total else NA_real_

  ## --- treatment variable name ----------------------------------------------
  treat_name <- attr(object, "treatment_var")

  ## --- per-treatment counts --------------------------------------------------
  tr_orig <- original_data[[treat_name]]
  tr_match <- object[[treat_name]]

  if (!is.factor(tr_orig)) {
    tr_orig <- factor(tr_orig)
  }
  if (!is.factor(tr_match)) {
    tr_match <- factor(tr_match, levels = levels(tr_orig))
  }

  tab_before <- table(tr_orig)
  tab_after <- table(tr_match)

  all_levels <- levels(tr_orig)

  n_before <- as.integer(tab_before[all_levels])
  n_before[is.na(n_before)] <- 0L

  n_after <- as.integer(tab_after[all_levels])
  n_after[is.na(n_after)] <- 0L

  retained_pct <- ifelse(n_before > 0, 100 * n_after / n_before, NA_real_)

  per_treatment <- data.frame(
    Treatment        = all_levels,
    n_before         = n_before,
    n_after          = n_after,
    Retained_percent = retained_pct,
    check.names      = FALSE,
    row.names        = NULL
  )

  ## --- extract matching parameters from match_gps() call ---------------------
  params <- NULL
  if (!is.null(call)) {
    call_list <- as.list(call)[-1] # drop function name
    call_list <- call_list[names(call_list) %nin% c("csmatrix")]
    params <- call_list
  }

  res <- list(
    call              = call,
    params            = params,
    n_total           = n_total,
    n_matched         = n_matched,
    n_unmatched       = n_unmatched,
    overall_retained  = overall_retained,
    per_treatment     = per_treatment,
    treatment_var     = treat_name,
    digits            = digits
  )

  class(res) <- "summary.matched"
  res
}

#' @export
print.summary.matched <- function(x, digits = NULL, ...) {
  # digits: explicit > stored > default
  if (is.null(digits)) {
    digits <- if (!is.null(x$digits)) x$digits else 1L
  }

  fmt_pct <- function(z) {
    ifelse(is.na(z), NA_character_, sprintf("%.*f%%", digits, z))
  }

  # Simple fixed-width table helper with safe character conversion
  print_table <- function(df) {
    df_chr <- data.frame(
      lapply(df, as.character),
      check.names = FALSE,
      row.names = NULL
    )
    colnames_df <- colnames(df_chr)

    cat("\n")
    cat(paste(sprintf("%-18s", colnames_df), collapse = " | "), "\n")
    cat(paste(rep("-", 20 * ncol(df_chr)), collapse = ""), "\n")

    apply(df_chr, 1, function(row) {
      cat(paste(sprintf("%-18s", row), collapse = " | "), "\n")
    })
    cat("\n")
  }

  cli::cli_h1("Summary of matched object")

  ## Matching settings ---------------------------------------------------------
  cli::cli_h2("Matching settings")

  params <- x$params
  if (!is.null(params) && length(params)) {
    cli::cli_ul()
    for (nm in names(params)) {
      val <- params[[nm]]
      val_str <- paste(deparse(val, width.cutoff = 60L), collapse = "")
      cli::cli_li("{.field {nm}} = {val_str}")
    }
    cli::cli_end()
  } else {
    cli::cli_text("{.emph No additional matching parameters recorded in
                  function_call.}")
  }

  ## Overall matching performance ---------------------------------------------
  cli::cli_h2("Overall matching performance")
  cli::cli_ul()
  cli::cli_li("Treatment variable: {.field {x$treatment_var}}")
  cli::cli_li("Total units in original_data: {x$n_total}")
  cli::cli_li("Units in matched data: {x$n_matched}")
  cli::cli_li("Units not matched: {x$n_unmatched}")
  cli::cli_li("Overall retention: {fmt_pct(x$overall_retained)}")
  cli::cli_end()

  ## Per-treatment retention ---------------------------------------------------
  pt <- x$per_treatment

  df <- data.frame(
    Treatment = pt$Treatment,
    "N before" = pt$n_before,
    "N after" = pt$n_after,
    "Retained (%)" = fmt_pct(pt$Retained_percent),
    check.names = FALSE,
    row.names = NULL
  )

  cli::cli_h2("Per-treatment retention")
  print_table(df)

  invisible(x)
}

#' @export
plot.matched <- function(x, ...) {
  # use summary.matched() programmatically
  s <- summary(x)
  pt <- s$per_treatment

  # prepare matrix: rows = before/after, cols = treatment levels
  counts_mat <- rbind(
    "Before matching" = pt$n_before,
    "After matching"  = pt$n_after
  )

  dots <- list(...)

  # default main/xlab/ylab if not provided in ...
  if (!"main" %in% names(dots)) {
    dots$main <- "Matched sample sizes by treatment"
  }
  if (!"xlab" %in% names(dots)) {
    dots$xlab <- sprintf("Treatment (%s)", s$treatment_var)
  }
  if (!"ylab" %in% names(dots)) {
    dots$ylab <- "Number of units"
  }
  if (!"names.arg" %in% names(dots)) {
    dots$names.arg <- pt$Treatment
  }
  if (!"ylim" %in% names(dots)) {
    max_y <- max(counts_mat, na.rm = TRUE)
    dots$ylim <- c(0, max_y * 1.1)
  }
  if (!"col" %in% names(dots)) {
    dots$col <- c("grey70", "white")
  }
  if (!"border" %in% names(dots)) {
    dots$border <- c("grey20", "black")
  }

  # always beside = TRUE for grouped bars
  dots$height <- counts_mat
  dots$beside <- TRUE

  old_par <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(old_par), add = TRUE)

  bar_pos <- do.call(graphics::barplot, dots)

  # legend in top-right by default
  graphics::legend(
    "topright",
    legend = rownames(counts_mat),
    fill   = if (!is.null(dots$col)) dots$col else c("grey70", "white"),
    border = if (!is.null(dots$border)) dots$border else c("grey20", "black"),
    bty    = "n"
  )

  invisible(x)
}

#' @export
as.data.frame.matched <- function(x, ...) {
  class(x) <- "data.frame"
  NextMethod()
}