#' @useDynLib GEMSS, .registration = TRUE
#' @importFrom Rcpp sourceCpp
NULL

#' Compute the Correlation Matrix for Specified Kernels
#'
#' @description
#' Calculates the correlation matrix between two sets of locations using
#' stationary kernels. This function supports Gaussian, Matern 5/2, and
#' Matern 3/2 correlation structures in a multivariate product form.
#'
#' @param X1 a matrix of input locations, where each row represents an observation.
#' @param X2 a matrix of design locations. If \code{NULL}, the correlation
#'   is calculated between \code{X1} and itself.
#' @param theta a \code{p x 1} numeric vector of lengthscale parameters
#'   corresponding to each dimension.
#' @param type a character string specifying the kernel type. Must be one of
#'   "\code{Gaussian}", "\code{Matern5_2}", or "\code{Matern3_2}".
#'
#' @details
#' For multivariate inputs, the correlation function is defined as the product
#' of univariate kernels across all \eqn{p} dimensions:
#' \deqn{C(\mathbf{x}, \mathbf{y}, \mathbf{\theta}) = \prod_{k=1}^{p} c_k(|x_k - y_k|, \theta_k)}
#'
#' The available univariate correlation functions \eqn{c(d)} (where \eqn{d = |x - y|}) are:
#' \itemize{
#'   \item "\code{Gaussian}": \deqn{c(d, \theta) = \exp\left(-\frac{d^2}{\theta}\right)}
#'   \item "\code{Matern5_2}": \deqn{c(d, \theta) = \left(1 + \frac{\sqrt{5}d}{\theta} + \frac{5d^2}{3\theta^2}\right) \exp\left(-\frac{\sqrt{5}d}{\theta}\right)}
#'   \item "\code{Matern3_2}": \deqn{c(d, \theta) = \left(1 + \frac{\sqrt{3}d}{\theta}\right) \exp\left(-\frac{\sqrt{3}d}{\theta}\right)}
#' }
#'
#' @return a numeric matrix of dimensions \code{nrow(X1) x nrow(X2)}.
#' @export
compute_kernel <- function(X1, X2 = NULL, theta, type) {
  if (!(type %in% c("Gaussian", "Matern5_2", "Matern3_2"))) stop("type should be one of 'Gaussian', 'Matern5_2', 'Matern3_2'")
  X1 <- as.matrix(X1)
  if(is.null(X2)){
    X2 <- as.matrix(X1)
  }
  return(compute_kernel_cpp(X1, X2, theta, type))
}


#' Gaussian Process Prediction
#'
#' @description
#' Performs Gaussian Process prediction for a given set of new design locations,
#' using hyperparameters provided in the parameters list.
#'
#' @param X_new a \code{m x p} matrix of new design locations to predict.
#' @param X a \code{n x p} matrix of training design locations.
#' @param Y a \code{n x 1} numeric vector of observed responses at training locations.
#' @param covtype a character string specifying the covariance kernel type. Must be
#'   one of "\code{Gaussian}", "\code{Matern5_2}", or "\code{Matern3_2}".
#' @param parameters a list of hyperparameters containing \code{theta}, \code{g},
#'   \code{sigma2}, and \code{beta0}.
#'
#' @return a list containing:
#' \itemize{
#'   \item \code{mean}: a \code{m x 1} numeric vector of predicted means at \code{X_new}.
#'   \item \code{sd2}: a \code{m x 1} numeric vector of predicted variances at \code{X_new}.
#' }
#'
#' @details
#' The function uses the standard GP prediction formulas:
#' \deqn{\mu(x_{new}) = \beta_0 + k(x_{new}, X) K^{-1} (Y - \beta_0)}
#' \deqn{s^2(x_{new}) = \sigma^2 [ (1 + g) - k(x_{new}, X) (K + gI)^{-1} k(X, x_{new}) ]}
#'
#' @export
gp_predict <- function(X_new, X, Y, covtype, parameters) {
  if (!(covtype %in% c("Gaussian", "Matern5_2", "Matern3_2"))) stop("covtype should be one of 'Gaussian', 'Matern5_2', 'Matern3_2'")
  gp_predict_cpp(as.matrix(X_new), as.matrix(X), as.numeric(Y), covtype, parameters)
}

#' Subdata Selection via GEMSS
#'
#' @description
#' Implements the Generalized Error-Minimizing Subsampling (GEMSS) algorithm to select subdata for
#' Gaussian Process models.
#'
#' @param X an \code{n x p} matrix of input design locations.
#' @param Y an \code{n x 1} numeric vector of observed responses.
#' @param ns target size of the final subdata.
#' @param covtype covariance kernel type, either "Gaussian", "Matern5_2", or "Matern3_2".
#' @param parameters a list of hyperparameters: \code{beta0} (mean),
#'   \code{theta} (length-scales), and \code{g} (nugget). If not provided,
#'   these are estimated via \code{hetGP::\link[hetGP]{mleHomGP}} on a pilot
#'   sample generated by \code{twinning::\link[twinning]{twin}}.
#' @param X_val,Y_val optional validation data for calculating the predictive R-squared.
#'   If omitted, a random sample of size \code{min(0.1*n, 5000)} is used.
#' @param c1 value between 0 and 1 of the first term in GEMSS criterion (\eqn{c_2 = 1 - c_1}).
#'   The default (\code{NULL}) uses the adaptive weighting proposed in the paper.
#' @param n_srs,n_top optional values to determine the candidate set for each iteration:
#' \itemize{
#'   \item \code{n_srs}: number of random data points (default is \code{ns/2}).
#'   \item \code{n_top}: number of top-ranked points from the previous iteration (default is \code{ns/2}).
#' }
#' @param verbose logical; if \code{TRUE}, progress is printed at each iteration.
#'
#' @return a list containing:
#' \itemize{
#'   \item \code{index}: the indices of the selected subdata.
#'   \item \code{r_sq}: a data.frame containing the subdata size and the corresponding predictive R-square.
#'   \item \code{parameters}: a list of hyperparameters containing \code{beta0}, \code{theta}, \code{g}, and \code{sigma2} (the plug-in estimator of the variance).
#' }
#'
#' @details
#' The GEMSS criterion is defined as:
#' \deqn{\mathcal{G} = c_1 \delta_1^2 + c_2 \delta_2}
#' where \eqn{\delta_1^2} is the squared prediction error and \eqn{\delta_2} is the
#' posterior variance. The weights \eqn{c_1} and \eqn{c_2 = 1 - c_1} control the
#' balance between prediction accuracy and space-filling properties.
#'
#' The default value for \eqn{c_1} is \eqn{3\delta_2^2 / (2\delta_1^4 + 3\delta_2^2)},
#' which is derived by minimizing the variance of the GEMSS criterion. Setting
#' \eqn{c_1 = 0} results in a space-filling subdata
#'
#' The algorithm initializes with a small random sample of size 2 and
#' sequentially adds the point from the candidate set that maximizes \eqn{\mathcal{G}}
#' until the subdata reaches size ns.
#'
#' The candidate set in each iteration is composed of two parts: a random
#' sample of size n_srs, and the n_top points that yielded the highest
#' \eqn{\mathcal{G}} values in the previous iteration.
#'
#' All GP hyperparameters except \eqn{\sigma^2} are estimated from an initial
#' pilot sample of size \eqn{ns} and remain fixed throughout the selection process.
#'
#'@references
#' Chang, M. C., Hua, S. Z., & Wu, C. F. J. (2026). GEMSS-Driven Subsampling
#' for Information Extraction and Redundancy Elimination. \emph{Technometrics},
#' 1–20. \doi{10.1080/00401706.2026.2670596}
#'
#' @examples
#' # --- Example 1: 1D Regression ---
#' fx <- function(x) cos((x - 0.8) * 2 * pi)^7 * sin(x) + 5 * sin(x) * (sin(x^2))^10
#' X <- matrix(runif(200), 200, 1)
#' Y <- apply(X, 1, fx) + rnorm(nrow(X), 0, 0.05)
#'
#' # Select 20 points using GEMSS
#' res <- gemss_select(X, Y, ns = 20, covtype = "Matern5_2")
#'
#' # Predict on a grid
#' x.grid <- as.matrix(seq(0, 1, 0.01))
#' y.pred <- gp_predict(x.grid, X[res$index, ], Y[res$index], "Matern5_2", res$parameters)
#'
#' # Visualize 1D Results
#' plot(X, Y, col = "grey", main = "Selected Subdata and GP Prediction")
#' points(X[res$index, ], Y[res$index], col = "red", pch = 19)
#' lines(x.grid, y.pred$mean, col = "red", lwd = 2)
#'
#' # --- Example 2: 2D Surface (Michalewicz function) ---
#' \donttest{
#' michalewicz <- function(xx) {
#'   x1 <- xx[1] * pi; x2 <- xx[2] * pi; m <- 10
#'   - (sin(x1) * (sin(x1^2 / pi))^(2 * m) + sin(x2) * (sin(2 * x2^2 / pi))^(2 * m))
#' }
#'
#' X2D <- matrix(runif(2000), 1000, 2)
#' Y2D <- apply(X2D, 1, michalewicz) + rnorm(nrow(X2D), 0, 0.005)
#'
#' res2D <- gemss_select(X2D, Y2D, ns = 80, covtype = "Matern5_2")
#'
#' # Visualization using ContourFunctions if available
#' if (requireNamespace("ContourFunctions", quietly = TRUE)) {
#'   ngrid <- 50 # Reduced grid size for faster example execution
#'   gx <- seq(0, 1, len = ngrid)
#'   grid <- as.matrix(expand.grid(x1 = gx, x2 = gx))
#'
#'   # plot the contour of Michalewicz function
#'   y_grid <- apply(grid, 1, michalewicz)
#'   ContourFunctions::cf_grid(gx, gx, matrix(y_grid, ngrid, ngrid),
#'                           bar = TRUE, main = "Michalewicz Function")
#'
#'   prep <- gp_predict(grid, X2D[res2D$index, ], Y2D[res2D$index],
#'   "Matern5_2", res2D$parameters)
#'   ContourFunctions::cf_grid(gx, gx, matrix(prep$mean, ngrid, ngrid),
#'                             bar = TRUE, main = "Prediction Using Selected Subdata",
#'                             afterplotfunc = function() {
#'                               points(X2D[res2D$index, ], col = 'blue', pch = 1, cex = 1.5, lwd = 2)
#'                             })
#' }
#' }
#'
#' @export
gemss_select <- function(X, Y, ns, covtype, parameters = NULL, X_val = NULL, Y_val = NULL, c1 = NULL, n_srs = NULL, n_top = NULL, verbose = TRUE) {
  if (nrow(X) != length(Y)) stop("Dimensions of X and Y do not match.")

  if (ns < 0 ) stop("ns should be positive integer")

  if (!(covtype %in% c("Gaussian", "Matern5_2", "Matern3_2"))) stop("covtype should be one of 'Gaussian', 'Matern5_2', 'Matern3_2'")

  if (is.null(parameters)) parameters <- list()

  if (!is.null(parameters$theta) && length(parameters$theta) != ncol(X)) {
    warning("Length of theta does not match ncol(X). Estimating theta instead.")
    parameters$theta <- NULL
  }

  if (is.null(parameters$theta) || is.null(parameters$g) || is.null(parameters$sigma2)) {
    known <- parameters
    parameters <- list()
    # Pilot Sample Selection
    ind_for_est <- twinning::twin(cbind(X, Y), r = floor(nrow(X) / ns))

    # Estimate only what is missing (known = parameters handles the logic)
    est_GP <- hetGP::mleHomGP(
      X = X[ind_for_est, , drop = FALSE],
      Z = Y[ind_for_est],
      covtype = covtype,
      known = known,
      maxit = 10000,
      noiseControl = list(g_bounds = c(sqrt(.Machine$double.eps), 10000))
    )

    # 3. Fill the gaps in parameters list
    parameters$theta <- est_GP$theta
    parameters$g     <- est_GP$g
    parameters$beta0 <- est_GP$beta0
  }

  if(is.null(X_val) || is.null(Y_val)){
    val_index <- sample(1:nrow(X), min(floor(0.1*nrow(X)), 5000))
    X_val <- X[val_index, , drop = FALSE]
    Y_val <- Y[val_index]

    not_val_ind <- setdiff(1:nrow(X), val_index)
    X <- X[not_val_ind, , drop = FALSE]
    Y <- Y[not_val_ind]
  }

  if(is.null(n_srs)) n_srs = ns/2
  if(is.null(n_top)) n_top = ns/2
  if(is.null(c1)){
    c1 = -1
  } else if (c1 > 1 || c1 < 0){
    stop("c1 is not between 0 and 1")
  }


  initial_idx <- sample(1:nrow(X), 2)


  results <- GEMSS_cpp_update_sig(
    X = as.matrix(X),
    Y = as.numeric(Y),
    X_val = as.matrix(X_val),
    Y_val = as.numeric(Y_val),
    initial_idx = initial_idx,
    ns = as.integer(ns),
    c1 = c1,
    n_srs = as.integer(n_srs),
    n_top = as.integer(n_top),
    theta = parameters$theta,
    nugget = parameters$g,
    beta0 = parameters$beta0,
    Cov_fun = covtype,
    print_result = verbose
  )
  parameters$sigma2 = results$sigma2

  list(
    index = not_val_ind[results$index],
    r_sq = data.frame(size = 3:ns, r_sq = round(results$r_sq[-c(1:2)], 4)),
    parameters = parameters
  )
}


#' Redundant Data Removal via GEMSS
#'
#' @description
#' Implements a backward elimination process to sequentially remove redundant
#' data points from a Gaussian Process model using the GEMSS criterion.
#'
#' @param X an \code{n x p} matrix of input design locations.
#' @param Y an \code{n x 1} numeric vector of observed responses.
#' @param X_val,Y_val validation data used to evaluate predictive performance.
#'   \code{X_val} must have the same number of columns as \code{X}.
#' @param n_remove an integer specifying the maximum number of data points to sequentially remove.
#' @param covtype a character string specifying the covariance kernel type. Must be
#'   one of "Gaussian", "Matern5_2", or "Matern3_2".
#' @param c1 a numeric value between 0 and 1 representing the weight of the first
#'   term in the GEMSS criterion (\eqn{c_2 = 1 - c_1}). The default (\code{NULL})
#'   uses the adaptive weighting proposed in the paper. A smaller \code{c1} places
#'   greater emphasis on the space-filling properties of the reduced subset.
#' @param threshold a numeric value specifying the minimum required value of predictive
#'   R-squared to justify removing points. Default is \code{0.01}.
#' @param verbose a logical value; if \code{TRUE}, progress is printed at each iteration.
#'
#' @return a list containing:
#' \itemize{
#'   \item \code{index}: an integer vector of the indices of the retained training data points.
#'   \item \code{remove}: an integer vector of the indices of the removed data points.
#'   \item \code{eval_matrix}: a matrix detailing the removal sequence, including the step \code{j},
#'     the \code{removed_id}, and the resulting out-of-sample \code{R2_pred}.
#' }
#'
#' @details
#' In each iteration, the leave-one-out GEMSS criterion is computed for every point
#' in the training set. The point that yields the smallest criterion value
#' is removed. The predictive R-squared
#' is then calculated on the validation set to evaluate the performance of the reduced dataset.
#'
#' If there exists a reduced dataset with a predictive R-squared value
#' greater than the specified \code{threshold}, the algorithm reports
#' the subset that maximizes the R-squared. Otherwise, no removal
#' is suggested.
#'
#' @references
#' Chang, M. C., Hua, S. Z., & Wu, C. F. J. (2026). GEMSS-Driven Subsampling
#' for Information Extraction and Redundancy Elimination. \emph{Technometrics},
#' 1–20. \doi{10.1080/00401706.2026.2670596}
#'@examples
#' # Generate 1D data with intentionally clustered (redundant) points
#' set.seed(123)
#' X_reg <- seq(0, 1, length.out = 20)
#' X_cluster <- rnorm(10, mean = 0.5, sd = 0.01) # Redundant points tightly packed around x=0.5
#' X <- as.matrix(c(X_reg, X_cluster))
#' Y <- sin(2 * pi * X) + rnorm(30, sd = 0.05)
#'
#' # Generate validation data
#' X_val <- as.matrix(seq(0.05, 0.95, length.out = 20))
#' Y_val <- sin(2 * pi * X_val) + rnorm(20, sd = 0.05)
#'
#' # Run GEMSS removal
#' res <- gemss_remove(X, Y, X_val, Y_val, n_remove = 10,
#' covtype = "Matern3_2", verbose = TRUE)
#'
#' # View the indices of the removed points
#' print(res$remove)
#' print(res$eval_matrix)
#'
#' # Plot the removal data points (red)
#' plot(X, Y, main = "GEMSS Data Removal")
#' points(X[res$remove], Y[res$remove], pch = 19, col = 2)
#'
#' # Compare GP predictions before and after removal
#' x_seq <- as.matrix(seq(0, 1, 0.02))
#' gp_bf <- hetGP::mleHomGP(X, Y, covtype = "Matern3_2")
#' lines(x_seq, predict(gp_bf, x_seq)$mean, col = 1, lty = 2)
#' gp_af <- hetGP::mleHomGP(X[-res$remove, , drop = FALSE],
#' Y[-res$remove], covtype = "Matern3_2")
#' lines(x_seq, predict(gp_af, x_seq)$mean, col = 3)
#'
#' @importFrom stats predict
#' @export
gemss_remove <- function(X, Y, X_val, Y_val, n_remove, covtype, c1 = NULL, threshold = 0.01, verbose = TRUE) {
  if (nrow(X) != length(Y)) stop("Dimensions of X and Y do not match.")

  if (n_remove < 0 ) stop("n_remove should be positive integer")

  if (!(covtype %in% c("Gaussian", "Matern5_2", "Matern3_2"))) stop("covtype should be one of 'Gaussian', 'Matern5_2', 'Matern3_2'")

  if(is.null(c1)){
    c1 = -1
  } else if (c1 > 1 || c1 < 0){
    stop("c1 is not between 0 and 1")
  }

  est_GP <- hetGP::mleHomGP(rbind(X, X_val), c(Y, Y_val), covtype = covtype, maxit = 10000, noiseControl = list(g_bounds = c(sqrt(.Machine$double.eps), 10000)))
  parameters <- est_GP[c(1,2,3,6)]

  cpp_res <- gemss_removal_cpp(X, Y,
                               theta = parameters$theta,
                               nugget = parameters$g,
                               beta0 = parameters$beta0,
                               sig2 = parameters$nu_hat,
                               Cov_fun = covtype,
                               n_remove = n_remove,
                               c1 = c1)

  removal_sequence <- cpp_res$removed_order
  index_remain <- cpp_res$index

  mse0 <- mean((Y_val - predict(hetGP::mleHomGP(X,Y, covtype = covtype), X_val)$mean)^2)
  eval_matrix <- matrix(NA, nrow = n_remove, ncol = 3)
  colnames(eval_matrix) <- c("j", "removed_id", "R2_pred")

  n_tr = nrow(X)
  index_tr <- c(index_remain, rev(removal_sequence))
  for (j in 1:n_remove) {
    point_to_remove <- removal_sequence[j]
    current_indices <- index_tr[1:(n_tr - j)]

    # Refit and predict
    fit <- hetGP::mleHomGP(X[current_indices, , drop = FALSE], Y[current_indices], covtype = covtype)
    pred <- predict(fit, X_val)$mean
    mspe <- mean((Y_val - pred)^2)

    eval_matrix[j, ] <- c(j, point_to_remove, 1 - mspe / mse0)
    if(verbose){
      cat(sprintf("j = %d | Removed: %d | R2: %.3f\n", j, point_to_remove, eval_matrix[j, 3]))
    }
  }

  # 4. Find the best stopping point
  best_j <- which.max(eval_matrix[,3])

  if(eval_matrix[best_j, 3] > threshold){
    message(paste0('Remove ', best_j, ' points'))
    remove <- eval_matrix[1:best_j, 2]
    index <- setdiff(1:nrow(X), remove)
    return(list(index = index, remove = remove, eval_matrix = eval_matrix))
  }else{
    message('No removal is suggested.')
    return(list(index = 1:nrow(X), remove = NULL,  eval_matrix = eval_matrix))
  }
}
