#' Cross-validate the rfok sensitivity parameter
#'
#' Cross-validates the sensitivity parameter \code{beta} for an \code{rfok}
#' model over a user-supplied or automatically generated grid.
#'
#' @param X a data frame, matrix, or vector of predictors. A vector is coerced
#'   to a one-column matrix.
#' @param y a numeric response vector.
#' @param beta optional numeric vector of candidate beta values. If \code{NULL},
#'   candidates are generated as \code{sd(y) * beta_mult}.
#' @param beta_mult numeric vector used to generate the beta grid when
#'   \code{beta = NULL}.
#' @param k Number of nearest neighbors to use for out-of-bag error estimators.
#' @param num_folds Number of cross-validation folds.
#' @param selection_rule Criterion used to choose beta. Options are
#'   \code{"crps"}, \code{"interval_score"}, and \code{"coverage_width"}.
#' @param conf Confidence level used for interval-based selection rules.
#' @param samples Number of conformal samples used when
#'   \code{selection_rule = "crps"}.
#' @param smoothing logical; should conformity scores be sampled or smoothed
#'   when generating conformal samples?
#' @param eps Tolerance for nearest neighbor search.
#' @param return_fit logical; if \code{TRUE}, refit \code{rfok()} on the full
#'   data using the selected beta and include the fitted model in the result.
#' @param verbose logical. should progress be printed?
#' @param ... additional parameters passed to \code{rfok()}.
#'
#' @details
#' The parameter \code{beta} controls the sensitivity of the conformal interval
#' widths to the local out-of-bag residual scale estimate. Smaller values make
#' intervals more locally adaptive, while larger values make intervals more
#' stable and more similar to globally scaled conformal intervals.
#'
#' This function evaluates candidate \code{beta} values using cross-validation.
#' For \code{selection_rule = "crps"}, predictive samples are generated for each
#' held-out observation and the candidate with the smallest average continuous
#' ranked probability score is selected.
#'
#' For \code{selection_rule = "interval_score"}, conformal intervals are
#' constructed at level \code{conf} and the candidate with the smallest average
#' interval score is selected. The interval score rewards narrow intervals but
#' penalizes observations that fall outside the interval.
#'
#' For \code{selection_rule = "coverage_width"}, the selected candidate is the
#' one with the smallest average interval width among candidates with empirical
#' coverage at least \code{conf}. If no candidate reaches the desired coverage,
#' the candidate with the largest empirical coverage is selected, with ties
#' broken by average interval width.
#'
#' @return A data frame with one row per candidate beta. The selected beta is
#'   marked in the \code{selected} column. If \code{return_fit = TRUE}, a list is
#'   returned with elements \code{results}, \code{best_beta}, and \code{fit}.
#'
#' @export
rfok_cv_beta <- function(X, y,
                         beta = NULL,
                         beta_mult = c(0, 1 / 100, 1 / 50, 1 / 30, 1 / 20, 1 / 10),
                         k = 5,
                         num_folds = 5,
                         selection_rule = c("crps", "interval_score", "coverage_width"),
                         conf = 0.95,
                         samples = 1000,
                         smoothing = TRUE,
                         eps = 0,
                         return_fit = FALSE,
                         verbose = TRUE,
                         ...) {

  selection_rule <- match.arg(selection_rule)

  # Validate and coerce X
  if (is.null(X)) {
    stop("`X` must not be NULL.", call. = FALSE)
  }

  if (is.vector(X) && !is.list(X)) {
    X <- matrix(X, ncol = 1)
  }

  if (!is.matrix(X) && !is.data.frame(X)) {
    stop("`X` must be a matrix, data frame, or vector coercible to a one-column matrix.",
         call. = FALSE)
  }

  if (nrow(X) == 0) {
    stop("`X` must have at least one row.", call. = FALSE)
  }

  if (ncol(X) == 0) {
    stop("`X` must have at least one column.", call. = FALSE)
  }

  if (anyNA(X)) {
    stop("`X` must not contain missing values.", call. = FALSE)
  }

  # Validate and coerce y
  if (missing(y) || is.null(y)) {
    stop("`y` must be supplied.", call. = FALSE)
  }

  if (is.matrix(y) || is.data.frame(y)) {
    if (ncol(y) != 1) {
      stop("`y` must be a vector or a single-column matrix/data frame.",
           call. = FALSE)
    }
    y <- y[, 1]
  }

  if (!is.atomic(y) || is.list(y)) {
    stop("`y` must be an atomic response vector.", call. = FALSE)
  }

  if (length(y) != nrow(X)) {
    stop("`y` must have length equal to `nrow(X)`.", call. = FALSE)
  }

  if (!is.numeric(y)) {
    stop("`y` must be numeric.", call. = FALSE)
  }

  if (anyNA(y)) {
    stop("`y` must not contain missing values.", call. = FALSE)
  }

  n <- nrow(X)

  if (n < 2) {
    stop("Cross-validation requires at least two observations.", call. = FALSE)
  }

  # Validate k
  if (length(k) != 1 || is.na(k) || !is.numeric(k)) {
    stop("`k` must be a single non-missing number.", call. = FALSE)
  }

  if (k < 1 || k != floor(k)) {
    stop("`k` must be a whole number greater than or equal to 1.", call. = FALSE)
  }

  k <- as.integer(k)

  # Validate num_folds
  if (length(num_folds) != 1 || is.na(num_folds) || !is.numeric(num_folds)) {
    stop("`num_folds` must be a single non-missing number.", call. = FALSE)
  }

  if (num_folds < 2 || num_folds != floor(num_folds)) {
    stop("`num_folds` must be a whole number greater than or equal to 2.",
         call. = FALSE)
  }

  num_folds <- as.integer(num_folds)

  if (num_folds > n) {
    warning("`num_folds` is greater than `nrow(X)`; using leave-one-out CV.",
            call. = FALSE)
    num_folds <- n
  }

  # Validate conf
  if (length(conf) != 1 || is.na(conf) || !is.numeric(conf)) {
    stop("`conf` must be a single non-missing number.", call. = FALSE)
  }

  if (conf <= 0 || conf >= 1) {
    stop("`conf` must be between 0 and 1.", call. = FALSE)
  }

  # Validate samples
  if (selection_rule == "crps") {
    if (length(samples) != 1 || is.na(samples) || !is.numeric(samples)) {
      stop("`samples` must be a single non-missing number.", call. = FALSE)
    }

    if (samples < 1 || samples != floor(samples)) {
      stop("`samples` must be a whole number greater than or equal to 1.",
           call. = FALSE)
    }

    samples <- as.integer(samples)
  }

  # Validate smoothing
  if (length(smoothing) != 1 || is.na(smoothing) || !is.logical(smoothing)) {
    stop("`smoothing` must be either TRUE or FALSE.", call. = FALSE)
  }

  # Validate eps
  if (length(eps) != 1 || is.na(eps) || !is.numeric(eps)) {
    stop("`eps` must be a single non-missing number.", call. = FALSE)
  }

  if (eps < 0) {
    stop("`eps` must be non-negative.", call. = FALSE)
  }

  # Construct beta grid
  if (is.null(beta)) {
    if (!is.numeric(beta_mult) || length(beta_mult) == 0) {
      stop("`beta_mult` must be a non-empty numeric vector.", call. = FALSE)
    }

    if (anyNA(beta_mult)) {
      stop("`beta_mult` must not contain missing values.", call. = FALSE)
    }

    if (any(beta_mult < 0)) {
      stop("`beta_mult` must contain non-negative values.", call. = FALSE)
    }

    beta <- stats::sd(y) * beta_mult
  }

  if (!is.numeric(beta) || length(beta) == 0) {
    stop("`beta` must be a non-empty numeric vector.", call. = FALSE)
  }

  if (anyNA(beta)) {
    stop("`beta` must not contain missing values.", call. = FALSE)
  }

  if (any(beta < 0)) {
    stop("`beta` must contain non-negative values.", call. = FALSE)
  }

  beta <- sort(unique(beta))

  # Generate folds
  fold_id <- sample(rep(seq_len(num_folds), length.out = n))

  results <- data.frame(
    beta = beta,
    mean_crps = NA_real_,
    coverage = NA_real_,
    mean_width = NA_real_,
    mean_interval_score = NA_real_,
    score = NA_real_
  )

  .rfok_msg(verbose,
    "Cross-validating ", length(beta), " beta value",
    if (length(beta) == 1) "" else "s",
    " across ", num_folds, " folds using selection_rule = \"",
    selection_rule, "\"."
  )

  for (b in seq_along(beta)) {
    beta_b <- beta[b]

    .rfok_msg(verbose,
      "[", b, "/", length(beta), "] beta = ",
      signif(beta_b, 5)
    )

    crps_values <- numeric(0)
    covered <- logical(0)
    widths <- numeric(0)
    interval_scores <- numeric(0)

    for (fold in seq_len(num_folds)) {
      .rfok_msg(verbose, "  fold ", fold, "/", num_folds)

      test_ind <- which(fold_id == fold)
      train_ind <- setdiff(seq_len(n), test_ind)

      fit <- rfok(
        X = X[train_ind, , drop = FALSE],
        y = y[train_ind],
        k = min(k, length(train_ind)),
        beta = beta_b,
        eps = eps,
        ...
      )

      # Always compute interval diagnostics.
      pred_int <- predict(
        fit,
        newdata = X[test_ind, , drop = FALSE],
        samples = NULL,
        conf = conf,
        eps = eps
      )

      pred <- pred_int$pred
      moe <- pred_int$moe[, 1]

      lower <- pred - moe
      upper <- pred + moe

      covered_fold <- y[test_ind] >= lower & y[test_ind] <= upper
      width_fold <- upper - lower
      score_fold <- .rfok_interval_score(y[test_ind], lower, upper, conf)

      covered <- c(covered, covered_fold)
      widths <- c(widths, width_fold)
      interval_scores <- c(interval_scores, score_fold)

      # Compute CRPS only when needed.
      if (selection_rule == "crps") {
        pred_samples <- predict(
          fit,
          newdata = X[test_ind, , drop = FALSE],
          samples = samples,
          conf = NULL,
          smoothing = smoothing,
          eps = eps
        )

        crps_values <- c(
          crps_values,
          vapply(
            seq_along(test_ind),
            function(j) .rfok_crps_from_samples(
              pred_samples[, j],
              y[test_ind[j]]
            ),
            numeric(1)
          )
        )
      }
    }

    results$coverage[b] <- mean(covered)
    results$mean_width[b] <- mean(widths)
    results$mean_interval_score[b] <- mean(interval_scores)

    if (selection_rule == "crps") {
      best_ind <- which.min(results$mean_crps)
      results$mean_crps[b] <- mean(crps_values)
      results$score[b] <- -results$mean_crps[b]
    }

    if (selection_rule == "interval_score") {
      best_ind <- which.min(results$mean_interval_score)
      results$score[b] <- -results$mean_interval_score[b]
    }

    if (selection_rule == "coverage_width") {
      valid <- which(results$coverage >= conf)
      if (length(valid) > 0) {
        best_ind <- valid[which.min(results$mean_width[valid])]
      } else {
        max_cov <- max(results$coverage)
        valid <- which(results$coverage == max_cov)
        best_ind <- valid[which.min(results$mean_width[valid])]
      }
      results$score[b] <- NA_real_
    }

    .rfok_msg(verbose,
      "  coverage = ", signif(results$coverage[b], 5),
      ", mean width = ", signif(results$mean_width[b], 5),
      ", mean interval score = ", signif(results$mean_interval_score[b], 5),
      if (selection_rule == "crps") {
        paste0(", mean CRPS = ", signif(results$mean_crps[b], 5))
      } else {
        ""
      }
    )
  }

  results$selected <- FALSE
  results$selected[best_ind] <- TRUE

  best_beta <- results$beta[best_ind]

  if (return_fit) {
    final_fit <- rfok(
      X = X,
      y = y,
      k = k,
      beta = best_beta,
      eps = eps,
      ...
    )

    return(list(
      results = results,
      best_beta = best_beta,
      fit = final_fit
    ))
  }

  results
}

.rfok_crps_from_samples <- function(samples, y, w = 0) {
  x <- samples
  M <- length(x)
  if(M < 2){
    stop("More samples are needed to estimate CRPS.")
  }
  term1 <- mean(abs(x - y))
  x <- sort(x)
  coef <- 2 * seq_len(M) - M - 1
  term2 <- 2 * sum(coef * x)
  term1 - (1 - w / M) * term2 / (2 * M * (M - 1))
}


.rfok_interval_score <- function(y, lower, upper, conf) {
  alpha <- 1 - conf

  width <- upper - lower
  lower_penalty <- (2 / alpha) * (lower - y) * (y < lower)
  upper_penalty <- (2 / alpha) * (y - upper) * (y > upper)

  width + lower_penalty + upper_penalty
}


.rfok_msg <- function(verbose, ...) {
  if(verbose)
    message(...)
}
