## Cold-start benchmark that separates loading from computing.
##
## Each method runs in a fresh R process. The child records four phases:
##   startup  R boot (measured by the parent as total minus what the child saw)
##   lib_s    attaching packages
##   load_s   getting the data into the form the method computes on
##   compute_s producing the 2,778-row answer from it
##
## For file-backed engines load_s is only the cost of opening a connection or
## building a lazy scan; the actual reading happens inside compute_s, which is
## the entire point of using them.

PQ      <- "C:/Users/Matthew/Downloads/polars_and_duckdb_tutorial/fake_pcawg_all.parquet"
THREADS <- 8L
GB      <- 2^30
OUT     <- commandArgs(trailingOnly = TRUE)[1]

suppressPackageStartupMessages({ library(callr); library(ps) })

child <- function(method, pq, threads) {
  Sys.setenv(POLARS_MAX_THREADS = as.character(threads))
  t0 <- proc.time()[["elapsed"]]

  pkgs <- switch(method,
    "base R (aggregate)"        = c("arrow"),
    "dplyr"                     = c("arrow", "dplyr"),
    "data.table"                = c("arrow", "data.table"),
    "collapse"                  = c("arrow", "collapse"),
    "polars (eager)"            = c("polars"),
    "polars (streaming scan)"   = c("polars"),
    "DuckDB (Parquet)"          = c("DBI", "duckdb"),
    "duckplyr"                  = c("duckplyr", "dplyr"),
    "arrow (Acero)"             = c("arrow", "dplyr")
  )
  for (p in pkgs) suppressPackageStartupMessages(
    library(p, character.only = TRUE, quietly = TRUE))
  if ("data.table" %in% pkgs) data.table::setDTthreads(threads)
  if ("arrow"      %in% pkgs) arrow::set_cpu_count(threads)
  t1 <- proc.time()[["elapsed"]]

  ## ---- load: get the data into the shape this method computes on -----------
  d <- switch(method,
    "base R (aggregate)"      = as.data.frame(arrow::read_parquet(pq)),
    "dplyr"                   = as.data.frame(arrow::read_parquet(pq)),
    "data.table"              = data.table::as.data.table(arrow::read_parquet(pq)),
    "collapse"                = as.data.frame(arrow::read_parquet(pq)),
    "polars (eager)"          = polars::pl$read_parquet(pq),
    "polars (streaming scan)" = polars::pl$scan_parquet(pq),
    "DuckDB (Parquet)"        = DBI::dbConnect(duckdb::duckdb(),
                                  config = list(threads = as.character(threads))),
    "duckplyr"                = duckplyr::read_parquet_duckdb(pq),
    "arrow (Acero)"           = arrow::open_dataset(pq)
  )
  ## touch it so lazy handles are really built, without forcing a read
  invisible(class(d))
  t2 <- proc.time()[["elapsed"]]

  ## ---- compute: produce one row per sample ---------------------------------
  res <- switch(method,

    "base R (aggregate)" = {
      cl  <- ave(d$cluster_id, d$samplename, FUN = min)
      isc <- d$cluster_id == cl
      a <- aggregate(list(n_snvs = d$cluster_id),
                     by = list(samplename = d$samplename), FUN = length)
      b <- aggregate(list(n_clonal = isc),
                     by = list(samplename = d$samplename), FUN = sum)
      r <- merge(a, b, by = "samplename")
      r$sMF <- (r$n_snvs - r$n_clonal) / r$n_snvs
      r
    },

    "dplyr" = dplyr::mutate(
        dplyr::summarize(d,
          n_snvs = dplyr::n(),
          n_clonal = sum(cluster_id == min(cluster_id)),
          .by = samplename),
        sMF = (n_snvs - n_clonal) / n_snvs),

    "data.table" = {
      r <- d[, .(n_snvs = .N,
                 n_clonal = sum(cluster_id == min(cluster_id))), by = samplename]
      r[, sMF := (n_snvs - n_clonal) / n_snvs][]
    },

    "collapse" = {
      cl <- collapse::fmin(d$cluster_id, g = d$samplename, TRA = "replace")
      n1 <- collapse::fnobs(d$cluster_id, g = d$samplename)
      n2 <- collapse::fsum(d$cluster_id == cl, g = d$samplename)
      data.frame(samplename = names(n1), n_snvs = as.numeric(n1),
                 n_clonal = as.numeric(n2),
                 sMF = as.numeric((n1 - n2) / n1))
    },

    "polars (eager)" = ,
    "polars (streaming scan)" = {
      pl <- polars::pl
      q <- d$group_by("samplename")$agg(
             pl$len()$alias("n_snvs"),
             (pl$col("cluster_id") == pl$col("cluster_id")$min())$sum()$alias("n_clonal")
           )$with_columns(
             ((pl$col("n_snvs") - pl$col("n_clonal")) / pl$col("n_snvs"))$alias("sMF"))
      if (method == "polars (streaming scan)") q$collect(engine = "streaming") else q
    },

    "DuckDB (Parquet)" = DBI::dbGetQuery(d, sprintf("
        SELECT samplename,
               COUNT(*) AS n_snvs,
               COUNT(*) FILTER (cluster_id = mn) AS n_clonal,
               (COUNT(*) - COUNT(*) FILTER (cluster_id = mn)) * 1.0 / COUNT(*) AS sMF
        FROM (SELECT samplename, cluster_id,
                     MIN(cluster_id) OVER (PARTITION BY samplename) AS mn
              FROM read_parquet('%s'))
        GROUP BY ALL", pq)),

    "duckplyr" = {
      s1 <- dplyr::summarize(d, n = dplyr::n(), .by = c(samplename, cluster_id))
      s2 <- dplyr::mutate(s1, mn = min(cluster_id), .by = samplename)
      s3 <- dplyr::summarize(s2,
              n_snvs = sum(n),
              n_clonal = sum(ifelse(cluster_id == mn, n, 0)), .by = samplename)
      dplyr::collect(dplyr::mutate(s3, sMF = (n_snvs - n_clonal) / n_snvs))
    },

    "arrow (Acero)" = {
      s1 <- dplyr::collect(dplyr::summarize(
              dplyr::group_by(d, samplename, cluster_id),
              n = dplyr::n(), .groups = "drop"))
      s2 <- dplyr::summarize(dplyr::group_by(s1, samplename),
              n_snvs = sum(n), n_clonal = n[which.min(cluster_id)], .groups = "drop")
      dplyr::mutate(s2, sMF = (n_snvs - n_clonal) / n_snvs)
    }
  )

  res <- as.data.frame(res)
  invisible(nrow(res))
  t3 <- proc.time()[["elapsed"]]

  if (method == "DuckDB (Parquet)") try(DBI::dbDisconnect(d, shutdown = TRUE), silent = TRUE)

  list(value = res[order(res$samplename), c("samplename", "n_snvs", "n_clonal", "sMF")],
       lib_s = t1 - t0, load_s = t2 - t1, compute_s = t3 - t2)
}

run_one <- function(method) {
  t_launch <- proc.time()[["elapsed"]]
  p <- callr::r_bg(child, args = list(method = method, pq = PQ, threads = THREADS),
                   supervise = TRUE)
  peak <- 0; n <- 0L; h <- NULL
  repeat {
    if (is.null(h)) h <- tryCatch(ps::ps_handle(p$get_pid()), error = function(e) NULL)
    if (!is.null(h)) {
      m <- tryCatch(ps::ps_memory_info(h)[["rss"]], error = function(e) NA_real_)
      if (!is.na(m)) { n <- n + 1L; if (m > peak) peak <- m }
    }
    if (!p$is_alive()) break
    Sys.sleep(0.02)
  }
  total <- proc.time()[["elapsed"]] - t_launch
  out <- tryCatch(p$get_result(), error = function(e) NULL)
  if (is.null(out)) {
    msg <- paste(utils::tail(p$read_all_error_lines(), 4), collapse = " | ")
    message("  FAILED: ", substr(msg, 1, 200))
    return(list(row = data.frame(method = method, total_s = NA, startup_s = NA,
      lib_s = NA, load_s = NA, compute_s = NA, peak_gb = NA, n_samples = NA,
      stringsAsFactors = FALSE), value = NULL))
  }
  startup <- max(0, total - out$lib_s - out$load_s - out$compute_s)
  list(row = data.frame(method = method, total_s = total, startup_s = startup,
         lib_s = out$lib_s, load_s = out$load_s, compute_s = out$compute_s,
         peak_gb = peak / GB, n_samples = n, stringsAsFactors = FALSE),
       value = out$value)
}

METHODS <- c("DuckDB (Parquet)", "duckplyr", "polars (streaming scan)",
             "tidypolars_skip", "arrow (Acero)", "polars (eager)",
             "collapse", "dplyr", "data.table", "base R (aggregate)")
METHODS <- setdiff(METHODS, "tidypolars_skip")

## Warm the page cache once so every method faces the same conditions.
invisible(readBin(PQ, "raw", n = file.size(PQ))); gc(full = TRUE)

## Whichever method runs first in a session pays to page in shared libraries that
## every later method then finds resident. Rather than quote that cost against one
## arbitrary method, run the suite once per method with the order rotated, so each
## takes a turn leading, and average.
rows <- list(); ref <- NULL
for (pass in seq_along(METHODS)) {
  order_this_pass <- METHODS[((seq_along(METHODS) + pass - 2L) %% length(METHODS)) + 1L]
  message("")
  message("#### pass ", pass, "/", length(METHODS),
          " leading with ", order_this_pass[1])
  for (m in order_this_pass) {
    r <- run_one(m)
    if (!is.null(r$value)) {
      v <- r$value
      v$n_snvs <- as.numeric(v$n_snvs); v$n_clonal <- as.numeric(v$n_clonal)
      v$sMF <- as.numeric(v$sMF)
      rownames(v) <- NULL
      if (is.null(ref)) { ref <- v; r$row$correct <- "reference" }
      else r$row$correct <- if (isTRUE(all.equal(ref, v, tolerance = 1e-8))) "match" else "MISMATCH"
    } else r$row$correct <- "error"
    r$row$pass <- pass
    r$row$position <- match(m, order_this_pass)
    message(sprintf("  %-24s pos %2d  total %6.2fs = startup %.2f + libs %.2f + load %.2f + compute %.2f | peak %5.2f GB [%s]",
                    m, r$row$position, r$row$total_s, r$row$startup_s, r$row$lib_s,
                    r$row$load_s, r$row$compute_s, r$row$peak_gb, r$row$correct))
    rows[[length(rows) + 1L]] <- r$row
    gc(full = TRUE)
  }
}

runs <- do.call(rbind, rows)
write.csv(runs, sub("[.]csv$", "_runs.csv", OUT), row.names = FALSE)

num <- c("total_s", "startup_s", "lib_s", "load_s", "compute_s", "peak_gb")
agg <- do.call(rbind, lapply(split(runs, runs$method), function(d) {
  data.frame(method = d$method[1],
             as.list(setNames(round(colMeans(d[num], na.rm = TRUE), 3), num)),
             total_sd = round(stats::sd(d$total_s, na.rm = TRUE), 3),
             startup_first = round(mean(d$startup_s[d$position == 1]), 3),
             startup_rest  = round(mean(d$startup_s[d$position != 1]), 3),
             n = nrow(d),
             correct = if (all(d$correct %in% c("match", "reference"))) "match" else "MISMATCH",
             stringsAsFactors = FALSE)
}))
agg <- agg[order(agg$total_s), ]
write.csv(agg, OUT, row.names = FALSE)
cat("
== mean of", length(METHODS), "passes, each method leading once ==
")
print(agg, row.names = FALSE, digits = 3)