#' check_patterns
#' 
#' Function to check if the TRUE FALSE patterns have at least two rows that are
#' opposite, we need this to be able to generate repressor regulators
#'
#' @param patterns tibble of TRUE FALSE values
#' @export
#' @return list of indices where the rows are opposite
#'
#' @examples
#' patterns <- tibble::tibble(one = c(TRUE, FALSE, TRUE, FALSE), 
#'               two = c(TRUE, TRUE, TRUE, TRUE), 
#'               three = c(FALSE, TRUE, FALSE, TRUE), 
#'               four = c(FALSE, TRUE, TRUE, TRUE))
#' opposite_indices <- check_patterns(patterns)
#' 
check_patterns <- function(patterns) {
  opposite_indices <- list()
  for (i in 1:(nrow(patterns) - 1)) {
    for (j in (i + 1):nrow(patterns)) {
      if (all(patterns[i, ] == !patterns[j, ])) {
        opposite_indices <- append(opposite_indices, list(c(i, j)))
      }
    }
  }
  return(opposite_indices)
}

#' make_cluster_patterns
#' 
#' Function to make the tibble with cluster combinations for the gene expression
#' patterns along the cells
#' This function is a slightly modified copy of the `make_cluster_patterns` function
#' from the `Acorde` package (v1.0.0), originally developed by Arzalluz-Luque A, 
#' Salguero P, Tarazona S, Conesa A. (2022). acorde unravels functionally 
#' interpretable networks of isoform co-usage from single cell data. Nature communications
#' 1828. DOI: 10.1038/s41467-022-29497-w. The original package is licensed under the 
#' GPL-3 license.
#'
#' @param numcells Number of different celltypes we are simulating
#' @param clusters OPTIONAL. Number of co-expression patterns the user wants
#'      to simulate
#'
#' @return A tibble with number of columns equal to number of celltypes, rows
#'  according to the number of TRUE/FALSE combinations corresponding to the
#'  gene expression patterns along the cells
#' @export
#' @examples
#' patterns <- make_cluster_patterns(numcells = 4, clusters = 8)
#' cell_types <- list('Treg' = c(1:10),'cDC' = c(11:20),'CD4_TEM' = c(21:30),
#' 'Memory_B' = c(31:40))
#' patterns <- make_cluster_patterns(numcells = length(cell_types), 
#'     clusters = 8)
#' 
make_cluster_patterns <- function(numcells = 4, clusters = 8) {
  patterns <- tibble::tibble()
  col_names <- paste0("Var", 1:numcells)
  
  # Loop until we find a suitable pattern matrix
  repeat {
    # Generate all possible patterns
    patterns <- tibble::tibble()
    for (i in 1:(2^numcells)) {
      # Make it a bigger probability to have FALSE than true
      # To be conservative with the expression values.
      boolArr <- runif(numcells) > 0.85
      patterns <- rbind(patterns, boolArr)
    }
    colnames(patterns) <- col_names
    
    # Sample the required number of clusters
    patterns_ret <- slice_sample(patterns, n = clusters)
    
    # Check for duplicate rows
    if (any(duplicated(patterns_ret))) {
      next
    }
    
    # Check for opposite patterns
    opposite_indices <- check_patterns(patterns_ret)
    if (length(opposite_indices) >= 1) {
      break
    }
  }
  
  return(list("patterns" = patterns_ret, "opposite_indices" = opposite_indices))
}

#' simulate coexpression
#' 
#' Adapted from ACORDE (https://github.com/ConesaLab/acorde) to adapt to our
#' data input type. Simulates coexpression of genes along celltypes
#'
#' This function is a slightly modified copy of the `simulate_coexpression` function
#' from the `Acorde` package (v1.0.0), originally developed by Arzalluz-Luque A, 
#' Salguero P, Tarazona S, Conesa A. (2022). acorde unravels functionally 
#' interpretable networks of isoform co-usage from single cell data. Nature communications
#' 1828. DOI: 10.1038/s41467-022-29497-w. The original package is licensed under the 
#' GPL-3 license.
#' @param sim_matrix Matrix with rows as features and columns as cells
#' @param feature_no Total number of features to be distributed between the 
#'    coexpression clusters
#' @param cellTypes list where the i-th element of the list contains the column 
#'    indices for i-th experimental conditions. List must be a named list.
#' @param patterns Tibble with TRUE FALSE depicting the cluster patterns to
#'    simulate. Generated by the user or by \code{make_cluster_patterns}.
#' @param cluster_size OPTIONAL. It may be inputted by the user. By default, 
#'    its the number of features divided by the number of patterns to generate.
#'
#' @return the simulated coexpression
#' @export

simulate_coexpression <- function(sim_matrix,
                                  feature_no, cellTypes,
                                  patterns,
                                  cluster_size = NULL){
  
  ## If no cluster size specified, make it by default
  if (is.null(cluster_size)){
    cluster_size <- feature_no/dim(patterns)[1]
  }
  
  # Define notin to filter later on
  `%notin%` <- Negate(`%in%`)
  
  # extract counts 
  normcounts <- as.data.frame(sim_matrix)
  # get cell ids in each cell type
  colData <- tibble::tibble(Cell = colnames(normcounts),
                            Group = rep(names(cellTypes), times = lengths(cellTypes)))
  
  group.list <- colData$Cell %>% split(colData$Group)
  # extract cell type (group) expr matrices
  normcounts.list <- purrr::map(group.list,
                                ~(normcounts[, as.character(.)] %>%
                                    tibble::rownames_to_column("feature")))
  
  # rank features by mean expression in each cell type
  normcounts.list <- purrr::map(normcounts.list,
                                ~dplyr::mutate(., mean = rowMeans(.[,-1])) %>%
                                  dplyr::arrange(dplyr::desc(mean))
                                %>% dplyr::select(-mean))
  
  # select top and bottom feature IDs for each cell type
  
  # top
  top_features.list <- purrr::map(normcounts.list,
                                  ~dplyr::select(., feature) %>%
                                    dplyr::rename(top = "feature") %>%
                                    dplyr::slice(., 1:feature_no) %>%
                                    tibble::as_tibble())
  
  # modify bottom feature no. to create range correctly
  feature_no.c <- nrow(normcounts) - (feature_no - 1)
  # bottom
  bottom_features.list <- purrr::map(normcounts.list,
                                     ~dplyr::select(., feature) %>%
                                       dplyr::rename(bottom = "feature") %>%
                                       dplyr::slice(., feature_no.c:nrow(normcounts)) %>%
                                       tibble::as_tibble())
  
  # create a two-col tibble with top/bottom features per group
  features.list <- purrr::map2(top_features.list, bottom_features.list,
                               dplyr::bind_cols)
  
  
  ## USE SUPPLIED PATTENRS TO SHUFFLE THE CELL TYPE MATRICES ##
  
  # match column names for patterns
  colnames(patterns) <- names(normcounts.list)
  
  # shuffle matrix for each cell type following cluster patterns
  # note that internal function shuffle_group_matrix() is used to perform
  # each individual shuffling operation
  
  # calculate no. of clusters based on size and number of features
  partitions <- feature_no / cluster_size
  partitions <- trunc(partitions)
  
  expr.list <- purrr::pmap(list(normcounts.list, features.list, patterns),
                           ~MOSim::shuffle_group_matrix(sim_data = ..1,
                                                        feature_ids = ..2,
                                                        group_pattern = ..3,
                                                        ngroups = partitions))
  
  feature_names <- expr.list[[1]]$feature
  
  # join cell type matrices into a single expression matrix
  expr.list <- purrr::map(expr.list, dplyr::select, -feature)
  coexpr.df <- dplyr::bind_cols(expr.list) %>% tibble::as_tibble()
  
  # merge all dataframes in the list by feature column, fillna the genes not 
  coexpr.df <- coexpr.df %>%
    #dplyr::mutate(feature = paste0("Feature", seq(1, nrow(coexpr.df)))) %>%
    dplyr::mutate(feature = feature_names) %>%
    dplyr::relocate(feature)
  
  # generate feature ID vectors for co-expression clusters
  clusters <- split(coexpr.df$feature,
                    cut(seq(1, nrow(coexpr.df)),
                        breaks = nrow(patterns), labels = FALSE))
  
  # New, not in acorde's simulate coexpression
  # Get the genes that we havent used into the table
  # Format the simulated dataframe as tibble with feature column and get the 
  # missing rows into the coexpression DF
  normcounts <- dplyr::mutate_all(normcounts, function(x) as.double(x))
  normcounts$feature <- rownames(normcounts)
  normcounts <- tibble::as_tibble(normcounts)
  # Sort the columnames according to the simulated dataframe
  normcounts <- normcounts[, colnames(coexpr.df)]
  
  # remove from the original simulated matrix the features co-regulated
  normcounts <- normcounts %>%
    dplyr::filter(feature %notin% coexpr.df$feature)
  
  ## Bring back other features not corregulated
  coexpr.df <- dplyr::bind_rows(coexpr.df, normcounts) %>% replace(is.na(.), 0)
  
  # build a list with results
  coexpr_sim <- list(sim_matrix = coexpr.df,
                     sim_clusters = clusters)
  return(coexpr_sim)
  
}


#' shuffle_group_matrix, Reorder cell type-specific expression matrix during 
#' co-expression simulation. Copied from ACORDE (https://github.com/ConesaLab/acorde)
#' to facilitate stability and running within our scripts
#' This function is a slightly modified copy of the `shuffle_group_matrix` function
#' from the `Acorde` package (v1.0.0), originally developed by Arzalluz-Luque A, 
#' Salguero P, Tarazona S, Conesa A. (2022). acorde unravels functionally 
#' interpretable networks of isoform co-usage from single cell data. Nature communications
#' 1828. DOI: 10.1038/s41467-022-29497-w. The original package is licensed under the 
#' GPL-3 license.
#' 
#' @description This function is used internally by \code{acorde} to perform
#' the shuffling of simulated features for an individual cell type, as part of
#' the co-expression simulation process. The function is called recursively by
#' \code{\link[MOSim:simulate_coexpression]{simulate_coexpression()}} to
#' perform the simulation on a full scRNA-seq matrix.
#'
#' @param sim_data A count matrix with features as rows and cells as columns.
#' Feature IDs must be included in an additional column named \code{feature}.
#' @param feature_ids A two-column \code{tibble} containing \code{top} and \code{bottom}
#' columns, each including the feature IDs of features to be used as highly or
#' lowly expressed when shuffling by the indicated expression pattern.
#' @param group_pattern A logical vector, containing \code{TRUE} to indicate that
#' high expression in that cell type is desired and \code{FALSE} if the opposite.
#' The vector must be ordered as the cell types in \code{sim_data}.
#' @param ngroups An integer indicating the number of groups that top and bottom
#' features should be divided into. It is computed by dividing the number
#' of features selected as highly/lowly expressed by the size of the clusters
#' that are to be generated.
#' @export
#' @return An expression matrix, with the same characteristics as \code{sim_data},
#' and a number of features defined as the total amount of top/bottom features
#' selected divided by the number of clusters for which co-expression patterns
#' where supplied.

shuffle_group_matrix <- function(sim_data, feature_ids, group_pattern, ngroups) {
  
  top <- feature_ids$top
  bottom <- feature_ids$bottom
  
  top <- sample(top)
  top.list <- split(top, cut(seq_along(top), breaks = ngroups, labels = FALSE, include.lowest = TRUE))
  
  bottom <- sample(bottom)
  bottom.list <- split(bottom, cut(seq_along(bottom), breaks = ngroups, labels = FALSE, include.lowest = TRUE))
  
  features_bound <- vector("list", length(group_pattern))
  features_bound[group_pattern] <- top.list[1:sum(group_pattern)]
  features_bound[!group_pattern] <- bottom.list[1:sum(!group_pattern)]
  features_bound <- unlist(features_bound)
  
  sim_data.mod <- sim_data[sim_data$feature %in% features_bound, ]
  sim_data.mod <- sim_data.mod[match(features_bound, sim_data.mod$feature), , drop = FALSE]
  
  return(sim_data.mod)
}


#' make_association_dataframe
#' 
#' This function generates a dataframe containing the information of the 
#' relationship between ATAC and RNA, based on the cluster groups, and then
#' tells the order the genes and peaks should be in the simulated dataframe
#' of the group
#'
#' @param group Group from which we are generating the association dataframe
#' @param genereggroup list of elements to generate the association dataframe
#'      such as clusters of each omic, indices of opposite clusters, which
#'      genes are activated, repressed, behavior of the features etc.
#' @param numtotalgenes total number of genes
#' @param numtotalpeaks total number of peaks
#' @param minFC FC below which is downregulated
#' @param maxFC FC above which is upregulated
#' @return a dataframe with all the information the user needs about each gene
#'      and the order of gene and peak names to rename them in the simulated
#'      datasets of the group
#' @export
#'
make_association_dataframe <- function(group, genereggroup, numtotalgenes, 
                                       numtotalpeaks, minFC, maxFC){
  # Start from the association list, now we have two columns Peak_ID and Gene_ID
  
  columns <- c("Gene_ID", "Peak_ID", "RegulatorEffect", "Gene_cluster", "Peak_cluster", 
               "Gene_DE", "Peak_DE", "Gene_FC", "Peak_FC")
  keep_remaining <- function(ori_clus, clus){
    clusters_a <- as.data.frame(lengths(ori_clus))
    clusters_a$cluster <- rownames(clusters_a)
    colnames(clusters_a) <- c("Freq", "cluster")
    clusters_ao <- as.data.frame(table(clus))
    colnames(clusters_ao) <- c("cluster", "Freq")
    clusters_a <- clusters_a[colnames(clusters_ao)]
    result <- clusters_a %>%
      dplyr::left_join(clusters_ao, by = "cluster", suffix = c(".a", ".ao")) %>%
      dplyr::mutate(Freq = Freq.a - ifelse(is.na(Freq.ao), 0, Freq.ao)) %>%
      dplyr::select(cluster, Freq)
    return(result)
  }
  
  keep_remaining2 <- function(ori_clus, clus){
    clusters_a <- as.data.frame(ori_clus)
    clusters_ao <- as.data.frame(table(clus))
    colnames(clusters_ao) <- c("cluster", "Freq")
    clusters_a <- clusters_a[colnames(clusters_ao)]
    result <- clusters_a %>%
      dplyr::left_join(clusters_ao, by = "cluster", suffix = c(".a", ".ao")) %>%
      dplyr::mutate(Freq = Freq.a - ifelse(is.na(Freq.ao), 0, Freq.ao)) %>%
      dplyr::select(cluster, Freq)
    
    return(result)
  }
  
  ## Make the repressed
  df2 <- data.frame(matrix(nrow = length(genereggroup[[paste0("GeneRepressed_G", 
                                                              group)]]$Gene_ID), ncol = length(columns)))
  colnames(df2) <- columns
  df2$Gene_ID <- genereggroup[[paste0("GeneRepressed_G", group)]]$Gene_ID
  df2$Peak_ID <- genereggroup[[paste0("GeneRepressed_G", group)]]$Peak_ID
  df2$RegulatorEffect <- rep("Repressor", length(df2[[1]]))
  Gene_cluster <- sample(unique(unlist(genereggroup$opposite_indices)), 
                          length(df2[[1]]), replace = TRUE)
  df2$Gene_cluster = Gene_cluster
  lookup <- unlist(lapply(genereggroup$opposite_indices, function(x) setNames(rev(x), x)))
  Peak_cluster <- as.numeric(sapply(Gene_cluster, function(x) lookup[as.character(x)]))
  df2$Peak_cluster <- Peak_cluster
  df2$Gene_DE <- rep("Down", length(df2[[1]]))
  df2$Peak_DE <- rep("Up", length(df2[[1]]))
  
  if (group > 1){
    df2$Gene_FC <- runif(n = length(df2[[1]]), min = 0.0001, max = minFC)
    df2$Peak_FC <- runif(n = length(df2[[1]]), min = maxFC, max = 100)
  }
  
  remaining_rna <- keep_remaining(genereggroup$`Clusters_scRNA-seq`, Gene_cluster)
  remaining_atac <- keep_remaining(genereggroup$`Clusters_scATAC-seq`, Peak_cluster)
  
  remaining_r <- as.numeric(rep(remaining_rna$cluster, remaining_rna$Freq))

  # Make the activated
  
  df1 <- data.frame(matrix(nrow = length(genereggroup[[paste0("GeneActivated_G", 
                                    group)]]$Gene_ID), ncol = length(columns)))
  colnames(df1) <- columns
  
  df1$Gene_ID <- genereggroup[[paste0("GeneActivated_G", group)]]$Gene_ID
  df1$Peak_ID <- genereggroup[[paste0("GeneActivated_G", group)]]$Peak_ID
  df1$RegulatorEffect <- rep("Activator", length(df1[[1]]))
  Gene_cluster <- sample(remaining_r, 
                         length(df1[[1]]), replace = FALSE)
  Peak_cluster <- Gene_cluster
  df1$Gene_cluster <- Gene_cluster
  df1$Peak_cluster <- Peak_cluster
  df1$Gene_DE <- rep("Up", length(df1[[1]]))
  df1$Peak_DE <- rep("Up", length(df1[[1]]))
  
  if (group > 1){
    df1$Gene_FC <- runif(n = length(df1[[1]]), min = maxFC, max = 100)
    df1$Peak_FC <- df1$Gene_FC
  }
  
  remaining_rna <- keep_remaining2(remaining_rna, Gene_cluster)
  remaining_atac <- keep_remaining2(remaining_atac, Peak_cluster)
  
  remaining_r <- as.numeric(rep(remaining_rna$cluster, remaining_rna$Freq))
  remaining_a <- as.numeric(rep(remaining_atac$cluster, remaining_atac$Freq))
  

  ## Add extras
  
  df3 <- data.frame(matrix(nrow = length(c(genereggroup[[paste0("GeneExtraUp_G", 
                    group)]], rep(NA, length(genereggroup[[paste0("FeatExtraUp_G", 
                    group)]])))), ncol = length(columns)))
  colnames(df3) <- columns
  df3$Gene_ID <- c(genereggroup[[paste0("GeneExtraUp_G", group)]], 
                      rep(NA, length(genereggroup[[paste0("FeatExtraUp_G", group)]])))
  df3$Peak_ID <- c(rep(NA, length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                      genereggroup[[paste0("FeatExtraUp_G", group)]])
  df3$RegulatorEffect <- rep("NE", length(df3[[1]]))
  df3$Gene_cluster <- c(rep(0, length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                      rep(NA, length(genereggroup[[paste0("FeatExtraUp_G", group)]])))
  df3$Peak_cluster <- c(rep(NA, length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                      rep(0, length(genereggroup[[paste0("FeatExtraUp_G", group)]])))
  df3$Gene_DE <- c(rep("Up", length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                      rep(NA, length(genereggroup[[paste0("FeatExtraUp_G", group)]])))
  df3$Peak_DE <- c(rep(NA, length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                      rep("Up", length(genereggroup[[paste0("FeatExtraUp_G", group)]])))

  if (group > 1){
    df3$Gene_FC <- c(runif(n = length(genereggroup[[paste0("GeneExtraUp_G", group)]]), min = maxFC, max = 100),
                        rep(NA, length(genereggroup[[paste0("FeatExtraUp_G", group)]])))
    df3$Peak_FC <- c(rep(NA, length(genereggroup[[paste0("GeneExtraUp_G", group)]])), 
                        runif(n = length(genereggroup[[paste0("FeatExtraUp_G", group)]]), min = maxFC, max = 100))
  }
  
  ## Extra down
  
  df4 <- data.frame(matrix(nrow = length(c(genereggroup[[paste0("GeneExtraDown_G", group)]], 
                                           rep(NA, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))), ncol = length(columns)))
  colnames(df4) <- columns
  df4$Gene_ID <- c(genereggroup[[paste0("GeneExtraDown_G", group)]], 
                      rep(NA, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
  df4$Peak_ID <- c(rep(NA, length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                      genereggroup[[paste0("FeatExtraDown_G", group)]])
  df4$RegulatorEffect <- rep("NE", length(df4[[1]]))
  df4$Gene_cluster <- c(rep(0, length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                      rep(NA, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
  df4$Peak_cluster <- c(rep(NA, length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                      rep(0, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
  df4$Gene_DE <- c(rep("Down", length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                      rep(NA, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
  df4$Peak_DE <- c(rep(NA, length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                      rep("Down", length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
  
  if (group > 1){
    df4$Gene_FC <- c(runif(n = length(genereggroup[[paste0("GeneExtraDown_G", group)]]), min = 0.0001, max = minFC),
                        rep(NA, length(genereggroup[[paste0("FeatExtraDown_G", group)]])))
    df4$Peak_FC <- c(rep(NA, length(genereggroup[[paste0("GeneExtraDown_G", group)]])), 
                        runif(n = length(genereggroup[[paste0("FeatExtraDown_G", group)]]), min = 0.0001, max = minFC))
  }
  
  ## All remaining
  
  df5 <- data.frame(matrix(nrow = length(c(genereggroup[[paste0("GeneRemaining_G", group)]], 
                                    rep(NA, length(genereggroup[[paste0("FeatRemaining_G", group)]])))), 
                           ncol = length(columns)))
  colnames(df5) <- columns
  df5$Gene_ID <- c(genereggroup[[paste0("GeneRemaining_G", group)]], 
                      rep(NA, length(genereggroup[[paste0("FeatRemaining_G", group)]])))
  df5$Peak_ID <- c(rep(NA, length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                      genereggroup[[paste0("FeatRemaining_G", group)]])
  df5$RegulatorEffect <- rep("NE", length(df5[[1]]))
  
  df5$Gene_cluster <- c(remaining_r, 
                      rep(0, length(genereggroup[[paste0("GeneRemaining_G", group)]]) - length(remaining_r)),
                      rep(NA, length(genereggroup[[paste0("FeatRemaining_G", group)]])))
  df5$Peak_cluster <- c(rep(NA, length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                      remaining_a, 
                      rep(0, length(genereggroup[[paste0("FeatRemaining_G", group)]]) - length(remaining_a)))
  
  df5$Gene_DE <- c(rep("NE", length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                      rep(NA, length(genereggroup[[paste0("FeatRemaining_G", group)]])))
  df5$Peak_DE <- c(rep(NA, length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                      rep("NE", length(genereggroup[[paste0("FeatRemaining_G", group)]])))
  
  if (group > 1){
    df5$Gene_FC <- c(rep(1, length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                        rep(NA, length(genereggroup[[paste0("FeatRemaining_G", group)]])))
    df5$Peak_FC <- c(rep(NA, length(genereggroup[[paste0("GeneRemaining_G", group)]])), 
                        rep(1, length(genereggroup[[paste0("FeatRemaining_G", group)]])))
  }
  
  # Concat dataframes of DE genes and features
  df <- rbind(df1, df2, df3, df4, df5)
  
  if (group == 1){
    df$Gene_FC <- ifelse(is.na(df$Gene_ID), NA, 1)
    df$Peak_FC <- ifelse(is.na(df$Peak_ID), NA, 1)
  }
  
  # Now we need to sort the association dataframe according to the clusters
  # But cluster 0 should be last
  custom_sort <- function(x) {
    ifelse(x == 0, Inf, x)
  }
  df <- df[order(custom_sort(df$Gene_cluster)),]
  
  # Remove duplicated atac peaks
  df <- df %>%
        group_by(Peak_ID) %>%
        filter(!(is.na(Gene_ID) & duplicated(Peak_ID)))
  
  dfGeneNames <- as.character(stats::na.omit(df$Gene_ID))
  
  df2 <- df[order(custom_sort(df$Peak_cluster)),]
  
  
  
  dfPeakNames <- unique(as.character(stats::na.omit(df2$Peak_ID)))
  
  return(list("associationMatrix" = df, "dfPeakNames" = dfPeakNames, "dfGeneNames" = dfGeneNames))
  
}


#' calculate_mean_per_list_df
#' 
#' Helper function to calculate mean expression per celltype
#'
#' @param df dataframe of expression where columns are cells
#' @param named_lists list of which cells belong to each celltype
#'
#' @export
#'
#' @examples
#' rna <- data.frame(c1 = c(1.5, 15.5, 3.5, 20.5), c2 = c(2, 15, 4, 20), 
#'           c3 = c(10, 1, 12, 13), c4 = c(11, 1, 13, 14))
#' cell_types <- list("ct1" = c(1,2), "ct2" = c(3, 4))
#' calculate_mean_per_list_df(rna, cell_types)
calculate_mean_per_list_df <- function(df, named_lists) {
  means <- list()
  for (name in names(named_lists)) {
    columns <- named_lists[[name]]
    means[[name]] <- rowMeans(df[, columns, drop = FALSE])
  }
  # Combine the list of means into a dataframe
  means_df <- do.call(cbind, means)
  # Add column names
  colnames(means_df) <- names(named_lists)
  return(as.data.frame(means_df))
}


#' match_gene_regulator
#' 
#' Helper function to make the most similar profiles possible between
#' gene and regulator
#'
#' @param rna dataframe of RNA expression
#' @param atac dataframe of ATAC expression
#' @param cell_types list of which cells belong to each celltype
#' @param associationList dataframe of two columns, Gene_ID and Peak_ID
#'
#' @export
#'
#' @examples
#' rna <- data.frame(c1 = c(1.5, 15.5, 3.5, 20.5), c2 = c(2, 15, 4, 20), 
#'         c3 = c(10, 1, 12, 13), c4 = c(11, 1, 13, 14), c5 = c(7, 0, 0, 0), 
#'         c6 = c(8, 1, 1, 1), c7 = c(8, 1, 1, 1))
#' rownames(rna) <- c('GenB', 'GenA', 'GenC', 'GenD')
#' associationList <- data.frame(Gene_ID = c('GenA', 'GenB', 'GenC', 'GenA'),
#'         Peak_ID = c('PeakA', 'PeakB', 'PeakC', 'PeakD'))
#' cell_types <- list("ct1" = c(1,2), "ct2" = c(3, 4), "ct3" = c(5, 6), "ct4" = c(7))
#' atac <- data.frame(c1 = c(3,20, 1,15, 1, 7, 1), c2 = c(4,20, 2,15, 0, 5, 1.5), 
#'         c3 = c(10, 13, 1, 12, 1, 14, 9), c4 = c(11, 14, 1, 13, 1, 4, 12), 
#'         c5 = c(0, 0, 0, 7, 1, 6, 6), c6 = c(1, 1, 1, 8, 0, 5, 8), 
#'         c7 = c(1, 1, 1, 8, 1, 5, 7))
#' rownames(atac) <- c('PeakB', "PeakC", "PeakF", "PeakD", "PeakE", "PeakA", "PeakG")
#' match_gene_regulator(rna, atac, cell_types, associationList)
match_gene_regulator <- function(rna, atac, cell_types, associationList){
  ## Prepare the dataframes:
  Gene_ID <- NULL
  dup_index <- NULL
  
  # Add a suffix to each duplicate occurrence in associationList
  associationList <- associationList %>%
    dplyr::group_by(Gene_ID) %>%
    dplyr::mutate(dup_index = row_number() - 1,
           Gene_ID = ifelse(dup_index > 0, paste0(Gene_ID, "_dup", dup_index), Gene_ID)) %>%
    dplyr::select(-dup_index) %>%
    dplyr::ungroup()
  
  # Check for duplicates in Gene_ID column
  duplicate_genes <- associationList$Gene_ID[grepl("_dup", associationList$Gene_ID)]
  
  # Duplicate rows in rna dataframe based on duplicates in associationList
  for (gene in duplicate_genes) {
    # Extract the base gene name
    base_gene <- sub("_dup[0-9]+$", "", gene)
    
    # Find the row in rna dataframe corresponding to the base gene
    rna_row <- rna[rownames(rna) == base_gene, ]
    
    # Append the row to the rna dataframe
    rna <- rbind(rna, rna_row)
    
    # Update the row names
    rownames(rna)[nrow(rna)] <- gene
  }
  
  
  ## Sort dataframes according to the association list order
  rna <- rna[order(match(rownames(rna), as.vector(associationList$Gene_ID))), , drop = FALSE]
  atac <- atac[order(match(rownames(atac), as.vector(associationList$Peak_ID))), , drop = FALSE]
  
  mean_rna <- MOSim::calculate_mean_per_list_df(rna, cell_types)
  mean_atac <- MOSim::calculate_mean_per_list_df(atac, cell_types)
  
  # Target values
  target_values <- as.numeric(as.data.frame(mean_rna)[[1]])
  
  # Function to find the most similar element and its index, avoiding used indexes
  find_most_similar <- function(values, target, used_indexes) {
    # Calculate the absolute differences, ignoring used indexes
    differences <- abs(values - target)
    differences[used_indexes] <- Inf  # Set used indexes to Inf to avoid selecting them
    min_index <- which.min(differences)
    return(list(element = values[min_index], index = min_index))
  }
  
  # Initialize list to store results
  results <- list(elements = numeric(length(target_values)), indexes = integer(length(target_values)))
  
  # Vector to keep track of used indexes
  used_indexes <- integer(0)
  
  # Find most similar elements and their indexes
  for (i in seq_along(target_values)) {
    res <- find_most_similar(mean_atac[[1]], target_values[i], used_indexes)
    results$elements[i] <- res$element
    results$indexes[i] <- res$index
    used_indexes <- c(used_indexes, res$index)
  }
  
  new_mean_atac <- data.frame(results$elements)
  colnames(new_mean_atac) <- names(mean_rna)[[1]]
  new_indexes_atac <- data.frame(results$indexes)
  colnames(new_indexes_atac) <- names(mean_rna)[[1]]
  
  for (j in 2:length(colnames(mean_atac))) {
    prev_ct <- names(mean_rna)[j - 1]
    curr_ct <- names(mean_rna)[j]
    
    differences_rna <- mean_rna[[prev_ct]] - mean_rna[[curr_ct]]
    results_curr <- data.frame(elements = numeric(length(differences_rna)), indexes = integer(length(differences_rna)))
    used_indexes_curr <- integer(0)
    
    for (i in seq_along(differences_rna)) {
      target_diff <- differences_rna[i]
      res <- find_most_similar(mean_atac[[curr_ct]], new_mean_atac[[prev_ct]][i] - target_diff, used_indexes_curr)
      results_curr$elements[i] <- res$element
      results_curr$indexes[i] <- res$index
      used_indexes_curr <- c(used_indexes_curr, res$index)
    }
    
    new_mean_atac[[curr_ct]] <- results_curr$elements
    new_indexes_atac[[curr_ct]] <- results_curr$indexes
  }
  
  # Get unique values in each column
  unique_values <- lapply(new_indexes_atac, unique)
  
  # Find missing numbers from 1 to 10 for each column
  missing_numbers <- lapply(unique_values, function(x) setdiff(1:nrow(mean_atac), x))
  
  # Add missing indexes
  new_indexes_atac <- rbind(new_indexes_atac, missing_numbers)
  
  ## Get the values from those indices from atac
  new_atac <- data.frame(matrix(ncol = ncol(atac), nrow = nrow(atac)))
  colnames(new_atac) <- colnames(atac)
  rownames(new_atac) <- rownames(atac)
  
  # Populate the new_atac data frame with the values from atac according to new_indexes_atac and cell_types
  for (cell_type in colnames(new_indexes_atac)) {
    # Get the indices for the current cell type
    indices <- new_indexes_atac[[cell_type]]
    
    # Get the columns in atac that correspond to the current cell type
    columns <- cell_types[[cell_type]]
    
    # Extract and reorder values for each column according to the indices
    for (col in columns) {
      new_atac[, col] <- atac[indices, col]
    }
  }
  
  # Function to remove the added rows
  remove_added_rows <- function(dataframe) {
    return(dataframe[!grepl("_dup[0-9]+$", rownames(dataframe)), ])
  }
  
  # Remove the added rows
  rna <- remove_added_rows(rna)
  atac <- new_atac
  return(list("atac" = atac, "rna" = as.data.frame(as.matrix(rna))))
}



#' match_gene_regulator_cluster
#'
#' @param rna rna expression dataframe
#' @param atac atac expression dataframe
#' @param cell_types list of which cells belong to each celltype
#' @param associationMatrix matrix of related genes and peaks
#'
#' @export
#'
#' @examples
#' rna <- data.frame(c1 = c(1.5, 15.5, 3.5, 20.5), c2 = c(2, 15, 4, 20), 
#'         c3 = c(10, 1, 12, 13), c4 = c(11, 1, 13, 14), c5 = c(7, 0, 0, 0), 
#'         c6 = c(8, 1, 1, 1), c7 = c(8, 1, 1, 1))
#' rownames(rna) <- c('GenB', 'GenA', 'GenC', 'GenD')
#' associationList <- data.frame(Gene_ID = c('GenA', 'GenB', 'GenC', 'GenA'),
#'         Peak_ID = c('PeakA', 'PeakB', 'PeakC', 'PeakD'),
#'         Gene_cluster = c(1, 2, 1, 2), Peak_cluster = c(1, 2, 1, 2))
#' cell_types <- list("ct1" = c(1,2), "ct2" = c(3, 4), "ct3" = c(5, 6), "ct4" = c(7))
#' atac <- data.frame(c1 = c(3,20, 1,15, 1, 7, 1), c2 = c(4,20, 2,15, 0, 5, 1.5), 
#'         c3 = c(10, 13, 1, 12, 1, 14, 9), c4 = c(11, 14, 1, 13, 1, 4, 12), 
#'         c5 = c(0, 0, 0, 7, 1, 6, 6), c6 = c(1, 1, 1, 8, 0, 5, 8), 
#'         c7 = c(1, 1, 1, 8, 1, 5, 7))
#' rownames(atac) <- c('PeakB', "PeakC", "PeakF", "PeakD", "PeakE", "PeakA", "PeakG")
#' match_gene_regulator_cluster(rna, atac, cell_types, associationList)
match_gene_regulator_cluster <- function(rna, atac, cell_types, associationMatrix){
  Gene_cluster <- NULL
  Peak_cluster <- NULL
  
  clusters <- unique(associationMatrix$Gene_cluster)
  
  # remove cluster 0 since these are the unorganized.
  clusters <- clusters[clusters != 0 & !is.na(clusters)]
  
  for (e in clusters){
    # Make mini association lists of each cluster and reorganize them
    miniAsoc <- associationMatrix %>%
      dplyr::filter((Gene_cluster == e & (Peak_cluster == e | is.na(Peak_cluster))) |
                      (Peak_cluster == e & (Gene_cluster == e | is.na(Gene_cluster))))
    
    miniATAC <- atac[row.names(atac) %in% as.character(stats::na.omit(miniAsoc$Peak_ID)), ]
    
    associationList <- stats::na.omit(miniAsoc[, c("Gene_ID", "Peak_ID")])
    
    miniRNA <- rna[row.names(rna) %in% as.character(associationList$Gene_ID), ]
    
    res <- MOSim::match_gene_regulator(miniRNA, miniATAC, cell_types, associationList)
    
    ## Update RNA and ATAC dataframes with the changes we have made
    atac <- as.data.frame(as.matrix(atac)) %>%
      tibble::rownames_to_column() %>%
      dplyr::rows_update(res$atac %>% tibble::rownames_to_column(), by = "rowname") %>%
      tibble::column_to_rownames()
    
    rna <- as.data.frame(as.matrix(rna)) %>%
      tibble::rownames_to_column() %>%
      dplyr::rows_update(res$rna %>% tibble::rownames_to_column(), by = "rowname") %>%
      tibble::column_to_rownames()
    
  }
  return(list("atac" = atac, "rna" = rna))
}



