#' @inheritParams MSstatsPreprocessBig
#' @keywords internal
MSstatsPreprocessBigSparklyr <- function(connection, input_file, output_file_name,
                                         max_feature_count = 20,
                                         filter_unique_peptides = FALSE,
                                         aggregate_psms = FALSE,
                                         filter_few_obs = FALSE,
                                         remove_annotation = FALSE) {
  tbl_name <- getTblName(remove_annotation)
  
  sparklyr::spark_read_csv(connection, "mstinput", path = input_file,
                           header = TRUE, delimiter = ",",
                           memory = FALSE, overwrite = TRUE,
                           repartition = 100)
  
  if (remove_annotation) {
    dplyr::tbl(connection, "mstinput") %>%
      dplyr::select(-Condition, -BioReplicate) %>%
      sparklyr::spark_write_table("mstinputnoannot", mode = "overwrite")
    DBI::dbRemoveTable(connection, "mstinput")
  }
  
  dplyr::tbl(connection, "mstinputnoannot") %>%
    dplyr::group_by(PeptideSequence) %>%
    dplyr::summarize(NumProteins = n_distinct(ProteinName)) %>%
    dplyr::filter(NumProteins > 1) %>%
    sparklyr::spark_write_table("sharedpeptides", mode = "overwrite")
  sparklyr::sdf_repartition(tbl(connection, "mstinputnoannot"), partition_by = "ProteinName")
  
  dplyr::tbl(connection, "mstinputnoannot") %>%
    dplyr::group_by(ProteinName, PeptideSequence,
                    PrecursorCharge, FragmentIon, ProductCharge, IsotopeLabelType,
                    Run) %>%
    dplyr::summarize(Intensity = max(Intensity)) %>%
    sparklyr::spark_write_table("mstinputagg", mode = "overwrite")
  DBI::dbRemoveTable(connection, "mstinputnoannot")
  
  dplyr::tbl(connection, "mstinputagg") %>%
    dplyr::group_by(ProteinName, PeptideSequence, PrecursorCharge, FragmentIon, ProductCharge, IsotopeLabelType) %>%
    dplyr::summarize(NumObs = sum(as.numeric(!is.na(Intensity) & Intensity > 1))) %>%
    dplyr::filter(NumObs <= 2) %>%
    dplyr::select(-NumObs) %>%
    sparklyr::spark_write_table(name = "less2obs", mode = "overwrite")
  
  dplyr::anti_join(
    dplyr::tbl(connection, "mstinputagg"),
    tbl(connection, "less2obs"),
    by <- c("ProteinName",
            "PeptideSequence", "PrecursorCharge", "FragmentIon",
            "ProductCharge", "IsotopeLabelType")
  ) %>%
    sparklyr::spark_write_table(name = "mstinputfilt", mode = "overwrite")
  DBI::dbRemoveTable(connection, "mstinputagg")
  DBI::dbRemoveTable(connection, "less2obs")
  
  dplyr::tbl(connection, "mstinputfilt") %>%
    dplyr::group_by(ProteinName, PeptideSequence, PrecursorCharge, FragmentIon, ProductCharge, IsotopeLabelType) %>%
    dplyr::summarize(Intensity = mean(Intensity)) %>%
    dplyr::ungroup() %>%
    sparklyr::spark_write_table("fortop50", mode = "overwrite", partition_by = "ProteinName")
  
  
  dplyr::tbl(connection, "fortop50") %>%
    dplyr::group_by(ProteinName, IsotopeLabelType) %>%
    dplyr::mutate(Rank = min_rank(-Intensity)) %>%
    dplyr::filter(Rank <= max_feature_count) %>%
    dplyr::select(-Rank, -Intensity) %>%
    sparklyr::spark_write_table("top50", mode = "overwrite", partition_by = "ProteinName")
  
  dplyr::inner_join(
    dplyr::tbl(connection, "mstinputfilt"),
    dplyr::tbl(connection, "top50"),
    by <- c("ProteinName", "IsotopeLabelType", "PeptideSequence", "PrecursorCharge",
            "FragmentIon", "ProductCharge")) %>%
    sparklyr::spark_write_table("mstprocessed", mode = "overwrite")
  DBI::dbRemoveTable(connection, "top50")
  DBI::dbRemoveTable(connection, "fortop50")
  DBI::dbRemoveTable(connection, "mstinputfilt")
  
  sparklyr::spark_write_csv(sparklyr::sdf_repartition(dplyr::tbl(connection, "mstprocessed"), 1),
                            output_file_name, mode = "overwrite")
  dplyr::tbl(connection, "mstprocessed")
}


#' @inheritParams MSstatsPreprocessBig
#' @keywords internal
MSstatsPreprocessBigArrow <- function(input_file,
                                      output_file_name,
                                      max_feature_count = 20,
                                      filter_unique_peptides = FALSE,
                                      aggregate_psms = FALSE,
                                      filter_few_obs = FALSE,
                                      calculateAnomalyScores = FALSE, 
                                      anomalyModelFeatures = c()) {
  input <- arrow::open_dataset(input_file, format = "csv")
  
  # dynamically cast any columns Arrow inferred as 'null'
  null_columns <- names(input$schema)[vapply(input$schema$fields, function(f) f$type$ToString() == "null", logical(1))]
  if (length(null_columns) > 0) {
    for (col in null_columns) {
      input <- dplyr::mutate(input, !!rlang::sym(col) := as.numeric(!!rlang::sym(col)))
    }
  }

  if ("PrecursorMz" %in% input$schema$names) {
    input <- dplyr::mutate(input, PrecursorMz = as.numeric(PrecursorMz))
  }
  
  input <- dplyr::mutate(input,
                         Feature = paste(PeptideSequence, PrecursorCharge,
                                         FragmentIon, ProductCharge, sep = "_"))
  feature_counts <- dplyr::group_by(input, ProteinName, Feature)
  feature_counts <- dplyr::summarize(feature_counts,
                                     MeanAbundance = mean(Intensity,
                                                          na.rm = TRUE))
  feature_counts <- dplyr::collect(feature_counts)
  
  feature_counts <- dplyr::mutate(
    feature_counts,
    feature_rank = dplyr::min_rank(dplyr::desc(MeanAbundance)))
  
  feature_counts <- dplyr::filter(feature_counts,
                                  feature_rank <= max_feature_count)
  
  feature_counts <- dplyr::select(feature_counts, -MeanAbundance, -feature_rank)
  input <- dplyr::inner_join(input, feature_counts, 
                             by = c("ProteinName", "Feature"))
  input <- dplyr::select(input, -Feature)
  
  arrow::write_dataset(input, .prefixedPath("topN_", output_file_name), format = "csv")
  
  if (filter_unique_peptides) {
    pp_df <- dplyr::select(input, ProteinName, PeptideSequence)
    pp_df <- dplyr::group_by(pp_df, PeptideSequence)
    pp_df <- dplyr::summarize(pp_df, NumProteins = dplyr::n_distinct(ProteinName))
    pp_df <- dplyr::filter(pp_df, NumProteins > 1)
    pp_df <- dplyr::select(pp_df, -NumProteins)
    input <- dplyr::anti_join(input, pp_df, by = "PeptideSequence")
  }
  if (aggregate_psms) {
    group_cols <- c("ProteinName", "PeptideSequence", "PrecursorCharge",
                    "FragmentIon", "ProductCharge", "IsotopeLabelType", "Run",
                    "Condition", "BioReplicate")
    
    max_per_group <- input %>%
      group_by(across(all_of(group_cols))) %>%
      summarise(max_intensity = max(Intensity, na.rm = TRUE), 
                .groups = "drop")
    
    input <- input %>%
      inner_join(max_per_group, by = group_cols) %>%
      filter(Intensity == max_intensity) %>%
      select(-max_intensity)
  }
  
  if (filter_few_obs) {
    input <- dplyr::mutate(input,
                             Feature = paste(PeptideSequence, PrecursorCharge,
                                             FragmentIon, ProductCharge, sep = "_"))
    input <- dplyr::group_by(input, ProteinName, Feature)
    observation_counts <- dplyr::summarize(input,
                                           NumObs = sum(!is.na(Intensity) &
                                                          Intensity > 0))
    observation_counts <- dplyr::filter(observation_counts, NumObs <= 2)
    observation_counts <- dplyr::select(observation_counts, -NumObs)
    input <- dplyr::anti_join(input, observation_counts,
                              by = c("ProteinName", "Feature"))
    input <- dplyr::select(input, -Feature)
  }
  
  arrow::write_dataset(input, output_file_name, format = "csv")
  input
}

getTblName <- function(remove_annotation) {
  if (remove_annotation) {
    "mstinput"
  } else {
    "mstinputnoannot"
  }
}
