
utils::globalVariables(c("v", "PC1", "PC2"))

#' @title Supervised Principal Component Analysis (SPCA)
#' @description True SPCA (Bair et al.): screening + PCA
#'
#' @param data Numeric data frame or matrix
#' @param response Numeric response vector
#' @param K Number of components (auto if NULL)
#' @param threshold Variance threshold for K selection
#' @param verbose Logical; if TRUE prints progress messages
#'
#' @return List of SPCA results
#'
#' @importFrom stats prcomp cor complete.cases sd
#' @importFrom psych cortest.bartlett KMO
#' @importFrom ggplot2 ggplot aes geom_col geom_line geom_point
#' @importFrom ggplot2 geom_text 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 scales percent_format
#' @importFrom grid unit arrow
#' @export
#'
#' @examples
#' data(mtcars)
#' Result <- supervised_pca(
#'   mtcars[, -1],
#'   mtcars$mpg
#' )

supervised_pca <- function(data,
                           response,
                           K = NULL,
                           threshold = 0.9,
                           verbose = FALSE) {

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

  if (!is.numeric(response)) {
    stop("response must be numeric")
  }

  df <- as.data.frame(data)

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

  if (ncol(df) < 2) {
    stop("Need at least 2 numeric predictors")
  }

  if (length(response) != nrow(df)) {
    stop("response length must match rows")
  }

  idx <- stats::complete.cases(df, response)

  df <- df[idx, , drop = FALSE]
  y <- response[idx]

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

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

      if (stats::sd(x) == 0) {
        return(0)
      }

      stats::cor(x, y)
    },
    numeric(1)
  )

  keep <- abs(cors) > 0.2

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

  df_sel <- df[, keep, drop = FALSE]

  if (verbose) {

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

    message(
      "Selected Variables: ",
      paste(names(df_sel),
            collapse = ", ")
    )
  }

  if (is.null(K)) {

    p_tmp <- stats::prcomp(
      df_sel,
      scale. = TRUE
    )

    ve <- p_tmp$sdev^2 /
      sum(p_tmp$sdev^2)

    K <- which(
      cumsum(ve) >= threshold
    )[1]

    if (verbose) {

      message(
        "\nSelected K: ",
        K
      )
    }
  }

  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(df_sel)) {
    stop("K cannot exceed variables")
  }

  cm <- stats::cor(df_sel)

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

  bartlett <- NULL
  kmo <- NULL

  if (!singular) {

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

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

  if (verbose) {

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

    if (!is.null(bartlett)) {

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

    } else {

      message(
        "Bartlett test not valid"
      )
    }

    if (!is.null(kmo)) {

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

    } else {

      message("KMO not valid")
    }
  }

  p <- stats::prcomp(
    df_sel,
    scale. = TRUE
  )

  ve_full <- p$sdev^2 /
    sum(p$sdev^2)

  pca_summary <- summary(p)

  load <- p$rotation

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

  theme_pretty <-
    ggplot2::theme_minimal(
      base_family = "serif",
      base_size = 16
    ) +
    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(ve_full)
      ),
      levels = paste0(
        "PC",
        seq_along(ve_full)
      )
    ),
    Var = ve_full,
    CumVar = cumsum(ve_full)
  )

  scree <- 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.2
    ) +
    ggplot2::geom_point(
      ggplot2::aes(y = CumVar),
      color = "#D55E00",
      size = 3
    ) +
    ggplot2::geom_text(
      ggplot2::aes(
        y = Var,
        label = paste0(
          round(Var * 100, 1),
          "%"
        )
      ),
      vjust = -0.5,
      family = "serif",
      fontface = "bold",
      size = 4
    ) +
    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 <- NULL

  if (ncol(scores) >= 2) {

    load2 <- as.data.frame(
      load[, 1:2, drop = FALSE]
    )

    load2$v <- rownames(load2)

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

    sf <- max(abs(score2))

    biplot <- ggplot2::ggplot() +

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

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

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

      ggplot2::labs(
        title = "Supervised PCA Biplot",
        x = "Component 1",
        y = "Component 2"
      ) +

      theme_pretty +

      ggplot2::coord_equal()

  } else {

    warning("Biplot not created")
  }

  invisible(list(
    K = K,
    selected_vars = names(df_sel),
    bartlett = bartlett,
    kmo = kmo,
    summary = pca_summary,
    loadings = load,
    scores = scores,
    variance = ve_full,
    scree_plot = scree,
    biplot = biplot
  ))
}
