# Module 3: Replicate consistency analysis -----------------------------------

# Mean value per compound per plate (compound wells only).
.compound_plate_means <- function(data, value, compound_col, plate_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, plate_col)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  agg <- stats::aggregate(
    df[[value]],
    by = list(compound_id = df[[compound_col]], plate_id = df[[plate_col]]),
    FUN = function(z) mean(.as_numeric(z), na.rm = TRUE)
  )
  names(agg)[3] <- "value"
  agg
}

#' Replicate coefficient of variation per compound
#'
#' Computes, for each compound, the coefficient of variation (CV, percent)
#' across its technical replicates (repeated wells, whether on the same plate or
#' across plates), a standard measure of assay reproducibility.
#'
#' @param data An `hts_norm` object or a long data frame of well-level values.
#' @param value Name of the value column on which to compute the CV. Default
#'   `"viability"` (more stable than inhibition, which can sit near zero).
#' @param compound_col Name of the compound column. Default `"compound_id"`.
#' @param group_cols Optional additional grouping columns (for example
#'   `"cell_line"`) that define what counts as one compound series.
#'
#' @return A data frame with one row per compound (and any `group_cols`):
#'   `compound_id`, the grouping columns, `n`, `mean`, `sd` and `cv` (percent).
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' cv <- calc_replicate_cv(norm)
#' head(cv[order(-cv$cv), ])
#'
#' @seealso [calc_replicate_correlation()], [filter_bad_replicate()]
#' @export
calc_replicate_cv <- function(data, value = "viability",
                              compound_col = "compound_id",
                              group_cols = NULL) {
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  if ("well_type" %in% names(df)) {
    df <- df[df$well_type == "compound", , drop = FALSE]
  }
  for (col in c(value, compound_col, group_cols)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  keys <- c(compound_col, group_cols)
  by_list <- stats::setNames(lapply(keys, function(k) df[[k]]), keys)
  v <- .as_numeric(df[[value]])
  agg_mean <- stats::aggregate(v, by_list, function(z) mean(z, na.rm = TRUE))
  agg_sd <- stats::aggregate(v, by_list, function(z) stats::sd(z, na.rm = TRUE))
  agg_n <- stats::aggregate(v, by_list, function(z) sum(!is.na(z)))
  out <- agg_mean
  names(out)[length(keys) + 1] <- "mean"
  out$sd <- agg_sd$x
  out$n <- agg_n$x
  out$cv <- 100 * out$sd / abs(out$mean)
  names(out)[1] <- "compound_id"
  out[, c("compound_id", group_cols, "n", "mean", "sd", "cv")]
}

#' Pairwise replicate correlation between plates
#'
#' Computes the Pearson correlation (and its \eqn{R^2}) between the per-compound
#' activity of every pair of replicate plates, the standard check that technical
#' replicates agree.
#'
#' @param data An `hts_norm` object or long data frame.
#' @param value Name of the value column. Default `"viability"`.
#' @param compound_col Name of the compound column. Default `"compound_id"`.
#' @param plate_col Name of the plate column. Default `"plate_id"`.
#'
#' @return A data frame with one row per plate pair: `plate_x`, `plate_y`, `n`,
#'   `pearson_r` and `r_squared`. The wide compound-by-plate matrix is attached
#'   as `attr(x, "wide")` for use by [plot_replicate_scatter()].
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' calc_replicate_correlation(norm)
#'
#' @seealso [calc_replicate_cv()], [plot_replicate_scatter()]
#' @export
calc_replicate_correlation <- function(data, value = "viability",
                                       compound_col = "compound_id",
                                       plate_col = "plate_id") {
  means <- .compound_plate_means(data, value, compound_col, plate_col)
  wide <- tidyr::pivot_wider(means, names_from = "plate_id",
    values_from = "value")
  wide <- as.data.frame(wide)
  plates <- setdiff(names(wide), "compound_id")
  if (length(plates) < 2) {
    stop("At least two plates are required to compute replicate correlation.",
      call. = FALSE)
  }
  pairs <- utils::combn(plates, 2)
  res <- lapply(seq_len(ncol(pairs)), function(j) {
    a <- pairs[1, j]
    b <- pairs[2, j]
    ok <- stats::complete.cases(wide[[a]], wide[[b]])
    r <- if (sum(ok) > 2) stats::cor(wide[[a]][ok], wide[[b]][ok]) else NA_real_
    data.frame(plate_x = a, plate_y = b, n = sum(ok),
      pearson_r = r, r_squared = r^2, stringsAsFactors = FALSE)
  })
  out <- do.call(rbind, res)
  attr(out, "wide") <- wide
  out
}

#' Drop compounds with poor replicate reproducibility
#'
#' Removes compounds whose replicate coefficient of variation exceeds a
#' threshold. When a correlation table is supplied it also checks that the
#' replicate plates are globally concordant, warning if any pair falls below the
#' \eqn{R^2} threshold.
#'
#' @param cv A data frame from [calc_replicate_cv()].
#' @param cv_max Maximum acceptable CV (percent). Default `15`.
#' @param cor_result Optional data frame from [calc_replicate_correlation()].
#' @param r2_min Minimum acceptable pairwise \eqn{R^2}. Default `0.8`.
#'
#' @return The subset of `cv` for compounds passing the CV threshold, with the
#'   removed compound ids stored in `attr(x, "removed")`.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' cv <- calc_replicate_cv(norm)
#' good <- filter_bad_replicate(cv, cv_max = 20)
#' attr(good, "removed")
#'
#' @seealso [calc_replicate_cv()], [calc_replicate_correlation()]
#' @export
filter_bad_replicate <- function(cv, cv_max = 15, cor_result = NULL,
                                 r2_min = 0.8) {
  if (!"cv" %in% names(cv)) {
    stop("`cv` must contain a 'cv' column (see `calc_replicate_cv()`).",
      call. = FALSE)
  }
  if (!is.null(cor_result) && "r_squared" %in% names(cor_result)) {
    bad_pairs <- cor_result[which(cor_result$r_squared < r2_min), , drop = FALSE]
    if (nrow(bad_pairs) > 0) {
      warning(nrow(bad_pairs), " plate pair(s) below R^2 = ", r2_min,
        "; replicate concordance is weak.", call. = FALSE)
    }
  }
  keep <- is.finite(cv$cv) & cv$cv <= cv_max
  out <- cv[keep, , drop = FALSE]
  attr(out, "removed") <- cv$compound_id[!keep]
  rownames(out) <- NULL
  out
}

# Default well-type -> colour map, mirroring the report palette
# (compounds blue, negative controls black, positive controls red, blanks gold).
.well_type_palette <- c(
  compound = "#0072B2", Experimental = "#0072B2", sample = "#0072B2",
  NC = "black", negative_ctrl = "black", DMSO = "black",
  PC = "#D55E00", positive_ctrl = "#D55E00", kill_ctrl = "#D55E00",
  Blank = "#E1B000", edgewell = "#E1B000", empty = "#E1B000"
)

# Build a compound/well-by-plate wide table, optionally retaining a type column
# so points can be coloured. Replicate plates share a layout, so wells are paired
# by their address; this keeps the control wells in the scatter, as in the
# source reps_corr_plot() workflow.
.replicate_wide <- function(data, value, plate_col, key_col, type_col) {
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  for (col in c(value, plate_col, key_col)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  v <- .as_numeric(df[[value]])
  by <- list(key = df[[key_col]], plate = df[[plate_col]])
  if (!is.null(type_col)) by <- c(by, list(type = df[[type_col]]))
  agg <- stats::aggregate(v, by, function(z) mean(z, na.rm = TRUE))
  names(agg)[ncol(agg)] <- "value"
  id_cols <- if (is.null(type_col)) "key" else c("key", "type")
  wide <- tidyr::pivot_wider(agg, id_cols = tidyr::all_of(id_cols),
    names_from = "plate", values_from = "value")
  as.data.frame(wide, stringsAsFactors = FALSE)
}

#' Replicate scatter plot(s) with regression fit and correlation
#'
#' Plots the activity of replicate plates against one another in the house style
#' used across the screening reports: points coloured by well type (compounds,
#' negative/positive controls, blanks) at 50% opacity, a linear regression line
#' with a shaded 95% confidence band, and a Pearson \eqn{R} (with p-value)
#' annotated on each panel. When more than two replicate plates are present and
#' no specific pair is requested, every pairwise combination (replicate 1 vs 2,
#' 1 vs 3, 2 vs 3, ...) is drawn together in a single faceted figure, mirroring
#' the `combn()`-based `reps_corr_plot()` workflow in the source analyses.
#'
#' @param data An `hts_norm`/long data frame, or a data frame from
#'   [calc_replicate_correlation()] (its wide matrix is used).
#' @param plate_x,plate_y Plate identifiers for the two axes. When both are
#'   `NULL` (default) all pairwise combinations of the available plates are
#'   drawn; supply both to draw a single pair.
#' @param value Name of the value column when `data` is well-level. Default
#'   `"viability"`.
#' @param plate_col Name of the plate column. Default `"plate_id"`.
#' @param key_col Column identifying matching observations across replicate
#'   plates. Defaults to the well-address column `"well"` when present (so that
#'   control wells are paired and displayed), otherwise the compound column.
#' @param compound_col Compound column, used as the pairing key when no `well`
#'   column is available. Default `"compound_id"`.
#' @param color_col Column used to colour points, typically the well-role column.
#'   Defaults to `"well_type"` when present; set to `NULL` to disable colouring.
#' @param palette Named vector mapping well-type levels to colours. Defaults to
#'   the report palette (compounds blue, negative controls black, positive
#'   controls red, blanks gold).
#' @param line_color,band_color Colour of the regression line and of its 95%
#'   confidence band. Default light blue on light grey, as in the reports.
#' @param show_identity Whether to draw the dotted \eqn{y = x} identity line.
#'   Default `TRUE`.
#'
#' @return A single [ggplot2::ggplot] object. With three or more plates it is
#'   faceted, one panel per replicate pair. The annotated \eqn{R} and p-value are
#'   computed over every point shown in the panel.
#'
#' @examples
#' norm <- norm_by_control(baseline_subtract(read_hts_plate(hts_primary_raw)))
#' # All pairwise replicate comparisons at once, points coloured by well type:
#' plot_replicate_scatter(norm)
#' # A single chosen pair:
#' plot_replicate_scatter(norm, plate_x = "P01", plate_y = "P02")
#'
#' @seealso [calc_replicate_correlation()]
#' @export
plot_replicate_scatter <- function(data, plate_x = NULL, plate_y = NULL,
                                    value = "viability",
                                    plate_col = "plate_id",
                                    key_col = NULL,
                                    compound_col = "compound_id",
                                    color_col = NULL,
                                    palette = .well_type_palette,
                                    line_color = "lightblue",
                                    band_color = "lightgray",
                                    show_identity = TRUE) {
  wide <- attr(data, "wide")
  type_present <- FALSE
  if (is.null(wide)) {
    df <- as.data.frame(data, stringsAsFactors = FALSE)
    if (is.null(color_col) && "well_type" %in% names(df)) color_col <- "well_type"
    if (is.null(key_col)) {
      key_col <- if ("well" %in% names(df)) "well" else compound_col
    }
    type_col <- if (!is.null(color_col) && color_col %in% names(df)) color_col else NULL
    wide <- .replicate_wide(df, value, plate_col, key_col, type_col)
    type_present <- !is.null(type_col)
  }
  reserved <- c("key", "type", "compound_id")
  plates <- setdiff(names(wide), reserved)
  if (length(plates) < 2) {
    stop("At least two plates are required for a replicate scatter plot.",
      call. = FALSE)
  }

  # Which plate pairs to draw: one requested pair, or every combination.
  if (!is.null(plate_x) || !is.null(plate_y)) {
    if (is.null(plate_x)) plate_x <- plates[1]
    if (is.null(plate_y)) plate_y <- plates[2]
    pairs <- matrix(c(plate_x, plate_y), nrow = 2)
  } else {
    pairs <- utils::combn(plates, 2)
  }

  type_vals <- if (type_present) wide$type else rep(NA_character_, nrow(wide))
  panels <- vector("list", ncol(pairs))
  labels <- vector("list", ncol(pairs))
  for (j in seq_len(ncol(pairs))) {
    a <- pairs[1, j]
    b <- pairs[2, j]
    for (p in c(a, b)) {
      if (!p %in% names(wide)) stop("Plate '", p, "' not found.", call. = FALSE)
    }
    lab_pair <- paste0(a, " vs ", b)
    d <- data.frame(x = wide[[a]], y = wide[[b]], type = type_vals,
      pair = lab_pair, stringsAsFactors = FALSE)
    d <- d[stats::complete.cases(d[, c("x", "y")]), , drop = FALSE]
    panels[[j]] <- d
    if (nrow(d) > 2) {
      ct <- stats::cor.test(d$x, d$y, method = "pearson")
      r <- unname(ct$estimate)
      p_lab <- if (ct$p.value < 0.001) "p < 0.001" else
        sprintf("p = %.3f", ct$p.value)
      lab <- sprintf("R = %.2f, %s", r, p_lab)
    } else {
      lab <- "R = NA"
    }
    labels[[j]] <- data.frame(pair = lab_pair, label = lab,
      stringsAsFactors = FALSE)
  }
  df <- do.call(rbind, panels)
  lab_df <- do.call(rbind, labels)
  # Preserve pair ordering for the facets.
  lev <- vapply(seq_len(ncol(pairs)),
    function(j) paste0(pairs[1, j], " vs ", pairs[2, j]), character(1))
  df$pair <- factor(df$pair, levels = lev)
  lab_df$pair <- factor(lab_df$pair, levels = lev)

  p <- ggplot2::ggplot(df, ggplot2::aes(x = .data$x, y = .data$y))
  if (isTRUE(show_identity)) {
    p <- p + ggplot2::geom_abline(slope = 1, intercept = 0,
      linetype = "dotted", color = "grey60")
  }
  if (type_present) {
    p <- p + ggplot2::geom_point(ggplot2::aes(color = .data$type),
      alpha = 0.5, size = 1.6)
    types <- sort(unique(stats::na.omit(df$type)))
    pal <- palette[types]
    if (anyNA(pal)) {
      pal[is.na(pal)] <- hts_pal(sum(is.na(pal)))
    }
    names(pal) <- types
    p <- p + ggplot2::scale_color_manual(values = pal, name = NULL,
      na.translate = FALSE)
  } else {
    p <- p + ggplot2::geom_point(alpha = 0.6, color = "#0072B2", size = 1.8,
      stroke = 0.2)
  }
  p <- p +
    ggplot2::geom_smooth(method = "lm", formula = y ~ x, se = TRUE,
      color = line_color, fill = band_color, linewidth = 0.8) +
    ggplot2::geom_text(data = lab_df, ggplot2::aes(label = .data$label),
      x = -Inf, y = Inf, hjust = -0.08, vjust = 1.4, size = 3.4,
      inherit.aes = FALSE) +
    theme_hts(legend = "right")

  if (ncol(pairs) > 1) {
    p <- p +
      ggplot2::facet_wrap(~ .data$pair, scales = "free") +
      ggplot2::labs(x = "Replicate (x-axis)", y = "Replicate (y-axis)")
  } else {
    p <- p + ggplot2::labs(x = pairs[1, 1], y = pairs[2, 1])
  }
  p
}

#' Histogram of replicate CV values
#'
#' Shows the distribution of per-compound replicate CVs across the library, with
#' the acceptance threshold marked.
#'
#' @param cv A data frame from [calc_replicate_cv()].
#' @param cv_max Threshold to draw as a reference line. Default `15`.
#' @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_cv_distribution(calc_replicate_cv(norm))
#'
#' @seealso [calc_replicate_cv()], [filter_bad_replicate()]
#' @export
plot_cv_distribution <- function(cv, cv_max = 15, bins = 30) {
  if (!"cv" %in% names(cv)) {
    stop("`cv` must contain a 'cv' column.", call. = FALSE)
  }
  df <- cv[is.finite(cv$cv), , drop = FALSE]
  ggplot2::ggplot(df, ggplot2::aes(x = .data$cv)) +
    ggplot2::geom_histogram(bins = bins, fill = .hts_col("compound"),
      color = "white", linewidth = 0.2, alpha = 0.9) +
    ggplot2::geom_vline(xintercept = cv_max, linetype = "dashed",
      color = .hts_col("reference"), linewidth = 0.6) +
    ggplot2::labs(x = "Replicate CV (%)", y = "Compounds") +
    theme_hts()
}
