
utils::globalVariables(c(
  "PC", "Var", "CumVar",
  "PC1", "PC2", "v"
))

#' @title Supervised Principal Component Regression
#'
#' @description
#' Performs supervised feature selection followed by
#' Principal Component Regression (PCR).
#' Includes Bartlett's test, KMO measure,
#' scree plot, biplot, and regression models.
#'
#' @param data A data frame or matrix
#' @param response Response variable name
#' @param K Number of principal components
#' @param threshold Cumulative variance threshold
#' for automatic component selection
#' @param verbose Logical; print summaries?
#'
#' @return A list containing PCA results,
#' regression models, diagnostics, and plots.
#'
#' @importFrom stats cor prcomp lm step complete.cases sd
#' @importFrom psych cortest.bartlett KMO
#' @importFrom ggplot2 ggplot aes geom_col geom_line
#' @importFrom ggplot2 geom_point geom_text
#' @importFrom ggplot2 scale_y_continuous labs
#' @importFrom ggplot2 theme_minimal theme
#' @importFrom ggplot2 element_text element_line
#' @importFrom ggplot2 element_blank coord_equal
#' @importFrom ggrepel geom_text_repel
#' @importFrom gridExtra grid.arrange
#' @importFrom scales percent_format
#' @importFrom grid unit arrow
#'
#' @export
#'
#' @examples
#' data(mtcars)
#' result <- supervised_pcreg(
#'   mtcars,
#'   response = "mpg"
#' )

supervised_pcreg <- function(data,
                             response,
                             K = NULL,
                             threshold = 0.90,
                             verbose = FALSE) {

  if (!is.data.frame(data) &&
      !is.matrix(data)) {
    stop("data must be a data frame or matrix")
  }

  df <- as.data.frame(data)

  df <- df[, vapply(
    df,
    is.numeric,
    logical(1)
  ), drop = FALSE]

  if (!is.character(response) ||
      length(response) != 1) {
    stop(
      "response must be a single character string"
    )
  }

  if (!(response %in% names(df))) {
    stop("Response variable not found")
  }

  Y <- df[[response]]

  X <- df[, setdiff(
    names(df),
    response
  ), drop = FALSE]

  if (ncol(X) < 2) {
    stop("Need at least two predictors")
  }

  idx <- stats::complete.cases(X, Y)

  X <- X[idx, , drop = FALSE]
  Y <- Y[idx]

  if (nrow(X) < 2) {
    stop("Not enough complete observations")
  }

  const <- vapply(
    X,
    function(x) stats::sd(x) == 0,
    logical(1)
  )

  if (any(const)) {

    if (verbose) {
      message(
        "Removed constant predictors: ",
        paste(names(X)[const],
              collapse = ", ")
      )
    }

    X <- X[, !const, drop = FALSE]
  }

  if (ncol(X) < 2) {
    stop(
      "Need at least two non-constant predictors"
    )
  }

  cors <- vapply(
    X,
    function(x) {

      r <- stats::cor(
        x,
        Y,
        use = "complete.obs"
      )

      ifelse(is.na(r), 0, r)
    },
    numeric(1)
  )

  keep <- abs(cors) > 0.20

  if (sum(keep) < 2) {
    stop("Too few variables selected")
  }

  X_sel <- X[, keep, drop = FALSE]

  if (verbose) {

    message("\n--- Supervised Selection ---")

    message(
      paste(
        names(X_sel),
        collapse = ", "
      )
    )
  }

  cm <- stats::cor(X_sel)

  singular <- tryCatch(
    base::det(cm) < 1e-10,
    error = function(e) TRUE
  )

  bart <- NULL
  kmo <- NULL

  if (!singular) {

    bart <- tryCatch(
      psych::cortest.bartlett(
        cm,
        n = nrow(X_sel)
      ),
      error = function(e) NULL
    )

    kmo <- tryCatch(
      psych::KMO(cm),
      error = function(e) NULL
    )
  }

  if (verbose) {

    message("\n--- PCA Suitability ---")

    if (!is.null(bart)) {

      message(
        "Bartlett p-value: ",
        format.pval(bart$p.value)
      )

    } else {

      message("Bartlett test not valid")
    }

    if (!is.null(kmo)) {

      message(
        "KMO: ",
        round(kmo$MSA, 3)
      )

    } else {

      message("KMO not valid")
    }
  }

  choose_K <- function(data,
                       threshold = 0.90) {

    pca0 <- stats::prcomp(
      data,
      scale. = TRUE
    )

    ve0 <- pca0$sdev^2 /
      sum(pca0$sdev^2)

    which(cumsum(ve0) >= threshold)[1]
  }

  if (is.null(K)) {

    K <- choose_K(
      X_sel,
      threshold = threshold
    )
  }

  if (!is.numeric(K) ||
      length(K) != 1) {
    stop("K must be a single numeric value")
  }

  K <- as.integer(K)

  if (K < 1) {
    stop("K must be >= 1")
  }

  if (K > ncol(X_sel)) {
    stop(
      "K cannot exceed number of selected predictors"
    )
  }

  if (verbose) {
    message("\nUsing K = ", K)
  }

  pca <- stats::prcomp(
    X_sel,
    scale. = TRUE
  )

  variance <- pca$sdev^2 /
    sum(pca$sdev^2)

  pca_summary <- summary(pca)

  loadings <- as.data.frame(
    pca$rotation[, 1:K, drop = FALSE]
  )

  scores <- as.data.frame(
    pca$x[, 1:K, drop = FALSE]
  )

  if (verbose) {

    message("\n--- PCA Summary ---")

    print(pca_summary$importance)

    message("\n--- PCA Loadings ---")

    print(round(loadings, 3))
  }

  reg_data <- data.frame(
    Y = Y,
    scores
  )

  full_model <- stats::lm(
    Y ~ .,
    data = reg_data
  )

  backward_model <- tryCatch(
    stats::step(
      full_model,
      direction = "backward",
      trace = 0
    ),
    error = function(e) full_model
  )

  if (verbose) {

    message("\n--- Full Model ---")

    print(summary(full_model))

    message(
      "\n--- Backward Selected Model ---"
    )

    print(summary(backward_model))
  }

  theme_pretty <-
    ggplot2::theme_minimal(
      base_family = "serif",
      base_size = 14
    ) +
    ggplot2::theme(
      plot.title =
        ggplot2::element_text(
          face = "bold",
          hjust = 0.5
        ),
      axis.title =
        ggplot2::element_text(
          face = "bold"
        ),
      axis.text =
        ggplot2::element_text(
          face = "bold",
          color = "black"
        ),
      axis.line =
        ggplot2::element_line(
          color = "black"
        ),
      axis.ticks =
        ggplot2::element_line(
          color = "black"
        ),
      panel.grid.minor =
        ggplot2::element_blank()
    )

  scree_data <- data.frame(
    PC = factor(
      paste0(
        "PC",
        seq_along(variance)
      ),
      levels = paste0(
        "PC",
        seq_along(variance)
      )
    ),
    Var = variance,
    CumVar = cumsum(variance)
  )

  scree_plot <- ggplot2::ggplot(
    scree_data,
    ggplot2::aes(x = PC)
  ) +
    ggplot2::geom_col(
      ggplot2::aes(y = Var),
      fill = "#4C72B0",
      width = 0.6
    ) +
    ggplot2::geom_line(
      ggplot2::aes(
        y = CumVar,
        group = 1
      ),
      color = "#D55E00",
      linewidth = 1
    ) +
    ggplot2::geom_point(
      ggplot2::aes(y = CumVar),
      color = "#D55E00",
      size = 2.5
    ) +
    ggplot2::geom_text(
      ggplot2::aes(
        y = Var,
        label = paste0(
          round(Var * 100, 1),
          "%"
        )
      ),
      vjust = -0.4,
      size = 4,
      family = "serif",
      fontface = "bold"
    ) +
    ggplot2::scale_y_continuous(
      labels = scales::percent_format(),
      limits = c(0, 1)
    ) +
    ggplot2::labs(
      title = "Supervised PCA Scree Plot",
      x = "Principal Components",
      y = "Variance Explained"
    ) +
    theme_pretty

  biplot_obj <- NULL

  if (K >= 2) {

    load2 <- loadings[, 1:2, drop = FALSE]

    load2$v <- rownames(load2)

    score2 <- scores[, 1:2, drop = FALSE]

    scale_factor <- max(abs(score2))

    biplot_obj <- ggplot2::ggplot() +

      ggplot2::geom_point(
        data = score2,
        ggplot2::aes(PC1, PC2),
        color = "#2C7BB6",
        size = 3,
        alpha = 0.7
      ) +

      ggplot2::geom_segment(
        data = load2,
        ggplot2::aes(
          x = 0,
          y = 0,
          xend = PC1 * scale_factor,
          yend = PC2 * scale_factor
        ),
        arrow = grid::arrow(
          length = grid::unit(
            0.25,
            "cm"
          )
        ),
        color = "#D7191C",
        linewidth = 1
      ) +

      ggrepel::geom_text_repel(
        data = load2,
        ggplot2::aes(
          PC1 * scale_factor,
          PC2 * scale_factor,
          label = v
        ),
        color = "#D7191C",
        family = "serif",
        fontface = "bold",
        size = 4
      ) +

      ggplot2::labs(
        title = "Supervised PCA Biplot",
        x = "PC1",
        y = "PC2"
      ) +

      ggplot2::coord_equal() +
      theme_pretty

    gridExtra::grid.arrange(
      scree_plot,
      biplot_obj,
      ncol = 2
    )

  } else {

    warning(
      "Biplot requires at least two components"
    )

    print(scree_plot)
  }

  invisible(list(
    K = K,
    selected_variables = names(X_sel),
    pca = pca,
    bartlett = bart,
    kmo = kmo,
    loadings = loadings,
    scores = scores,
    variance = variance,
    full_model = summary(full_model),
    backward_model =
      summary(backward_model),
    scree_plot = scree_plot,
    biplot = biplot_obj
  ))
}
