# Check that two SummarizedExperiment objects can be combined
#
# Checks many matching criteria, basically pheno data needs to have similar
# special columns,
# the number of samples needs to be the same and feature data need to have the
# same columns names. Throws an error if any of the criteria is not fulfilled.
#
# @param x,y SummarizedExperiment objects
.check_match <- function(x, y) {
  # Lots of checks to ensure that everything goes smoothly

  # Amount of samples must be equal
  if (nrow(colData(x)) != nrow(colData(y))) {
    warning("Unequal amount of samples")
  }
  # Resulting feature ID must be unique
  feature_id <- c(rowData(x)$Feature_ID, rowData(y)$Feature_ID)
  if (!.all_unique(feature_id)) {
    stop("Merge would result in duplicated feature ID")
  }
  common <- intersect(colnames(x), colnames(y))
  if (!identical(colData(x)[common, "Injection_order"], 
                 colData(y)[common, "Injection_order"])) {
    stop("Injection orders of common samples are not identical")
  }
  if (!identical(colData(x)$Sample_ID, colData(y)$Sample_ID)) {
    warning("Sample IDs are not identical")
    samples_x <- setdiff(colnames(x), colnames(y))
    samples_y <- setdiff(colnames(y), colnames(x))
    log_text("Merging objects with unequal amounts of samples.")
    log_text("Samples only in first object:")
    log_text(paste0(paste(samples_x, collapse = ", "), "\n"))
    log_text("Samples only in second object:")
    log_text(paste0(paste(samples_y, collapse = ", "), "\n"))
  }


  overlap_cols <- intersect(colnames(colData(x)), colnames(colData(y))) |>
    setdiff(c("Sample_ID", "Injection_order"))

  if (length(overlap_cols)) {
    for (overlap_col in overlap_cols) {
      if (!identical(colData(x)[common, overlap_col], 
                     colData(y)[common, overlap_col])) {
        stop("Columns named ", overlap_col, 
             " in pheno data have different content")
      }
    }
  }

  if (!identical(colnames(rowData(x)), colnames(rowData(y)))) {
    stop("Feature data have different column names")
  }
}

# Merge two SummarizedExperiment objects together
.merge_mode_helper <- function(x, y) {
  # Create dummy injection order if original ones differ
  common <- intersect(colnames(x), colnames(y))
  if (!identical(colData(x)[common, "Injection_order"], 
      colData(y)[common, "Injection_order"])) {
    log_text(paste0("Injection order differs between modes.",
                    "Creating dummy injection order"))
    x_modes <- unique(rowData(x)$Split)
    # Save original injection order for first mode
    if (length(x_modes) == 1) {
      colData(x)[, paste0(x_modes, "_Injection_order")] <- x$Injection_order
    }
    dummy_injection <- as.numeric(-seq_along(colData(x)$Sample_ID))
    names(dummy_injection) <- x$Sample_ID
    colData(x)$Injection_order <- dummy_injection
    # Save original injection order in other modes
    colData(y)[, paste0(unique(rowData(y)$Split), "_Injection_order")] <-
      y$Injection_order
    # Update dummy injection
    y_in_x <- y$Sample_ID %in% x$Sample_ID
    new_io <- seq(from = min(dummy_injection) - 1, length.out = sum(!y_in_x))
    names(new_io) <- y$Sample_ID[!y_in_x]
    dummy_injection <- append(dummy_injection, new_io)

    colData(y)$Injection_order <- dummy_injection[match(colData(y)$Sample_ID,
                                                        names(dummy_injection))]
    log_text("Dummy injection order (row numbers) created")
  }
  # Check that the match is ok
  .check_match(x, y)

  merged_coldata <- dplyr::full_join(as.data.frame(colData(x)), 
                                     as.data.frame(colData(y)),
                                     by = intersect(colnames(colData(x)),
                                                    colnames(colData(y)))) |>
    S4Vectors::DataFrame()
  rownames(merged_coldata) <- merged_coldata$Sample_ID
  merged_rowdata <- rbind(rowData(x), rowData(y)) |>
    S4Vectors::DataFrame()
  if (identical(colnames(assay(x)), colnames(assay(y)))) {
    merged_assay <- rbind(assay(x), assay(y))
  } else {
    merged_assay <- dplyr::bind_rows(as.data.frame(assay(x)),
                                     as.data.frame(assay(y))) |> 
      as.matrix()
    rownames(merged_assay) <- rownames(merged_rowdata)
  }

  merged_object <- SummarizedExperiment(assays = merged_assay, 
                                        colData = merged_coldata,
                                        rowData = merged_rowdata)

  merged_object
}

# Convert objects in ... to a list
# OR if a list is given in the first place, preserve that list
.to_list <- function(...) {
  # Combine the objects to a list
  objects <- list(...)
  # If a list is given in the first place, it should move to top level
  if (length(objects) == 1) {
    if (is(objects[[1]], "list")) {
      objects <- objects[[1]]
    }
  }

  objects
}

#' Merge SummarizedExperiment objects together
#'
#' Merges two or more SummarizedExperiment objects together. Can be used to
#' merge analytical modes or batches.
#'
#' @param ... \code{
#' \link[SummarizedExperiment:SummarizedExperiment-class]{SummarizedExperiment}}
#' objects or a list of objects
#' @param merge what to merge? features is used for combining analytical modes,
#' samples is used for batches
#' @param assay.type character, assay to be used in case of multiple assays.
#' The same assay needs to be present in all objects to be merged, and the
#' resultant object contains this single assay.
#'
#' @return A merged SummarizedExperiment object.
#'
#' @details When merging samples, sample IDs that begin with "QC" or "Ref" are
#' combined so that they have running numbers on them. This means that if both
#' batches have samples called "QC_1", this will not result in an error,
#' but the sample IDs will be adjusted so that they are unique. Column names in
#' the feature data that are shared between batches but have different content
#' are renamed by adding a suffix to avoid data loss. The suffix is the index
#' of the batch in the input list. The most common flag value across batches
#' is included in the "Flag" column in the merged object.
#'
#' @examples
#' # Merge analytical modes
#' data(hilic_neg_sample, hilic_pos_sample, rp_neg_sample, rp_pos_sample)
#' merged <- merge_notame_sets(
#'   hilic_neg_sample, hilic_pos_sample,
#'   rp_neg_sample, rp_pos_sample
#' )
#' # Merge batches
#' batch1 <- toy_notame_set[, toy_notame_set$Batch == 1]
#' batch2 <- toy_notame_set[, toy_notame_set$Batch == 2]
#' merged <- merge_notame_sets(batch1, batch2, merge = "samples")
#'
#' @export
merge_notame_sets <- function(
  ...,
  merge = c("features", "samples"),
  assay.type = NULL
) {
  merge <- match.arg(merge)
  # Combine the objects to a list
  objects <- .to_list(...)
  # Class check
  if (!all(vapply(objects, class, character(1)) == "SummarizedExperiment")) {
    stop("The arguments should only contain SummarizedExperiment objects")
  }
  # Check assay.type and prepare objects for merge
  from_list <- lapply(objects, .get_from_name, assay.type)
  if (!length(unique(from_list)) == 1) {
    stop(
      "The same assay must be present in all objects or alternatively,
          use objects with a single assay"
    )
  } else {
    from <- unlist(unique(from_list))
  }
  objects <- lapply(objects, function(object) {
    assays(object) <- assays(object)[from]
    object
  })

  # Choose merging function
  if (merge == "features") {
    merged <- Reduce(.merge_mode_helper, objects)
  } else {
    merged <- .merge_batch_helper(objects)
    merged <- .recreate_flag(merged)
  }
  merged
}

.rowdata_batch_helper <- function(list_of_dfs) {
  # Find all common columns
  list_of_dfs <- lapply(list_of_dfs, as.data.frame)
  all_cols <- lapply(list_of_dfs, colnames)
  common_cols <- Reduce(intersect, all_cols)

  # Find columns with same name but different content
  # If not all contents are identical, mark for suffixing
  nonmatching_cols_idx <- vapply(
    common_cols,
    function(col) {
      col_list <- lapply(list_of_dfs, `[[`, col)
      !all(vapply(col_list[-1], identical, logical(1), col_list[[1]]))
    },
    logical(1)
  )

  # Add suffix to nonmatching columns in each dataframe
  nonmatching_cols <- common_cols[nonmatching_cols_idx]
  list_of_dfs <- mapply(
    function(df, idx) {
      cols_to_rename <- colnames(df) %in% nonmatching_cols
      colnames(df)[cols_to_rename] <- paste0(
        colnames(df)[cols_to_rename],
        "_",
        idx
      )
      df
    },
    list_of_dfs,
    seq_along(list_of_dfs),
    SIMPLIFY = FALSE
  )

  # Merge all dataframes
  merged <- Reduce(
    function(x, y) {
      dplyr::full_join(x, y, by = common_cols[!nonmatching_cols_idx])
    },
    list_of_dfs
  )
  merged <- S4Vectors::DataFrame(merged)
  merged
}

# Helper function to recreate the Flag column after merging batches, since the
# original Flag columns are renamed to avoid data loss, we need to combine them
# back into a single Flag column that contains all the information.
.recreate_flag <- function(se) {
  flags <- grep("Flag", colnames(rowData(se)))
  # If we have less than 2 flag columns, we can just return the original object
  # since the merge of Flag columns was successful
  if (length(flags) < 2) {
    return(se)
  }
  # Compute the most common flag for each feature across batches
  # and use that as the final flag for the merged object
  flag_mat <- rowData(se)[, flags, drop = FALSE]
  unique_flags <- c(sort(unique(as.vector(as.matrix(flag_mat)))))
  counts <- apply(flag_mat, 1, function(row) {
    tabulate(factor(row, levels = unique_flags), nbins = length(unique_flags))
  }, simplify = FALSE)
  most_common <- unname(vapply(counts, function(count) {
    if (all(count == 0)) {
      return(NA_character_)
    }
    unique_flags[which.max(count)]
  }, character(1)))
  rowData(se)$Flag <- most_common
  se
}

.check_batch_coldata <- function(cd) {
  if (anyDuplicated(cd$Sample_ID)) {
    log_text(paste0(
      "Found duplicated sample IDs when merging, ",
      "renaming QC and Ref samples"
    ))
    qc_idx <- grepl("QC", cd$Sample_ID)
    cd$Sample_ID[qc_idx] <- paste0(
      "QC_",
      seq_len(sum(qc_idx))
    )
    ref_idx <- grepl("Ref", cd$Sample_ID)
    cd$Sample_ID[ref_idx] <- paste0(
      "Ref_",
      seq_len(sum(ref_idx))
    )
    rownames(cd) <- cd$Sample_ID
  }
  cd
}

.merge_batch_helper <- function(objects) {
  # Handle column data
  cds <- lapply(objects, colData)
  merged_coldata <- do.call(rbind, cds)
  merged_coldata <- .check_batch_coldata(merged_coldata)

  # Handle assay data, bind_rows to preserve all features
  assays <- lapply(objects, assay)
  assays <- lapply(assays, \(df) as.data.frame(t(df)))
  merged_assay <- t(as.matrix(dplyr::bind_rows(assays)))
  colnames(merged_assay) <- rownames(merged_coldata)

  # Handle row data
  rds <- lapply(objects, rowData)
  merged_rowdata <- .rowdata_batch_helper(rds)

  merged_object <- SummarizedExperiment(
    assays = merged_assay,
    colData = merged_coldata,
    rowData = merged_rowdata
  )

  merged_object
}
