#' Check the sample distribution of the dataset and
#' give the suggestion if need the class weight or not
#'
#' @param labels The label for distribution
#' @param threshold The threshold for the ratio (0.5) if it is the
#' imbalanced dataset
#' @param plot Choose to have the figures or not
#' @param outdir The output directory where plots will be saved
#' (default: tempdir()).
#'
#' @importFrom tibble tibble
#'
#' @return A A \code{CrcBiomeScreen} object. object with updated slots.
#' @export
#'
#' @examples
#' # Small toy example for runnable demonstration
#' @examples
#' train_labels <- factor(c("control", "CRC", "control", "CRC"))
#' checkClassBalance(train_labels)

checkClassBalance <- function(labels,
                              outdir = tempdir(),
                              threshold = 0.5,
                              plot = TRUE) {
    # Convert to factor
    labels <- as.factor(labels)
    class_counts <- table(labels)
    class_props <- prop.table(class_counts)

    # Identify minority classes
    is_imbalanced <- any(class_props < threshold | class_props > (1 - threshold))
    suggestion <- if (is_imbalanced) {
      "Unbalanced: Using class weights is recommended."
    } else {
      "Balanced: Class weights are likely unnecessary."
    }

    # Optional plot
    if (plot) {
      pdf(file.path(outdir, "class_balance_plot.pdf"), width = 8, height = 6)
      barplot(class_counts,
        col = "steelblue",
        main = "Class Distribution",
        ylab = "Sample Count",
        xlab = "Class Label"
      )
      abline(h = mean(class_counts), col = "red", lty = 2)
      dev.off()
    }

    # Print summary
    message("Class counts:\n")
    message(class_counts)
    message("Class proportions:\n")
    message(round(class_props, 3))
    message("Balance status:\n", suggestion, "\n")

    # Return result
    return(list(
      class_counts = class_counts,
      class_proportions = class_props,
      is_imbalanced = is_imbalanced,
      suggestion = suggestion
    ))
}
