library(Seurat)
library(SingleCellExperiment)
library(dplyr)
library(ggplot2)
library(nichenetr)
library(multinichenetr)

# Set Folders
wdir <- "./MultiNicheNet"

obj <- readRDS(file = paste(wdir, "Full_dataset.rds", sep = "/"))

organism = "mouse"
lr_network = readRDS(file = paste(wdir, "lr_network_mouse_21122021.rds", sep = "/"))
lr_network = lr_network %>% dplyr::rename(ligand = from, receptor = to) %>% distinct(ligand, receptor) %>%
  mutate(ligand = make.names(ligand), receptor = make.names(receptor))
ligand_target_matrix = readRDS(file = paste(wdir, "ligand_target_matrix_nsga2r_final_mouse.rds", sep = "/"))
colnames(ligand_target_matrix) = colnames(ligand_target_matrix) %>% make.names()
rownames(ligand_target_matrix) = rownames(ligand_target_matrix) %>% make.names()

# Annotation Custom
obj@meta.data$AnnotationCART_MAGIC <- obj@meta.data$Annotation
obj@meta.data <- mutate(obj@meta.data, AnnotationCART_MAGIC = case_when(obj@meta.data$Annotation %in% c("TAM1", "TAM2_Il1b", "TAM3_AltAct", "DC1", "DC2") ~ obj@meta.data$Annotation,
                                                                        obj@meta.data$CART_MAGIC == "CART_MAGIC" ~ "CART_MAGIC",
                                                                        TRUE ~ "NA"))

sce <- Seurat::as.SingleCellExperiment(obj, assay = "RNA")
sce <- alias_to_symbol_SCE(sce, "mouse") %>% makenames_SCE()

sample_id = "orig.ident"
group_id = "PROP.Condition"
celltype_id = "AnnotationCART_MAGIC"
covariates = NA
batches = NA

senders_oi <- c("TAM1","TAM2_Il1b","TAM3_AltAct", "DC1", "DC2")
receivers_oi <- c("CART_MAGIC")

sce = sce[, SummarizedExperiment::colData(sce)[,celltype_id] %in% c(senders_oi, receivers_oi)]
SummarizedExperiment::colData(sce)$orig.ident = SummarizedExperiment::colData(sce)$orig.ident %>% make.names()

min_cells = 10
abundance_expression_info = get_abundance_expression_info(sce = sce,
                                                          sample_id = sample_id,
                                                          group_id = group_id,
                                                          celltype_id = celltype_id,
                                                          min_cells = min_cells,
                                                          senders_oi = senders_oi,
                                                          receivers_oi = receivers_oi,
                                                          lr_network = lr_network,
                                                          batches = batches)

contrasts_oi = c("'P2A-Empty','Ortho-Empty','IFN-Empty'")
contrast_tbl = tibble(contrast = c("P2A-Empty","Ortho-Empty","IFN-Empty"),
                      group = c("P2A", "Ortho", "IFN"))

DE_info = get_DE_info(sce = sce,
                      sample_id = sample_id,
                      group_id = group_id,
                      celltype_id = celltype_id,
                      batches = batches,
                      covariates = covariates,
                      contrasts_oi = contrasts_oi,
                      min_cells = min_cells)
DE_info$celltype_de$de_output_tidy %>% arrange(p_adj) %>% head()
celltype_de = DE_info$celltype_de$de_output_tidy

sender_receiver_de = combine_sender_receiver_de(
  sender_de = celltype_de,
  receiver_de = celltype_de,
  senders_oi = senders_oi,
  receivers_oi = receivers_oi,
  lr_network = lr_network
)

logFC_threshold = 0.50
p_val_threshold = 0.05
fraction_cutoff = 0.05
p_val_adj = FALSE 
top_n_target = 250
verbose = TRUE
cores_system = 4
n.cores = min(cores_system, sender_receiver_de$receiver %>% unique() %>% length())

ligand_activities_targets_DEgenes = suppressMessages(suppressWarnings(get_ligand_activities_targets_DEgenes(
  receiver_de = celltype_de,
  receivers_oi = receivers_oi,
  ligand_target_matrix = ligand_target_matrix,
  logFC_threshold = logFC_threshold,
  p_val_threshold = p_val_threshold,
  p_val_adj = p_val_adj,
  top_n_target = top_n_target,
  verbose = verbose, 
  n.cores = n.cores
)))

prioritizing_weights_DE = c("de_ligand" = 1,
                            "de_receptor" = 1)
prioritizing_weights_activity = c("activity_scaled" = 2)
prioritizing_weights_expression_specificity = c("exprs_ligand" = 2,
                                                "exprs_receptor" = 2)
prioritizing_weights_expression_sufficiency = c("frac_exprs_ligand_receptor" = 1)
prioritizing_weights_relative_abundance = c( "abund_sender" = 0,
                                             "abund_receiver" = 0)

prioritizing_weights = c(prioritizing_weights_DE, 
                         prioritizing_weights_activity, 
                         prioritizing_weights_expression_specificity,
                         prioritizing_weights_expression_sufficiency, 
                         prioritizing_weights_relative_abundance)

sender_receiver_tbl = sender_receiver_de %>% dplyr::distinct(sender, receiver)

metadata_combined = SummarizedExperiment::colData(sce) %>% tibble::as_tibble()

if(!is.na(batches)){
  grouping_tbl = metadata_combined[,c(sample_id, group_id, batches)] %>% tibble::as_tibble() %>% dplyr::distinct()
  colnames(grouping_tbl) = c("sample","group",batches)
} else {
  grouping_tbl = metadata_combined[,c(sample_id, group_id)] %>% tibble::as_tibble() %>% dplyr::distinct()
  colnames(grouping_tbl) = c("sample","group")
}

prioritization_tables = suppressMessages(generate_prioritization_tables(
  sender_receiver_info = abundance_expression_info$sender_receiver_info,
  sender_receiver_de = sender_receiver_de,
  ligand_activities_targets_DEgenes = ligand_activities_targets_DEgenes,
  contrast_tbl = contrast_tbl,
  sender_receiver_tbl = sender_receiver_tbl,
  grouping_tbl = grouping_tbl,
  prioritizing_weights = prioritizing_weights,
  fraction_cutoff = fraction_cutoff, 
  abundance_data_receiver = abundance_expression_info$abundance_data_receiver,
  abundance_data_sender = abundance_expression_info$abundance_data_sender
))
prioritization_tables$group_prioritization_tbl %>% head(20)

lr_target_prior_cor = lr_target_prior_cor_inference(prioritization_tables$group_prioritization_tbl$receiver %>% unique(),
                                                    abundance_expression_info,
                                                    celltype_de, grouping_tbl,
                                                    prioritization_tables,
                                                    ligand_target_matrix,
                                                    logFC_threshold = logFC_threshold,
                                                    p_val_threshold = p_val_threshold,
                                                    p_val_adj = p_val_adj)

multinichenet_output = list(
  celltype_info = abundance_expression_info$celltype_info,
  celltype_de = celltype_de,
  sender_receiver_info = abundance_expression_info$sender_receiver_info,
  sender_receiver_de =  sender_receiver_de,
  ligand_activities_targets_DEgenes = ligand_activities_targets_DEgenes,
  prioritization_tables = prioritization_tables,
  grouping_tbl = grouping_tbl,
  lr_target_prior_cor = lr_target_prior_cor
) 
multinichenet_output = make_lite_output(multinichenet_output)

saveRDS(multinichenet_output, paste(wdir, "Full_WPRE_MNNout_CARTvsEmpty.rds", sep = "/"))

# Plots circos
prioritized_tbl_oi_all = get_top_n_lr_pairs(multinichenet_output$prioritization_tables, 75, rank_per_group = FALSE)
prioritized_tbl_oi = multinichenet_output$prioritization_tables$group_prioritization_tbl %>%
  filter(id %in% prioritized_tbl_oi_all$id) %>%
  distinct(id, sender, receiver, ligand, receptor, group) %>% left_join(prioritized_tbl_oi_all)
prioritized_tbl_oi$prioritization_score[is.na(prioritized_tbl_oi$prioritization_score)] = 0
senders_receivers = union(prioritized_tbl_oi$sender %>% unique(), prioritized_tbl_oi$receiver %>% unique()) %>% sort()
colors_sender = RColorBrewer::brewer.pal(n = length(senders_receivers), name = 'Spectral') %>% magrittr::set_names(senders_receivers)
colors_receiver = RColorBrewer::brewer.pal(n = length(senders_receivers), name = 'Spectral') %>% magrittr::set_names(senders_receivers)
circos_list = make_circos_group_comparison(prioritized_tbl_oi, colors_sender, colors_receiver)

