# Module 5: Dose response curve fitting and AUC ------------------------------

# Build a stable per-curve key from the compound and any grouping columns.
.curve_keys <- function(df, compound_col, group_cols) {
  keys <- c(compound_col, group_cols)
  parts <- lapply(keys, function(k) as.character(df[[k]]))
  do.call(paste, c(parts, sep = " | "))
}

#' Import gradient dose-response data
#'
#' Standardises a secondary (dose-response) screening table into the columns
#' expected by the fitting engine, matching each compound to its concentration
#' gradient and response, ready for [fit_4pl_curve()].
#'
#' @param data A data frame, or a path to a csv/txt/xlsx file.
#' @param compound_col,concentration_col,response_col Names of the compound,
#'   concentration and response columns in `data`.
#' @param group_cols Optional grouping columns kept alongside each curve (for
#'   example `"cell_line"`).
#'
#' @return An `hts_dose` object (data frame) with columns `compound_id`,
#'   `concentration`, `response`, any `group_cols`, and canonical class.
#'
#' @examples
#' dr <- import_dose_response(hts_dose_response, response_col = "viability",
#'   group_cols = "cell_line")
#' head(dr)
#'
#' @seealso [fit_4pl_curve()]
#' @export
import_dose_response <- function(data, compound_col = "compound_id",
                                 concentration_col = "concentration",
                                 response_col = "viability",
                                 group_cols = NULL) {
  if (is.character(data) && length(data) == 1) {
    data <- .read_hts_file(data)
  }
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  for (col in c(compound_col, concentration_col, response_col, group_cols)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  out <- data.frame(
    compound_id = as.character(df[[compound_col]]),
    concentration = .as_numeric(df[[concentration_col]]),
    response = .as_numeric(df[[response_col]]),
    stringsAsFactors = FALSE
  )
  for (g in group_cols) out[[g]] <- df[[g]]
  .as_hts(out, "hts_dose")
}

#' Fit four-parameter logistic curves per compound
#'
#' Fits a 4PL dose-response model to every compound (optionally split by
#' grouping columns such as cell line), reusing the package's self-contained
#' fitter, and records the goodness of fit \eqn{R^2}. No external fitting
#' package is required.
#'
#' @param data An `hts_dose` object from [import_dose_response()] (or a long
#'   data frame with `compound_id`, `concentration`, `response`).
#' @param response,concentration,compound_col Column names. Default
#'   `"response"`, `"concentration"`, `"compound_id"`.
#' @param group_cols Optional grouping columns.
#' @param rsq_min Curves with \eqn{R^2} below this are flagged (not dropped
#'   here; use [filter_low_quality_curve()]). Default `0.85`.
#'
#' @return An `hts_drc` object: a list with `meta` (one row per curve: keys,
#'   `n`, `rsq`, `converged`, `pass`), `fits` (named list of `dsr_4pl` models)
#'   and `data` (the input). `group_cols` is stored as an attribute.
#'
#' @examples
#' dr <- import_dose_response(hts_dose_response, response_col = "viability",
#'   group_cols = "cell_line")
#' drc <- fit_4pl_curve(dr, group_cols = "cell_line")
#' head(drc$meta)
#'
#' @seealso [extract_drc_params()], [calc_drc_auc()], [plot_single_drc()]
#' @export
fit_4pl_curve <- function(data, response = "response",
                          concentration = "concentration",
                          compound_col = "compound_id", group_cols = NULL,
                          rsq_min = 0.85) {
  df <- as.data.frame(data, stringsAsFactors = FALSE)
  for (col in c(response, concentration, compound_col, group_cols)) {
    if (!col %in% names(df)) {
      stop("Column '", col, "' not found in `data`.", call. = FALSE)
    }
  }
  df$.curve <- .curve_keys(df, compound_col, group_cols)
  keys <- unique(df$.curve)

  fits <- list()
  meta_rows <- list()
  for (k in keys) {
    sub <- df[df$.curve == k, , drop = FALSE]
    x <- .as_numeric(sub[[concentration]])
    y <- .as_numeric(sub[[response]])
    ok <- is.finite(x) & is.finite(y)
    x <- x[ok]
    y <- y[ok]
    row <- data.frame(curve_id = k,
      compound_id = as.character(sub[[compound_col]][1]),
      stringsAsFactors = FALSE)
    for (g in group_cols) row[[g]] <- sub[[g]][1]
    row$n <- length(x)
    if (length(x) < 4 || length(unique(x)) < 4) {
      row$rsq <- NA_real_
      row$converged <- FALSE
      fits[[k]] <- NULL
    } else {
      fit <- tryCatch(fit_dose_response(
        data.frame(response = y, concentration = x),
        "response", "concentration"), error = function(e) NULL)
      if (is.null(fit)) {
        row$rsq <- NA_real_
        row$converged <- FALSE
      } else {
        sst <- sum((y - mean(y))^2)
        rsq <- if (sst > 0) 1 - fit$sse / sst else NA_real_
        row$rsq <- rsq
        row$converged <- isTRUE(fit$convergence)
        fits[[k]] <- fit
      }
    }
    meta_rows[[k]] <- row
  }
  meta <- do.call(rbind, meta_rows)
  meta$pass <- is.finite(meta$rsq) & meta$rsq >= rsq_min
  rownames(meta) <- NULL

  structure(list(meta = meta, fits = fits,
    data = .as_hts(df[, setdiff(names(df), ".curve"), drop = FALSE],
      "hts_dose")),
    group_cols = group_cols, rsq_min = rsq_min, class = "hts_drc")
}

#' Extract dose-response parameters
#'
#' Pulls the fitted potency and efficacy parameters from every curve of an
#' `hts_drc` object.
#'
#' @param drc An `hts_drc` object from [fit_4pl_curve()].
#'
#' @return A data frame with one row per curve: keys plus `ic50`, `emax` (`top`),
#'   `bottom`, `hill`, `rsq` and `pass`. Curves that failed to fit yield `NA`
#'   parameters.
#'
#' @examples
#' drc <- fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line")
#' head(extract_drc_params(drc))
#'
#' @seealso [fit_4pl_curve()], [calc_drc_auc()]
#' @export
extract_drc_params <- function(drc) {
  stopifnot(inherits(drc, "hts_drc"))
  meta <- drc$meta
  params <- lapply(meta$curve_id, function(k) {
    fit <- drc$fits[[k]]
    if (is.null(fit)) {
      return(data.frame(ic50 = NA_real_, emax = NA_real_, bottom = NA_real_,
        hill = NA_real_))
    }
    cf <- stats::coef(fit)
    data.frame(ic50 = cf[["ic50"]], emax = cf[["top"]], bottom = cf[["bottom"]],
      hill = cf[["hill"]])
  })
  out <- cbind(meta, do.call(rbind, params))
  out[, c(setdiff(names(meta), c("converged", "pass")),
    "ic50", "emax", "bottom", "hill", "pass")]
}

#' Area under the dose-response curve
#'
#' Computes the area under each fitted curve by trapezoidal integration of the
#' fitted 4PL over the tested concentration range (on the log10 scale), a
#' robust single-number summary of overall potency.
#'
#' @param drc An `hts_drc` object from [fit_4pl_curve()].
#'
#' @return The `hts_drc` object with an `auc` column added to its `meta`, also
#'   returned invisibly as a data frame via `drc$meta`.
#'
#' @examples
#' drc <- fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line")
#' calc_drc_auc(drc)$meta
#'
#' @seealso [compute_auc()], [extract_drc_params()]
#' @export
calc_drc_auc <- function(drc) {
  stopifnot(inherits(drc, "hts_drc"))
  auc <- vapply(drc$meta$curve_id, function(k) {
    fit <- drc$fits[[k]]
    if (is.null(fit)) return(NA_real_)
    x <- fit$data$concentration
    x <- x[x > 0]
    if (length(x) < 2) return(NA_real_)
    lx <- log10(range(x))
    grid_l <- seq(lx[1], lx[2], length.out = 200)
    grid <- 10^grid_l
    yhat <- stats::predict(fit, newdata = grid)
    # Trapezoid over the log10-concentration axis.
    sum(diff(grid_l) * (utils::head(yhat, -1) + utils::tail(yhat, -1)) / 2)
  }, numeric(1))
  drc$meta$auc <- unname(auc)
  drc
}

#' Drop low-quality dose-response curves
#'
#' Removes curves that failed to fit or whose \eqn{R^2} falls below the
#' threshold, leaving only trustworthy dose-response models.
#'
#' @param drc An `hts_drc` object.
#' @param rsq_min Minimum acceptable \eqn{R^2}. Defaults to the value stored at
#'   fit time.
#'
#' @return The filtered `hts_drc` object; removed curve ids are stored in
#'   `attr(drc, "removed")`.
#'
#' @examples
#' drc <- fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line")
#' clean <- filter_low_quality_curve(drc)
#' nrow(clean$meta)
#'
#' @seealso [fit_4pl_curve()]
#' @export
filter_low_quality_curve <- function(drc, rsq_min = NULL) {
  stopifnot(inherits(drc, "hts_drc"))
  if (is.null(rsq_min)) rsq_min <- attr(drc, "rsq_min")
  if (is.null(rsq_min)) rsq_min <- 0.85
  keep <- is.finite(drc$meta$rsq) & drc$meta$rsq >= rsq_min
  removed <- drc$meta$curve_id[!keep]
  drc$meta <- drc$meta[keep, , drop = FALSE]
  rownames(drc$meta) <- NULL
  drc$fits <- drc$fits[drc$meta$curve_id]
  attr(drc, "removed") <- removed
  drc
}

# Assemble points + fitted line for one curve.
.drc_curve_frame <- function(drc, curve_id, n_points = 200) {
  fit <- drc$fits[[curve_id]]
  if (is.null(fit)) return(NULL)
  obs <- fit$data
  x <- obs$concentration[obs$concentration > 0]
  grid <- 10^seq(log10(min(x)), log10(max(x)), length.out = n_points)
  list(
    obs = obs[obs$concentration > 0, , drop = FALSE],
    fit = data.frame(concentration = grid,
      response = stats::predict(fit, newdata = grid))
  )
}

#' Plot a single dose-response curve
#'
#' Plots one compound's observed points and fitted 4PL curve on a log10
#' concentration axis, annotating IC50, AUC and \eqn{R^2}.
#'
#' @param drc An `hts_drc` object (ideally after [calc_drc_auc()]).
#' @param curve_id Curve identifier (from `drc$meta$curve_id`). Defaults to the
#'   first curve.
#' @param color Curve colour. Default `"#0072B2"`.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @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"))
#' plot_single_drc(drc, drc$meta$curve_id[1])
#'
#' @seealso [plot_batch_drc_overlay()]
#' @export
plot_single_drc <- function(drc, curve_id = NULL, color = "#0072B2") {
  stopifnot(inherits(drc, "hts_drc"))
  if (is.null(curve_id)) curve_id <- drc$meta$curve_id[1]
  cf <- .drc_curve_frame(drc, curve_id)
  if (is.null(cf)) stop("No fitted model for curve '", curve_id, "'.",
    call. = FALSE)
  meta <- drc$meta[drc$meta$curve_id == curve_id, , drop = FALSE]
  fit <- drc$fits[[curve_id]]
  ic50 <- stats::coef(fit)[["ic50"]]
  auc <- if ("auc" %in% names(meta)) meta$auc[1] else NA_real_
  rsq <- meta$rsq[1]
  subtitle <- sprintf("IC50: %.3g   AUC: %s   R^2: %.3f", ic50,
    if (is.na(auc)) "NA" else sprintf("%.2f", auc), rsq)

  p <- ggplot2::ggplot() +
    ggplot2::geom_point(data = cf$obs,
      ggplot2::aes(x = .data$concentration, y = .data$response),
      color = color, size = 2.5, alpha = 0.85, stroke = 0.2) +
    ggplot2::geom_line(data = cf$fit,
      ggplot2::aes(x = .data$concentration, y = .data$response),
      color = color, linewidth = 1.1, lineend = "round") +
    ggplot2::geom_vline(xintercept = ic50, linetype = "dashed",
      color = .hts_col("reference"), linewidth = 0.6) +
    ggplot2::scale_x_log10() +
    ggplot2::labs(x = "Concentration", y = "Response",
      title = curve_id, subtitle = subtitle) +
    theme_hts()
  p
}

#' Overlay multiple dose-response curves
#'
#' Overlays the fitted curves of several compounds for side-by-side comparison
#' of activity.
#'
#' @param drc An `hts_drc` object.
#' @param curve_ids Curve identifiers to overlay. Defaults to all fitted curves.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @examples
#' drc <- fit_4pl_curve(import_dose_response(hts_dose_response,
#'   response_col = "viability", group_cols = "cell_line"),
#'   group_cols = "cell_line")
#' plot_batch_drc_overlay(drc, drc$meta$curve_id[1:4])
#'
#' @seealso [plot_single_drc()]
#' @export
plot_batch_drc_overlay <- function(drc, curve_ids = NULL) {
  stopifnot(inherits(drc, "hts_drc"))
  if (is.null(curve_ids)) curve_ids <- names(drc$fits)
  curve_ids <- curve_ids[!vapply(drc$fits[curve_ids], is.null, logical(1))]
  if (length(curve_ids) == 0) {
    stop("No fitted curves available to overlay.", call. = FALSE)
  }
  parts <- lapply(curve_ids, function(k) {
    cf <- .drc_curve_frame(drc, k)
    if (is.null(cf)) return(NULL)
    cf$fit$curve <- k
    cf$fit
  })
  line_df <- do.call(rbind, parts)
  ggplot2::ggplot(line_df, ggplot2::aes(x = .data$concentration,
    y = .data$response, color = .data$curve)) +
    ggplot2::geom_line(linewidth = 1, lineend = "round") +
    ggplot2::scale_x_log10() +
    ggplot2::scale_color_viridis_d(name = NULL, option = "D", end = 0.9) +
    ggplot2::labs(x = "Concentration", y = "Response") +
    theme_hts(legend = "right")
}

#' IC50 versus AUC scatter
#'
#' Scatter plot of IC50 against AUC across curves, a compact view for triaging
#' strong, moderate and weak compounds. IC50 marks the potency (the
#' concentration for half-maximal effect) while AUC integrates potency *and*
#' efficacy over the whole tested range; plotting them together shows whether the
#' two agree and flags compounds that are potent but only weakly efficacious (or
#' vice versa). A regression trend line and the Pearson correlation (computed on
#' the log10 IC50, matching the log x-axis) are drawn on the figure by default.
#'
#' @param params A data frame from [extract_drc_params()] joined with an `auc`
#'   column (for example `merge(extract_drc_params(drc), calc_drc_auc(drc)$meta)`),
#'   or the `meta` of an `hts_drc` after [calc_drc_auc()] plus IC50.
#' @param auc_col,ic50_col Column names. Default `"auc"`, `"ic50"`.
#' @param add_trend Add a linear regression trend line with a 95% confidence
#'   band. Default `TRUE`.
#' @param annotate_cor Annotate the Pearson correlation coefficient and p-value
#'   on the plot. Default `TRUE`.
#'
#' @return A [ggplot2::ggplot] object.
#'
#' @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"))
#' params <- merge(extract_drc_params(drc), drc$meta[, c("curve_id", "auc")])
#' plot_ic50_auc_cor(params)
#'
#' @seealso [extract_drc_params()], [calc_drc_auc()]
#' @export
plot_ic50_auc_cor <- function(params, auc_col = "auc", ic50_col = "ic50",
                              add_trend = TRUE, annotate_cor = TRUE) {
  for (col in c(auc_col, ic50_col)) {
    if (!col %in% names(params)) {
      stop("Column '", col, "' not found in `params`.", call. = FALSE)
    }
  }
  df <- params[is.finite(params[[auc_col]]) & is.finite(params[[ic50_col]]) &
    params[[ic50_col]] > 0, , drop = FALSE]
  p <- ggplot2::ggplot(df,
    ggplot2::aes(x = .data[[ic50_col]], y = .data[[auc_col]]))
  if (isTRUE(add_trend) && nrow(df) > 2) {
    p <- p + ggplot2::geom_smooth(method = "lm", formula = y ~ x, se = TRUE,
      color = .hts_col("reference"), fill = "#E5E5E5", linewidth = 0.8)
  }
  p <- p +
    ggplot2::geom_point(color = .hts_col("accent"), size = 2.6, alpha = 0.8,
      stroke = 0.2) +
    ggplot2::scale_x_log10() +
    ggplot2::labs(x = "IC50", y = "AUC") +
    theme_hts()
  if (isTRUE(annotate_cor) && nrow(df) > 2) {
    ct <- suppressWarnings(stats::cor.test(log10(df[[ic50_col]]),
      df[[auc_col]], method = "pearson"))
    lab <- sprintf("R = %.2f, p %s", unname(ct$estimate),
      if (ct$p.value < 0.001) "< 0.001" else sprintf("= %.3f", ct$p.value))
    x_pos <- min(df[[ic50_col]], na.rm = TRUE)
    y_pos <- max(df[[auc_col]], na.rm = TRUE)
    p <- p + ggplot2::annotate("text", x = x_pos, y = y_pos, label = lab,
      hjust = 0, vjust = 1, size = 3.4, color = "#333333")
  }
  p
}
