#' Design crosses to produce high-performing and stable progeny
#' #'
#' Selects pairs of candidate parents by maximising an index that reflects
#' expected progeny performance for the main effect, with a penalty applied
#' for imbalance across rotated higher-order factors.
#' 
#' @param df A data frame containing individual IDs and rotated slopes. The first
#'   column must be named \code{id}, followed by slope columns of the form
#'   \code{slp1}, \code{slp2}, ..., \code{slpk}.
#' @param H_mat A matrix of corresponding rotated environmental covariates.
#' @param n_sel Number of individuals to select. Must be an even number because
#'   individuals are selected as pairs.
#' @param lambda A scalar giving the penalty weight applied to the stability term.
#'   Larger values place more emphasis on stability.
#' @param n_top Number of top-ranked individuals, based on the first slope term,
#'   to retain before forming candidate pairs.
#' @param n_balance_factors Number of higher-order factors used to calculate the
#'   stability penalty. If \code{NULL}, all available balance factors are used.
#' @param max_use Maximum number of times each individual can be used across
#'   selected pairs.
#'
#' @return A list containing:
#' \itemize{
#'   \item \code{selected_ids}: IDs of selected individuals.
#'   \item \code{selected_df}: Rows of \code{df} corresponding to selected individuals.
#'   \item \code{pair_ids}: A data frame containing the selected pairs.
#'   \item \code{pair_scores}: A data frame containing all candidate pairs and their
#'     main-effect scores, stability penalties, and total scores.
#' }
#'
#' @examples
#' # Select 100 individuals as 50 pairs using all available balance factors
#' ss <- dCross(
#'   df = slope_df,
#'   H_mat = H_mat,
#'   n_sel = 100,
#'   lambda = 1,
#'   n_top = 1000,
#'   max_use = 1
#' )
#'
#' # Examine selected pairs
#' ss$pair_ids
#'
#' # Examine scores for all candidate pairs
#' ss$pair_scores
#'
#' @export

dCross <- function(df, H_mat, n_sel = 100, lambda = 1, n_top = 1000,
                   n_balance_factors = NULL,
                   max_use = 1) {
  
  if (!"id" %in% colnames(df) || colnames(df)[1] != "id") {
    stop("Data frame must contain a column named 'id' in the first position.")
  }
  
  if (n_sel %% 2 != 0) {
    stop("Parameter 'n_sel' must be an even number, because individuals are paired.")
  }  
  
  n <- nrow(df)             # Number of candidates
  X <- as.matrix(df[, -1])  # drop id column 
  k <- ncol(X)              # total number of genetic values
  
  # 1. Select top individuals based on X1
  top_indices <- order(-X[, 1])[1:min(n_top, n)]
  df_top <- df[top_indices, ]
  X_top <- X[top_indices, ]
  n_top_actual <- nrow(df_top)
  
  # 2. Form all unique pairs
  pairs <- which(upper.tri(matrix(1, n_top_actual, n_top_actual)), arr.ind = TRUE)
  npairs <-  nrow(pairs)
  
  # 3. Consider only n_balance traits
  balance_factors <- X_top[, -1, drop = FALSE] 
  num_balance_factors <- ncol(balance_factors)
  
  if (is.null(n_balance_factors)) {
    n_balance_factors <- num_balance_factors
    message("Using all ", n_balance_factors, " balance trait(s).")
  } else if (n_balance_factors > num_balance_factors) {
    message("Only ", num_balance_factors, " balance trait(s) detected; changing n_balance_factors to ", num_balance_factors)
    n_balance_factors <- num_balance_factors
  }
  
  
  # 4. Calculate pair scores
  main_score <- numeric(npairs)
  stab <- numeric(npairs)
  
  X_avg_all <- (X_top[pairs[,1], , drop = FALSE] +
                  X_top[pairs[,2], , drop = FALSE]) / 2
  
  # Main effect (vector length npairs)
  main_weight <- mean(H_mat[, 1])
  main_score  <- X_avg_all[, 1] * main_weight
  
  # stability  penalty using the first n_balance_factors 
  if (n_balance_factors > 0) {
    X_avg_stab    <- X_avg_all[, 1 + seq_len(n_balance_factors), drop = FALSE]   # npairs x nBF
    H_mat_stab    <- H_mat[, 1 + seq_len(n_balance_factors),      drop = FALSE]   # nEnv x nBF
    proj         <- X_avg_stab %*% t(H_mat_stab)                                 # npairs x nEnv
    stab         <- sqrt(rowMeans(proj^2))
  } else {
    stab <- rep(0, nrow(X_avg_all))
  }
  
  total_score <- main_score - lambda * stab
  
  pair_scores <- data.frame(
    id1 = df_top$id[pairs[, 1]],
    id2 = df_top$id[pairs[, 2]],
    main_score = main_score,
    stab = stab,
    total_score = total_score
  )
  
  # 5. Greedy Pair Selection with max_use
  pair_order <- order(-total_score)
  selected_pairs <- list()
  used_ids <- setNames(rep(0, n_top_actual), df_top$id)  # usage count
  
  for (i in pair_order) {
    if (length(selected_pairs) >= n_sel / 2) break
    id1 <- df_top$id[pairs[i, 1]]
    id2 <- df_top$id[pairs[i, 2]]
    
    # Ensure usage limit is not exceeded
    if (used_ids[id1] < max_use && used_ids[id2] < max_use) {
      selected_pairs[[length(selected_pairs) + 1]] <- c(id1, id2)
      used_ids[id1] <- used_ids[id1] + 1
      used_ids[id2] <- used_ids[id2] + 1
    }
  }
  
  if (length(selected_pairs) < n_sel / 2) {
    warning("Could not fill all pairs due to filtering and max_use constraints.")
  }
  
  selected_pairs_df <- do.call(rbind, lapply(selected_pairs, function(x) data.frame(id1 = x[1], id2 = x[2])))
  selected_df <- df[df$id %in% names(used_ids)[used_ids > 0], ]
  
  list(
    selected_ids = names(used_ids)[used_ids > 0],
    selected_df = selected_df,
    pair_ids = selected_pairs_df,
    pair_scores = pair_scores
  )
}
