#' Conformal Random Forest with Out-of-bag KNN
#'
#' A conformal RF based on Johansson et al. (2014). Rather than using a traditional calibration set, conformal intervals are constructed using out-of-bag samples for each tree. K-nearest neighbors is used to calibrate the prediction intervals for each prediction location.
#'
#' @param X a data frame or a matrix of predictors
#' @param y a response vector
#' @param k Number of nearest neighbors to use for out-of-bag error estimaters
#' @param beta Sensitivity parameter
#' @param eps Tolerance for nearest neighbor approach. When \code{eps = 0}, exact nearest neighbors are used.
#' @param ... additional parameters passed to randomForest
#' @return An object with class "rfok"
#' @details
#' This function fits a random forest and constructs conformal prediction
#' intervals using out-of-bag (OOB) predictions. For each training observation,
#' an OOB prediction is computed using only trees for which that observation was
#' not included in the bootstrap sample. These OOB predictions are used to form
#' residuals without requiring a separate calibration set.
#'
#' Local scale estimates are computed using the \code{k} nearest neighbors in
#' the predictor space. For observation \eqn{i}, the local scale estimate
#' \eqn{\mu_i} is the average absolute OOB residual among its nearest neighbors.
#' The non-conformity score is
#'
#' \deqn{\alpha_i = \frac{|y_i - \hat{y}^{OOB}_i|}{\mu_i + \beta}.}
#'
#' The parameter \code{k} controls the locality of the scale estimate. Smaller
#' values make the intervals more locally adaptive, while larger values produce
#' smoother, more stable local scale estimates.
#'
#' The parameter \code{beta} is a non-negative stabilization parameter added to
#' the local scale estimate. Smaller values make the intervals more sensitive to
#' local changes in the OOB residuals. Larger values reduce this sensitivity and
#' make the intervals behave more like globally scaled conformal intervals. The
#' default, \code{sd(y) / 30}, is a scale-aware value intended to provide mild
#' stabilization.
#'
#' The parameter \code{eps} is passed to the nearest-neighbor search. Setting
#' \code{eps = 0} uses exact nearest neighbors. Larger values allow approximate
#' nearest-neighbor search and may be faster for large data sets.
#' @references Johansson, U., Boström, H., Löfström, T., & Linusson, H. (2014). Regression conformal prediction with random forests. Machine learning, 97, 155-176.
#' @examples
#' # Friedman function
#' f <-  function(x){
#'   10 * sin(pi * x[1] * x[2]) + 20 * (x[3] - 0.5)^2 + 10 * x[4] + 5 * x[5]
#' }
#' X <- matrix(runif(250), nrow=50, ncol=5)
#' y <- apply(X, 1, f)
#' fit <- rfok(X, y)
#'
#' Xnew <- matrix(runif(250), nrow=50, ncol=5)
#' predict(fit, Xnew)
#' @importFrom graphics par abline segments
#' @importFrom stats predict quantile rbinom runif sd
#' @export
rfok <- function(X, y, k = 5, beta = sd(y) / 30, eps = 0, ...) {
  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 (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)
  }
  if (anyNA(X)) {
    stop("`X` must not contain missing values.", call. = FALSE)
  }
  if (length(k) != 1 || is.na(k)) {
    stop("`k` must be a single non-missing number.", call. = FALSE)
  }
  if (!is.numeric(k)) {
    stop("`k` must be numeric.", call. = FALSE)
  }
  if (k < 1) {
    stop("`k` must be at least 1.", call. = FALSE)
  }
  if (k != floor(k)) {
    stop("`k` must be a whole number.", call. = FALSE)
  }
  k <- as.integer(k)
  if (k > nrow(X)) {
    warning("`k` is greater than `nrow(X)`; using `k = nrow(X)` instead.",
            call. = FALSE)
    k <- nrow(X)
  }
  if (length(beta) != 1 || is.na(beta)) {
    stop("`beta` must be a single non-missing number.", call. = FALSE)
  }
  if (!is.numeric(beta)) {
    stop("`beta` must be numeric.", call. = FALSE)
  }
  if (beta < 0) {
    stop("`beta` must be non-negative.", call. = FALSE)
  }
  if (length(eps) != 1 || is.na(eps)) {
    stop("`eps` must be a single non-missing number.", call. = FALSE)
  }
  if (!is.numeric(eps)) {
    stop("`eps` must be numeric.", call. = FALSE)
  }
  if (eps < 0) {
    stop("`eps` must be non-negative.", call. = FALSE)
  }
  # End input validation

  n <- nrow(X)

  # Exact nn search?
  if(eps == 0){
    searchtype = "standard"
  }else{
    searchtype = "priority"
  }

  # Fit RF model
  fit <- randomForest::randomForest(X, y, keep.inbag=TRUE, ...)

  # Get predictions
  preds <- predict(fit, newdata=X, predict.all=TRUE)
  yhat <- preds$aggregate
  preds <- preds$individual

  # Get out of bag predictions
  out_of_bag_sets <- apply(fit$inbag, 1, function(xx) which(xx==0))
  yhat_oob <- rep(NA, n)
  for(i in seq_len(n)){
    oob_i <- out_of_bag_sets[[i]]
    if(length(oob_i) == 0){
      yhat_oob[i] <- yhat[i]
    } else {
      yhat_oob[i] <- mean(preds[i, oob_i])
    }
  }

  # Get k nearest neighbors for each point
  neighbors <- RANN::nn2(X, X, k=min(k, nrow(X)), searchtype=searchtype, eps=eps)$nn.idx

  # Estimate the mus
  mu <- rep(NA, n)
  for(i in 1:n){
    ind <- neighbors[i,]
    mu[i] <- mean(abs(y[ind] - yhat_oob[ind]))
  }

  # Calculate non-conformity scores
  alpha <- abs(y - yhat_oob)/(mu + beta)
  #alpha_s <- quantile(alpha, conf)

  # Return object
  object <- list(
    fit = fit,
    k = k,
    beta = beta,
    eps = eps,
    alpha = alpha,
    oob_error = yhat_oob - y,
    X = X,
    y = y,
    call = match.call()
  )
  class(object) <- "rfok"
  return(object)
}


#' Predict method for rfok
#'
#' A conformal RF based on Johansson et al. (2014). Rather than using a traditional calibration set, conformal intervals are constructed using out-of-bag samples for each tree. K-nearest neighbors is used to calibrate the prediction intervals for each prediction location.
#'
#' @param object Object returned by \code{rfok()}
#' @param newdata a data frame or matrix of new data
#' @param samples Number of samples from the predictive distribution.
#' @param conf a vector of levels for desired confidence intervals. Ignored unless samples is NULL.
#' @param smoothing logical; should conformity scores be sampled or smoothed (using quantiles)
#' @param eps Tolerance for nearest neighbor approach. When \code{eps = 0}, exact nearest neighbors are used.
#' @param ... Additional, ignored, parameters
#' @return A matrix of predictions with 1 column per test point.
#' @details
#' Prediction intervals are constructed by estimating a local scale for each
#' prediction point using the \code{k} nearest training observations. The
#' conformal margin of error is the requested quantile of the training
#' non-conformity scores multiplied by this local scale.
#'
#' If \code{samples = NULL}, the function returns point predictions and margins
#' of error for the confidence levels supplied in \code{conf}. For example,
#' \code{conf = 0.95} uses the 0.95 quantile of the non-conformity scores.
#'
#' If \code{samples} is a positive integer, the function returns simulated
#' predictive samples. When \code{smoothing = TRUE}, signed non-conformity scores
#' are generated using quantiles of the symmetrized empirical distribution. When
#' \code{smoothing = FALSE}, non-conformity scores are sampled directly with
#' replacement and assigned random signs.
#' @examples
#' # Friedman function
#' f <-  function(x){
#'   10 * sin(pi * x[1] * x[2]) + 20 * (x[3] - 0.5)^2 + 10 * x[4] + 5 * x[5]
#' }
#' X <- matrix(runif(250), nrow=50, ncol=5)
#' y <- apply(X, 1, f)
#' fit <- rfok(X, y)
#'
#' Xnew <- matrix(runif(250), nrow=50, ncol=5)
#' predict(fit, Xnew)
#' @export
predict.rfok <- function(object, newdata=NULL, samples=1000, conf=NULL, smoothing=TRUE, eps=0, ...){
  if(is.null(newdata)){
    newdata <- object$X
  }

  # Exact nn search?
  if(eps == 0){
    searchtype = "standard"
  }else{
    searchtype = "priority"
  }


  pred <- predict(object$fit, newdata)
  n <- length(pred)

  if(is.null(conf) & is.null(samples)){
    return(pred)
  }

  # Get nearest neighbors
  neighbors <- RANN::nn2(object$X, newdata, k=min(object$k, nrow(object$X)),
                         searchtype=searchtype, eps=eps)$nn.idx

  # Estimate mus
  mu <- rep(NA, n)
  for(i in 1:n){
    ind <- neighbors[i,]
    mu[i] <- mean(abs(object$oob_error[ind]))
  }

  if(is.null(samples)){
    # Generate predictions and confidence intervals
    moe <- matrix(NA, nrow=length(pred), ncol=length(conf))
    colnames(moe) <- gsub("\\.", "_", paste(conf))
    for(i in seq_along(conf)){
      moe[,i] <- quantile(object$alpha, conf[i])*(mu + object$beta)
    }
    out <- list(pred=pred, moe=moe)
    return(out)
  }else{
    # Generate predictive samples
    preds <- matrix(NA, nrow=samples, ncol=n)
    for(i in 1:n){
      if(smoothing){
        alpha_delta <- quantile(c(object$alpha, -object$alpha), runif(samples))
      }else{
        alpha_delta <- sample(object$alpha, samples, replace=TRUE) * (2*rbinom(samples,1,0.5)- 1)
      }
      error_term <- alpha_delta*(mu[i] + object$beta)
      preds[,i] <- pred[i] + error_term
    }
  }
  return(preds)
}

#' Plot method for rfok
#'
#' A conformal RF based on Johansson et al. (2014). Rather than using a traditional calibration set, conformal intervals are constructed using out-of-bag samples for each tree. K-nearest neighbors is used to calibrate the prediction intervals for each prediction location.
#'
#' @param x Object returned by \code{rfok()}
#' @param conf Confidence level for plotting
#' @param ... Optional parameters to pass to \code{plot()}
#' @return Invisibly returns \code{x}.
#' @examples
#' # Friedman function
#' f <-  function(x){
#'   10 * sin(pi * x[1] * x[2]) + 20 * (x[3] - 0.5)^2 + 10 * x[4] + 5 * x[5]
#' }
#' X <- matrix(runif(250), nrow=50, ncol=5)
#' y <- apply(X, 1, f)
#' fit <- rfok(X, y)
#'
#' plot(fit)
#' @export
plot.rfok <- function(x, conf=0.95, ...){
  hold_par <- par(no.readonly = TRUE)
  on.exit(par(hold_par), add = TRUE)
  par(mfrow = c(1, 2))
  object <- x
  a <- 1-conf
  xx <- object$X
  pp <- predict(object)
  yhat <- colMeans(pp)
  bounds <- apply(pp, 2, quantile, probs = c(a/2, 1-a/2))

  par(mfrow=c(1,2))
  plot(object$y, yhat, ylim=range(pp), main="Predictions", ...)
  abline(0, 1, lwd=2, col='orange')
  segments(object$y, bounds[1,], object$y, bounds[2,])

  plot(object$fit, main="MSE")
  invisible(x)
}

#' Print method for rfok objects
#'
#' @param x Object returned by \code{rfok()}.
#' @param ... Additional arguments, currently ignored.
#'
#' @return Invisibly returns \code{x}.
#'
#' @export
#' @method print rfok
print.rfok <- function(x, ...) {
  cat("Conformal Random Forest with Out-of-bag KNN\n")
  cat("\n")

  cat("Training observations:", length(x$y), "\n")
  cat("Predictors:", ncol(x$X), "\n")
  cat("K nearest neighbors:", x$k, "\n")
  cat("Beta:", signif(x$beta, 4), "\n")

  if (!is.null(x$fit$ntree)) {
    cat("Trees:", x$fit$ntree, "\n")
  }

  cat("\n")
  cat("OOB error summary:\n")
  print(summary(x$oob_error))

  cat("\n")
  cat("Nonconformity score summary:\n")
  print(summary(x$alpha))

  invisible(x)
}


#' Summary method for rfok objects
#'
#' @param object Object returned by \code{rfok()}.
#' @param ... Additional arguments, currently ignored.
#'
#' @return An object of class \code{"summary.rfok"}.
#'
#' @export
#' @method summary rfok
summary.rfok <- function(object, ...) {
  out <- list(
    call = object$call,
    n = length(object$y),
    p = ncol(object$X),
    k = object$k,
    beta = object$beta,
    oob_error = summary(object$oob_error),
    alpha = summary(object$alpha)
  )
  class(out) <- "summary.rfok"
  out
}

#' @export
#' @method print summary.rfok
print.summary.rfok <- function(x, ...) {
  cat("Summary of rfok object\n\n")
  cat("Training observations:", x$n, "\n")
  cat("Predictors:", x$p, "\n")
  cat("K nearest neighbors:", x$k, "\n")
  cat("Beta:", signif(x$beta, 4), "\n")

  cat("\nOOB error summary:\n")
  print(x$oob_error)

  cat("\nNonconformity score summary:\n")
  print(x$alpha)

  invisible(x)
}
