# Module 2: Normalisation and activity calculation ---------------------------

#' Normalise wells to plate controls
#'
#' Converts raw or background-subtracted signals to percent inhibition and
#' percent viability using each plate's own negative (NC) and positive (PC)
#' controls, so that readouts are comparable across plates. Inhibition is
#' \eqn{100 \times (m_{NC} - x) / (m_{NC} - m_{PC})} and viability is its
#' complement relative to the control window.
#'
#' @param hts_raw An `hts_raw` object (typically after [baseline_subtract()]).
#' @param signal_col Name of the signal column to normalise. Default
#'   `"signal"`.
#' @param nc_label,pc_label Well-type labels for the negative and positive
#'   controls. Defaults `"NC"` and `"PC"`.
#'
#' @return An `hts_norm` object: the input data with added numeric columns
#'   `inhibition` (percent) and `viability` (percent). Multiple plates are
#'   normalised independently.
#'
#' @examples
#' raw <- baseline_subtract(read_hts_plate(hts_primary_raw))
#' norm <- norm_by_control(raw)
#' summary(norm$inhibition)
#'
#' @seealso [merge_plate_data()], [plot_plate_heatmap_inhibition()]
#' @export
norm_by_control <- function(hts_raw, signal_col = "signal",
                            nc_label = "NC", pc_label = "PC") {
  if (!signal_col %in% names(hts_raw)) {
    stop("Column '", signal_col, "' not found in `hts_raw`.", call. = FALSE)
  }
  df <- as.data.frame(hts_raw, stringsAsFactors = FALSE)
  df$inhibition <- NA_real_
  df$viability <- NA_real_
  for (p in unique(as.character(df$plate_id))) {
    idx <- which(as.character(df$plate_id) == p)
    sig <- .as_numeric(df[[signal_col]][idx])
    m_nc <- mean(sig[df$well_type[idx] == nc_label], na.rm = TRUE)
    m_pc <- mean(sig[df$well_type[idx] == pc_label], na.rm = TRUE)
    window <- m_nc - m_pc
    if (!is.finite(window) || window == 0) {
      warning("Plate '", p, "' has a degenerate control window; ",
        "normalisation set to NA.", call. = FALSE)
      next
    }
    df$inhibition[idx] <- 100 * (m_nc - sig) / window
    df$viability[idx] <- 100 * (sig - m_pc) / window
  }
  .as_hts(df, c("hts_norm", "hts_raw"))
}

#' Combine normalised plates into one data set
#'
#' Row-binds several normalised plates (or a list of them) into a single
#' `hts_norm` object holding the activity of the whole screen, keeping only the
#' columns shared by all inputs.
#'
#' @param ... One or more `hts_norm`/data-frame objects, or a single list of
#'   them.
#'
#' @return A combined `hts_norm` object.
#'
#' @examples
#' n1 <- norm_by_control(read_hts_plate(
#'   hts_primary_raw[hts_primary_raw$plate_id == "P01", ]))
#' n2 <- norm_by_control(read_hts_plate(
#'   hts_primary_raw[hts_primary_raw$plate_id == "P02", ]))
#' merged <- merge_plate_data(n1, n2)
#' table(merged$plate_id)
#'
#' @seealso [norm_by_control()]
#' @export
merge_plate_data <- function(...) {
  args <- list(...)
  if (length(args) == 1 && is.list(args[[1]]) &&
    !is.data.frame(args[[1]])) {
    args <- args[[1]]
  }
  if (length(args) == 0) {
    stop("No plates supplied to `merge_plate_data()`.", call. = FALSE)
  }
  args <- lapply(args, as.data.frame, stringsAsFactors = FALSE)
  common <- Reduce(intersect, lapply(args, names))
  if (length(common) == 0) {
    stop("The supplied plates share no common columns.", call. = FALSE)
  }
  out <- do.call(rbind, lapply(args, function(d) d[, common, drop = FALSE]))
  rownames(out) <- NULL
  .as_hts(out, c("hts_norm", "hts_raw"))
}

#' Percent-inhibition plate heatmap
#'
#' Draws a plate's normalised percent-inhibition values in plate layout. A thin
#' wrapper around [plot_plate_heatmap()] for `hts_norm` objects.
#'
#' @param hts_norm An `hts_norm` object from [norm_by_control()].
#' @param plate Plate identifier to show. Defaults to the first plate.
#' @param value Name of the value column to colour by. Default `"inhibition"`.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' plot_plate_heatmap_inhibition(norm, plate = "P01")
#'
#' @seealso [plot_plate_heatmap()], [plot_inhibition_hist()]
#' @export
plot_plate_heatmap_inhibition <- function(hts_norm, plate = NULL,
                                          value = "inhibition") {
  if (is.null(plate)) plate <- as.character(hts_norm$plate_id)[1]
  df <- as.data.frame(hts_norm)
  plot_plate_heatmap(df, fill = value, well_col = "well",
    plate = plate, plate_col = "plate_id",
    control_col = if ("well_type" %in% names(df)) "well_type" else NULL,
    title = paste0("% ", value, ": ", plate))
}

#' Histogram of compound inhibition on a plate
#'
#' Shows the distribution of percent inhibition across the compound wells of a
#' plate (controls excluded), the standard view for judging overall hit density.
#'
#' @param hts_norm An `hts_norm` object.
#' @param plate Optional plate identifier; if `NULL` (default) all compound
#'   wells are pooled.
#' @param value Name of the value column. Default `"inhibition"`.
#' @param compound_label Well-type label marking compound wells. Default
#'   `"compound"`.
#' @param bins Number of histogram bins. Default `30`.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' plot_inhibition_hist(norm, plate = "P01")
#'
#' @seealso [norm_by_control()]
#' @export
plot_inhibition_hist <- function(hts_norm, plate = NULL, value = "inhibition",
                                 compound_label = "compound", bins = 30) {
  df <- as.data.frame(hts_norm)
  if ("well_type" %in% names(df)) {
    df <- df[df$well_type == compound_label, , drop = FALSE]
  }
  if (!is.null(plate)) {
    df <- df[as.character(df$plate_id) == plate, , drop = FALSE]
  }
  if (!value %in% names(df) || nrow(df) == 0) {
    stop("No compound '", value, "' values to plot.", call. = FALSE)
  }
  ggplot2::ggplot(df, ggplot2::aes(x = .data[[value]])) +
    ggplot2::geom_histogram(bins = bins, fill = .hts_col("compound"),
      color = "white", linewidth = 0.2, alpha = 0.9) +
    ggplot2::labs(x = paste0("% ", value), y = "Compounds",
      title = if (is.null(plate)) NULL else paste0("Plate ", plate)) +
    theme_hts()
}
