# Module 4: Primary screen hit selection -------------------------------------

# Per-compound activity summary from well-level data.
.compound_activity <- function(data, value, compound_col,
                               compound_label = "compound") {
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  if ("well_type" %in% names(df)) {
    df <- df[df$well_type == compound_label, , drop = FALSE]
  }
  for (col in c(value, compound_col)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  agg <- stats::aggregate(
    .as_numeric(df[[value]]),
    by = list(compound_id = df[[compound_col]]),
    FUN = function(z) mean(z, na.rm = TRUE)
  )
  names(agg)[2] <- "activity"
  agg
}

#' Select primary-screen hits
#'
#' Flags active compounds from single-concentration primary-screen data using
#' either a fixed activity threshold or a percentile rank of the whole library.
#'
#' @param data An `hts_norm` object or long data frame.
#' @param value Name of the activity column. Default `"inhibition"`.
#' @param method Selection mode: `"threshold"` (default) keeps compounds with
#'   activity at or above `threshold`; `"percentile"` keeps the top fraction.
#' @param threshold Activity cut-off for `method = "threshold"`. Default `50`.
#' @param percentile Upper fraction to keep for `method = "percentile"` (for
#'   example `0.95` keeps the top 5 percent). Default `0.95`.
#' @param compound_col Name of the compound column. Default `"compound_id"`.
#'
#' @return A data frame with one row per compound: `compound_id`, `activity`,
#'   `is_hit`. The applied cut-off is stored in `attr(x, "cutoff")` and the
#'   method in `attr(x, "method")`.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' hits <- select_primary_hit(norm, threshold = 50)
#' sum(hits$is_hit)
#'
#' @seealso [summarize_primary_hit()], [plot_primary_inhibition_rank()]
#' @export
select_primary_hit <- function(data, value = "inhibition",
                               method = c("threshold", "percentile"),
                               threshold = 50, percentile = 0.95,
                               compound_col = "compound_id") {
  method <- match.arg(method)
  act <- .compound_activity(data, value, compound_col)
  if (method == "threshold") {
    cutoff <- threshold
  } else {
    cutoff <- stats::quantile(act$activity, probs = percentile, na.rm = TRUE)
  }
  act$is_hit <- is.finite(act$activity) & act$activity >= cutoff
  act <- act[order(-act$activity), , drop = FALSE]
  rownames(act) <- NULL
  attr(act, "cutoff") <- unname(cutoff)
  attr(act, "method") <- method
  act
}

#' Summarise primary-screen hits
#'
#' Reports the number of compounds screened, the number of hits and the hit
#' rate.
#'
#' @param hits A data frame from [select_primary_hit()].
#'
#' @return A one-row data frame with `n_compounds`, `n_hits`, `hit_rate`
#'   (fraction) and `cutoff`.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' summarize_primary_hit(select_primary_hit(norm))
#'
#' @seealso [select_primary_hit()]
#' @export
summarize_primary_hit <- function(hits) {
  if (!"is_hit" %in% names(hits)) {
    stop("`hits` must contain an 'is_hit' column.", call. = FALSE)
  }
  n <- nrow(hits)
  n_hit <- sum(hits$is_hit, na.rm = TRUE)
  cutoff <- attr(hits, "cutoff")
  if (is.null(cutoff)) cutoff <- NA_real_
  data.frame(
    n_compounds = n,
    n_hits = n_hit,
    hit_rate = if (n > 0) n_hit / n else NA_real_,
    cutoff = cutoff,
    stringsAsFactors = FALSE
  )
}

#' Ranked primary-activity curve
#'
#' Draws the whole library ranked by activity, with the hit cut-off marked, the
#' standard view of how many compounds clear the primary threshold.
#'
#' @param hits A data frame from [select_primary_hit()].
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' plot_primary_inhibition_rank(select_primary_hit(norm))
#'
#' @seealso [select_primary_hit()]
#' @export
plot_primary_inhibition_rank <- function(hits) {
  if (!"activity" %in% names(hits)) {
    stop("`hits` must contain an 'activity' column.", call. = FALSE)
  }
  df <- hits[order(-hits$activity), , drop = FALSE]
  df$rank <- seq_len(nrow(df))
  cutoff <- attr(hits, "cutoff")
  p <- ggplot2::ggplot(df, ggplot2::aes(x = .data$rank, y = .data$activity)) +
    ggplot2::geom_point(ggplot2::aes(color = .data$is_hit), size = 1.6,
      alpha = 0.8, stroke = 0.2) +
    ggplot2::scale_color_manual(values = c("FALSE" = .hts_col("nonhit"),
      "TRUE" = .hts_col("hit")), name = "Hit") +
    ggplot2::labs(x = "Compound rank", y = "Activity") +
    theme_hts()
  if (!is.null(cutoff) && is.finite(cutoff)) {
    p <- p + ggplot2::geom_hline(yintercept = cutoff, linetype = "dashed",
      color = .hts_col("reference"), linewidth = 0.6)
  }
  p
}

#' Hit-count bar chart
#'
#' Bar chart of the number of hit versus non-hit compounds from a primary
#' screen.
#'
#' @param hits A data frame from [select_primary_hit()].
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' plot_hit_bar_count(select_primary_hit(norm))
#'
#' @seealso [summarize_primary_hit()]
#' @export
plot_hit_bar_count <- function(hits) {
  if (!"is_hit" %in% names(hits)) {
    stop("`hits` must contain an 'is_hit' column.", call. = FALSE)
  }
  df <- as.data.frame(table(is_hit = ifelse(hits$is_hit, "Hit", "Non-hit")))
  names(df) <- c("category", "count")
  ggplot2::ggplot(df, ggplot2::aes(x = .data$category, y = .data$count,
    fill = .data$category)) +
    ggplot2::geom_col(width = 0.6, color = "white", linewidth = 0.3) +
    ggplot2::geom_text(ggplot2::aes(label = .data$count), vjust = -0.4,
      color = "#333333", size = 3.6) +
    ggplot2::scale_fill_manual(values = c("Hit" = .hts_col("hit"),
      "Non-hit" = .hts_col("nonhit")), guide = "none") +
    ggplot2::scale_y_continuous(expand = ggplot2::expansion(mult = c(0, 0.1))) +
    ggplot2::labs(x = NULL, y = "Compounds") +
    theme_hts()
}

#' Select primary-screen hits by distance from the library distribution
#'
#' Treats the per-compound activity of the whole library as an approximately
#' normal "inactive" distribution and flags compounds sitting far out in the
#' tail. Each compound receives a z-score (how many standard deviations it lies
#' from the library centre) and is binned by whether it clears 1, 2 or 3 sigma.
#' By default the centre and spread are estimated *robustly* (median and MAD),
#' the HTS-standard choice: a handful of genuine hits would otherwise inflate a
#' plain mean/SD and mask themselves (with too wide an SD nothing reaches 3
#' sigma).
#'
#' @param data An `hts_norm` object or long data frame.
#' @param value Name of the activity column. Default `"inhibition"`.
#' @param method How to estimate the distribution centre and spread:
#'   `"robust"` (default) uses the median and MAD (scaled by 1.4826 to match the
#'   SD of a normal); `"sd"` uses the plain mean and standard deviation.
#' @param n_sigma Number of sigma a compound must clear to be called a hit.
#'   Default `3`.
#' @param direction Tail to test: `"greater"` (default, high activity),
#'   `"less"` (low activity) or `"two.sided"` (either tail, scored on `abs(z)`).
#' @param compound_col Name of the compound column. Default `"compound_id"`.
#'
#' @return A data frame with one row per compound: `compound_id`, `activity`,
#'   `sigma` (signed z-score), `sigma_level` (0-3, how many sigma it clears in
#'   the tested direction), `band` (an ordered-factor label) and `is_hit`
#'   (`sigma_level >= n_sigma`). The distribution `center`, `scale`, `method`,
#'   `n_sigma` and `direction` are stored as attributes.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' hits <- select_sigma_hits(norm, n_sigma = 3)
#' sum(hits$is_hit)
#'
#' @seealso [summarize_sigma_hits()], [plot_sigma_hits()], [select_primary_hit()]
#' @export
select_sigma_hits <- function(data, value = "inhibition",
                              method = c("robust", "sd"),
                              n_sigma = 3,
                              direction = c("greater", "less", "two.sided"),
                              compound_col = "compound_id") {
  method <- match.arg(method)
  direction <- match.arg(direction)
  act <- .compound_activity(data, value, compound_col)
  x <- .as_numeric(act$activity)
  if (method == "robust") {
    center <- stats::median(x, na.rm = TRUE)
    scale <- stats::mad(x, center = center, constant = 1.4826, na.rm = TRUE)
  } else {
    center <- mean(x, na.rm = TRUE)
    scale <- stats::sd(x, na.rm = TRUE)
  }
  z <- if (is.finite(scale) && scale > 0) (x - center) / scale else
    rep(NA_real_, length(x))
  dev <- switch(direction, greater = z, less = -z, two.sided = abs(z))
  lvl <- ifelse(is.finite(dev), pmin(3L, pmax(0L, as.integer(floor(dev)))),
    NA_integer_)
  band_lab <- c("< 1 sigma", ">= 1 sigma", ">= 2 sigma", ">= 3 sigma")
  act$sigma <- z
  act$sigma_level <- lvl
  act$band <- factor(band_lab[lvl + 1L], levels = band_lab, ordered = TRUE)
  act$is_hit <- is.finite(dev) & lvl >= n_sigma
  act <- act[order(-dev), , drop = FALSE]
  rownames(act) <- NULL
  attr(act, "center") <- unname(center)
  attr(act, "scale") <- unname(scale)
  attr(act, "method") <- method
  attr(act, "n_sigma") <- n_sigma
  attr(act, "direction") <- direction
  act
}

#' Summarise sigma-based hits
#'
#' Counts how many compounds clear 1, 2 and 3 sigma from the library centre and
#' compares each count to what a perfectly normal, hit-free distribution would
#' predict, so enrichment of the tail over chance is obvious.
#'
#' @param hits A data frame from [select_sigma_hits()].
#'
#' @return A data frame with one row per sigma level (1, 2, 3): `sigma`,
#'   `n_beyond` (compounds clearing it), `frac_beyond`, `expected_frac`
#'   (normal-theory tail probability) and `expected_n` (expected count).
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' summarize_sigma_hits(select_sigma_hits(norm))
#'
#' @seealso [select_sigma_hits()], [plot_sigma_hits()]
#' @export
summarize_sigma_hits <- function(hits) {
  if (!"sigma_level" %in% names(hits)) {
    stop("`hits` must come from `select_sigma_hits()`.", call. = FALSE)
  }
  direction <- attr(hits, "direction")
  if (is.null(direction)) direction <- "greater"
  n <- sum(!is.na(hits$sigma_level))
  tail_p <- function(k) {
    if (direction == "two.sided") 2 * stats::pnorm(k, lower.tail = FALSE) else
      stats::pnorm(k, lower.tail = FALSE)
  }
  ks <- 1:3
  n_beyond <- vapply(ks,
    function(k) sum(hits$sigma_level >= k, na.rm = TRUE), integer(1))
  exp_frac <- vapply(ks, tail_p, numeric(1))
  data.frame(
    sigma = ks,
    n_beyond = n_beyond,
    frac_beyond = if (n > 0) n_beyond / n else NA_real_,
    expected_frac = exp_frac,
    expected_n = exp_frac * n,
    stringsAsFactors = FALSE
  )
}

#' Distribution plot of sigma-based hits
#'
#' Histogram of the per-compound activity with the library centre and the 1, 2
#' and 3 sigma cut-offs marked, bars coloured by how many sigma each compound
#' clears. Optionally overlays the fitted normal curve so the departure from
#' normality (the hit tail) stands out.
#'
#' @param hits A data frame from [select_sigma_hits()].
#' @param bins Number of histogram bins. Default `30`.
#' @param show_normal Overlay the fitted normal density curve. Default `TRUE`.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' plot_sigma_hits(select_sigma_hits(norm))
#'
#' @seealso [select_sigma_hits()], [summarize_sigma_hits()]
#' @export
plot_sigma_hits <- function(hits, bins = 30, show_normal = TRUE) {
  if (!all(c("activity", "band") %in% names(hits))) {
    stop("`hits` must come from `select_sigma_hits()`.", call. = FALSE)
  }
  center <- attr(hits, "center")
  scale <- attr(hits, "scale")
  direction <- attr(hits, "direction")
  if (is.null(direction)) direction <- "greater"
  x <- .as_numeric(hits$activity)
  df <- data.frame(activity = x, band = hits$band)
  rng <- range(x[is.finite(x)])
  bw <- (rng[2] - rng[1]) / bins
  if (!is.finite(bw) || bw <= 0) bw <- 1
  sigma_cols <- c("< 1 sigma" = "#B4B4B4", ">= 1 sigma" = "#56B4E9",
    ">= 2 sigma" = "#E69F00", ">= 3 sigma" = "#D55E00")
  ks <- 1:3
  pos <- switch(direction,
    greater = center + ks * scale,
    less = center - ks * scale,
    two.sided = c(center - ks * scale, center + ks * scale))
  lab_pos <- switch(direction,
    less = center - ks * scale,
    center + ks * scale)
  p <- ggplot2::ggplot(df, ggplot2::aes(x = .data$activity)) +
    ggplot2::geom_histogram(ggplot2::aes(fill = .data$band), binwidth = bw,
      boundary = center, color = "white", linewidth = 0.2) +
    ggplot2::scale_fill_manual(values = sigma_cols, drop = FALSE,
      name = "Distance from centre") +
    ggplot2::scale_y_continuous(expand = ggplot2::expansion(mult = c(0, 0.08)))
  if (isTRUE(show_normal) && is.finite(scale) && scale > 0) {
    n <- sum(is.finite(x))
    p <- p + ggplot2::stat_function(fun = function(z)
      stats::dnorm(z, mean = center, sd = scale) * n * bw,
      color = .hts_col("neutral"), linewidth = 0.8, na.rm = TRUE)
  }
  p +
    ggplot2::geom_vline(xintercept = center, color = .hts_col("reference"),
      linewidth = 0.6) +
    ggplot2::geom_vline(xintercept = pos, linetype = "dashed",
      color = .hts_col("reference"), linewidth = 0.4) +
    ggplot2::annotate("text", x = lab_pos, y = Inf, label = paste0(ks, " sigma"),
      angle = 90, vjust = -0.4, hjust = 1.05, size = 2.9,
      color = .hts_col("reference")) +
    ggplot2::labs(x = "Per-compound activity", y = "Compounds") +
    theme_hts()
}
