# Loads the necessary libraries for data analysis and visualization
library(reshape2)
library(ggplot2)
library(dplyr)
library(RColorBrewer)
library(gplots)
library(pheatmap)
library(Seurat)

# Loads the GSEA results from an RDS file
gsea_res <- readRDS("gsea_intracluster_OAT_NOA_vs_FER.rds")

# Removes comparisons between NOA and OAT
index<- c(3, 6 ,9 ,12, 15, 18, 21, 24, 27)
gsea_res<-gsea_res[-index] 

# Divides the results into two groups: OAT vs FER and NOA vs FER
gsea_OAT_vs_FER<- gsea_res[-c(2,4,6,8,10,12,14,16,18)]
gsea_NOA_vs_FER <- gsea_res[-c(1,3,5,7,9,11,13,15,17)]
# Combines the two groups
gsea_res <- append(gsea_OAT_vs_FER, gsea_NOA_vs_FER)

# Initializes empty list and dataframe
gsea_df_melt <- list()
counter <- 1

gsea_df <- NULL

# Extracts GSEA results for each comparison
for(i in names(gsea_res)){
  print(i)
  if(is.null(gsea_res[[i]])){
    next()
  }
  else if (nrow(gsea_res[[i]]@result) == 0){
    next
  }
  else{
    gsea_df[[i]] <- gsea_res[[i]]@result
  }
}


# Defines the correct order for the classifications
right_classification <- c("CD4+ Tfh_OAT_vs_FER", "CD4+ Th2_OAT_vs_FER", "CD4+ Th1/Th17_OAT_vs_FER", "CD4+ CD8+_OAT_vs_FER", "CD4+ CD8+ TEM_OAT_vs_FER", "CD8+ TCM(Gzk+)_OAT_vs_FER", "CD8+ TE(TOX+)_OAT_vs_FER", "CD8+ TE (GzK+)_OAT_vs_FER", "CD8+ TEM_OAT_vs_FER","CD4+ Tfh_NOA_vs_FER", "CD4+ Th2_NOA_vs_FER", "CD4+ Th1/Th17_NOA_vs_FER", "CD4+ CD8+_NOA_vs_FER", "CD4+ CD8+ TEM_NOA_vs_FER", "CD8+ TCM(Gzk+)_NOA_vs_FER", "CD8+ TE(TOX+)_NOA_vs_FER", "CD8+ TE (GzK+)_NOA_vs_FER", "CD8+ TEM_NOA_vs_FER")


# Gets the indices to reorder the dataframe
order_indices <- match(right_classification, names(gsea_df))

# Reorder the dataframe
gsea_df<- gsea_df[order_indices]

# Extracts information from each dataframe
for (i in names(gsea_df)) {
  cluster <- unlist(strsplit(x = i, split = "_"))[1]
  comparison <- unlist(sub(x = i, "^.+?_", ""))
  gsea_df[[i]]$cluster <- cluster
  gsea_df[[i]]$comparison <- comparison
  gsea_df[[i]]$ID <- gsub(gsea_df[[i]]$ID, pattern = "HALLMARK_", replacement = "")
  gsea_df_melt[[i]] <- melt(data = gsea_df[[i]], id.vars = c("ID","p.adjust","comparison", "cluster"),  measure.vars = c("NES"))
  gsea_df_melt[[i]]$value <- round(gsea_df_melt[[i]]$value, digits = 2)
  counter <- counter + 1
}

# Combines all dataframes into a single dataframe
total <- bind_rows(gsea_df, .id = "Classification")
total_melt <- bind_rows(gsea_df_melt, .id = "Classification")

#min(total_melt$value, na.rm = T)#-2.4
#max(total_melt$value, na.rm = T)#2.37

# Defines the levels of the factors for ordering variables
total_melt$ID <- factor(total_melt$ID, levels = term_right)
total_melt$Classification <- factor(total_melt$Classification, levels = unique(right_classification))

# Transforms the data into a format suitable for visualization
total_cast_melt <- melt(dcast(data = total_melt, formula = ID ~ Classification, value.var = 'value'), id.vars = 'ID' ) 

# Extracts the cluster and comparison from the column labels
total_cast_melt$cluster <- stringr::str_split_fixed(total_cast_melt$variable, "_", 2)[,1]
total_cast_melt$comparison <- stringr::str_split_fixed(total_cast_melt$variable, "_", 2)[,2]

# Defines the GSEA terms and their groups

term_right<-c("NOTCH_SIGNALING",
              "MTORC1_SIGNALING",
              "KRAS_SIGNALING_DN",
              "WNT_BETA_CATENIN_SIGNALING",
              "TGF_BETA_SIGNALING",
              "PI3K_AKT_MTOR_SIGNALING",
              "EPITHELIAL_MESENCHYMAL_TRANSITION",
              "MYOGENESIS",
              "APICAL_JUNCTION",
              
              "COAGULATION",
              "COMPLEMENT",
              "ALLOGRAFT_REJECTION",
              "IL2_STAT5_SIGNALING",
              "IL6_JAK_STAT3_SIGNALING",
              "INFLAMMATORY_RESPONSE",
              "INTERFERON_GAMMA_RESPONSE",
              "INTERFERON_ALPHA_RESPONSE",
              "TNFA_SIGNALING_VIA_NFKB",
              
              "G2M_CHECKPOINT",
              "MITOTIC_SPINDLE",
              "MYC_TARGETS_V2",
              "MYC_TARGETS_V1",
              
              "APOPTOSIS",
              "UV_RESPONSE_DN",
              "UV_RESPONSE_UP",
              "P53_PATHWAY",
              "UNFOLDED_PROTEIN_RESPONSE",
              "PROTEIN_SECRETION",
              
              "GLYCOLYSIS",
              "HYPOXIA",
              "CHOLESTEROL_HOMEOSTASIS",
              "ESTROGEN_RESPONSE_EARLY",
              "ESTROGEN_RESPONSE_LATE",
              "FATTY_ACID_METABOLISM",
              "BILE_ACID_METABOLISM")


group_letter <- c(rep("A_CellSignaling", 9), rep("B_ImmuneResponse", 9), rep("C_CellCycle", 4), rep("D_CellStress", 6), rep("E_Metabolism", 7))
terms_group <- data.frame(term_right = term_right, group_letter = group_letter)

# Merges the data with the information about the GSEA term groups
total_cast_melt <- merge(total_cast_melt, terms_group, by.x = "ID", by.y = "term_right", all.x = T)
# Define the order of comparison
total_cast_melt$comparison <- factor(total_cast_melt$comparison, levels = c("OAT_vs_FER", "NOA_vs_FER"))
# Define the order of clusters
total_cast_melt$cluster <- factor(total_cast_melt$cluster, levels = c("CD4+ Tfh","CD4+ Th2","CD4+ Th1/Th17","CD4+ CD8+","CD4+ CD8+ TEM","CD8+ TCM(Gzk+)","CD8+ TE(TOX+)","CD8+ TE (GzK+)","CD8+ TEM"))
# Define the order of groups (pathways)
total_cast_melt$group_letter <- factor(total_cast_melt$group_letter, levels = c("E_Metabolism","D_CellStress","C_CellCycle","B_ImmuneResponse","A_CellSignaling"))

#sort(unique(total_cast_melt$group_letter))

# Create the GSEA tile plot using ggplot2
gsea_tile_plot <-ggplot(data = total_cast_melt, mapping = aes(x = cluster, y = ID))+  
  facet_grid(group_letter~comparison, scales = "free", space = "free") + 
  geom_tile(aes(fill=value), width=0.80, height=0.80, color="white",lwd = 0.1,linetype = 5, size=0.20) + labs(fill = "NES") + 
  ggtitle(label = "GSEA hallmarks", subtitle = " Intracluster -  padj < 0.1")+
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1, size=10),
        strip.text.x = element_text(size = 16),
        strip.text.y = element_text(size =0),
        strip.background = element_blank(),
        plot.title = element_text(hjust = 0, size = 18),
        plot.subtitle = element_text(hjust = 0, size = 12),
        axis.ticks = element_blank(),
        axis.text.y = element_text(size = 9),
        axis.title = element_blank(),
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.spacing.x = unit(.3, "cm"),
        panel.spacing.y = unit(.3,"cm"),
        panel.background = element_rect(fill = "white")) +
  scale_fill_gradient2(na.value = '#ECECED', 
                       limits=c(-2.5, 2.5), breaks=seq(-3,3,by=1),
                       low='#005295',high='#770000', mid = 'white') # +

# Save the plot to a file
ggsave(filename = "GSEA_tile_plot.png", width = 7, height = 6, plot =gsea_tile_plot)
