# Module 6: AUC matrix and global activity heatmap ---------------------------

#' Build a compound-by-sample AUC matrix
#'
#' Reshapes long AUC results into a numeric matrix with compounds as rows and
#' samples (for example cell lines or batches) as columns, optionally
#' standardised to Z-scores for comparability.
#'
#' @param data A data frame with compound, sample and AUC columns (typically
#'   `drc$meta` after [calc_drc_auc()]).
#' @param row_col Compound column. Default `"compound_id"`.
#' @param col_col Sample column. Default `"cell_line"`.
#' @param value_col AUC column. Default `"auc"`.
#' @param zscore Standardisation: `"none"` (default), `"row"` or `"column"`.
#'
#' @return A numeric matrix (compounds x samples). The standardisation used is
#'   stored in `attr(x, "zscore")`.
#'
#' @examples
#' \donttest{
#' drc <- calc_drc_auc(fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line"))
#' build_auc_matrix(drc$meta)
#' }
#'
#' @seealso [cluster_auc_matrix()], [plot_auc_heatmap()]
#' @export
build_auc_matrix <- function(data, row_col = "compound_id",
                             col_col = "cell_line", value_col = "auc",
                             zscore = c("none", "row", "column")) {
  zscore <- match.arg(zscore)
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  for (col in c(row_col, col_col, value_col)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  wide <- tidyr::pivot_wider(
    df[, c(row_col, col_col, value_col)],
    names_from = tidyr::all_of(col_col),
    values_from = tidyr::all_of(value_col),
    values_fn = function(z) mean(z, na.rm = TRUE)
  )
  wide <- as.data.frame(wide)
  rn <- as.character(wide[[row_col]])
  mat <- as.matrix(wide[, setdiff(names(wide), row_col), drop = FALSE])
  rownames(mat) <- rn
  z <- function(v) {
    s <- stats::sd(v, na.rm = TRUE)
    if (!is.finite(s) || s == 0) return(v - mean(v, na.rm = TRUE))
    (v - mean(v, na.rm = TRUE)) / s
  }
  if (zscore == "row") {
    mat <- t(apply(mat, 1, z))
  } else if (zscore == "column") {
    mat <- apply(mat, 2, z)
  }
  attr(mat, "zscore") <- zscore
  mat
}

#' Hierarchically cluster an AUC matrix
#'
#' Performs hierarchical clustering of the compounds (rows) and samples
#' (columns) of an AUC matrix and cuts the compound tree into groups.
#'
#' @param mat A numeric matrix from [build_auc_matrix()].
#' @param k Number of compound clusters to cut. Default `3`.
#' @param distance Distance method for [stats::dist()]. Default `"euclidean"`.
#' @param method Linkage method for [stats::hclust()]. Default `"complete"`.
#'
#' @return A list with `row_hclust`, `col_hclust` (or `NULL` if too few
#'   columns), `clusters` (named integer vector of compound cluster ids) and
#'   `k`.
#'
#' @examples
#' drc <- calc_drc_auc(fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line"))
#' cl <- cluster_auc_matrix(build_auc_matrix(drc$meta), k = 2)
#' table(cl$clusters)
#'
#' @seealso [extract_cluster_hit()], [plot_cluster_tree()]
#' @export
cluster_auc_matrix <- function(mat, k = 3, distance = "euclidean",
                               method = "complete") {
  mat <- mat[stats::complete.cases(mat), , drop = FALSE]
  if (nrow(mat) < 2) {
    stop("At least two complete compound rows are required to cluster.",
      call. = FALSE)
  }
  row_h <- stats::hclust(stats::dist(mat, method = distance), method = method)
  col_h <- if (ncol(mat) >= 2) {
    stats::hclust(stats::dist(t(mat), method = distance), method = method)
  } else {
    NULL
  }
  k <- min(k, nrow(mat))
  clusters <- stats::cutree(row_h, k = k)
  list(row_hclust = row_h, col_hclust = col_h, clusters = clusters, k = k)
}

#' Extract the most active compound cluster
#'
#' Returns the compounds belonging to the cluster with the most extreme mean
#' activity, that is the lowest mean AUC (strongest potency) by default.
#'
#' @param mat The AUC matrix used for clustering.
#' @param cluster_result The list returned by [cluster_auc_matrix()].
#' @param direction `"low"` (default) selects the lowest-AUC cluster (most
#'   active); `"high"` selects the highest-AUC cluster.
#'
#' @return A character vector of compound ids in the selected cluster.
#'
#' @examples
#' drc <- calc_drc_auc(fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line"))
#' m <- build_auc_matrix(drc$meta)
#' extract_cluster_hit(m, cluster_auc_matrix(m, k = 2))
#'
#' @seealso [cluster_auc_matrix()]
#' @export
extract_cluster_hit <- function(mat, cluster_result, direction = c("low",
                                "high")) {
  direction <- match.arg(direction)
  clusters <- cluster_result$clusters
  mat <- mat[names(clusters), , drop = FALSE]
  cluster_mean <- tapply(rowMeans(mat, na.rm = TRUE), clusters,
    function(z) mean(z, na.rm = TRUE))
  target <- if (direction == "low") {
    names(cluster_mean)[which.min(cluster_mean)]
  } else {
    names(cluster_mean)[which.max(cluster_mean)]
  }
  names(clusters)[clusters == as.integer(target)]
}

#' AUC clustering heatmap
#'
#' Draws a clustered heatmap of the AUC matrix. Uses \pkg{ComplexHeatmap} (with
#' \pkg{circlize} for the colour scale) when available for a publication-quality
#' figure with a diverging navy-white-firebrick palette matching
#' [plot_delta_auc_heatmap()], otherwise falls back to base [stats::heatmap()].
#'
#' @param mat A numeric matrix from [build_auc_matrix()].
#' @param title Plot title. Default `"AUC heatmap"`.
#' @param cluster_rows,cluster_columns Whether to cluster rows/columns. Default
#'   `TRUE`.
#'
#' @return Invisibly, the drawn object. Called for its side effect (a plot on
#'   the active graphics device).
#'
#' @examples
#' \donttest{
#' drc <- calc_drc_auc(fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line"))
#' plot_auc_heatmap(build_auc_matrix(drc$meta, zscore = "row"))
#' }
#'
#' @seealso [build_auc_matrix()], [plot_cluster_tree()]
#' @export
plot_auc_heatmap <- function(mat, title = "AUC heatmap",
                             cluster_rows = TRUE, cluster_columns = TRUE) {
  mat <- mat[stats::complete.cases(mat), , drop = FALSE]
  if (requireNamespace("ComplexHeatmap", quietly = TRUE) &&
    requireNamespace("circlize", quietly = TRUE)) {
    rng <- range(mat, na.rm = TRUE)
    mid <- if (rng[1] < 0 && rng[2] > 0) 0 else stats::median(mat, na.rm = TRUE)
    col_fun <- circlize::colorRamp2(
      c(rng[1], mid, rng[2]), c("navy", "white", "firebrick3")
    )
    ht <- ComplexHeatmap::Heatmap(mat, name = "AUC",
      col = col_fun,
      column_title = title,
      column_names_rot = 45,
      border = TRUE,
      cluster_rows = cluster_rows, cluster_columns = cluster_columns)
    ComplexHeatmap::draw(ht)
    return(invisible(ht))
  }
  stats::heatmap(mat, main = title,
    Rowv = if (cluster_rows) NULL else NA,
    Colv = if (cluster_columns) NULL else NA,
    scale = "none")
  invisible(mat)
}

#' Plot the compound clustering dendrogram
#'
#' Draws the hierarchical clustering tree of compounds from an AUC matrix.
#'
#' @param cluster_result The list returned by [cluster_auc_matrix()].
#' @param main Plot title. Default `"Compound clustering"`.
#'
#' @return Invisibly `NULL`; called for its side effect (a base-graphics plot).
#'
#' @examples
#' drc <- calc_drc_auc(fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line"))
#' m <- build_auc_matrix(drc$meta)
#' plot_cluster_tree(cluster_auc_matrix(m, k = 2))
#'
#' @seealso [cluster_auc_matrix()]
#' @export
plot_cluster_tree <- function(cluster_result, main = "Compound clustering") {
  if (is.null(cluster_result$row_hclust)) {
    stop("`cluster_result` has no row dendrogram.", call. = FALSE)
  }
  graphics::plot(cluster_result$row_hclust, main = main, xlab = "", sub = "")
  invisible(NULL)
}
