# Written by: Reilly Pidgeon

# Script to gather and analyze the relative abundance (ra) and gene families (gf) of specific taxa and genes in fecal metagenomics samples
# Make sure to run this on the Compute Canada Clusters since it'll crash your computer
# This was run on the Narval cluster (Digital Research Alliance of Canada)

  # Info on curatedMetagenomicData can be found here: 
  # Bioconductor info: https://bioconductor.riken.jp/packages/3.14/data/experiment/manuals/curatedMetagenomicData/man/curatedMetagenomicData.pdf
  # All available studies: https://waldronlab.io/curatedMetagenomicData/articles/available-studies.html
  # General: https://waldronlab.io/curatedMetagenomicData/articles/curatedMetagenomicData.html

# The commented lines below were put in a shell script (remove the first #)
  # #!/bin/bash
  # #SBATCH --time=6:00:00
  # #SBATCH --mem-per-cpu=100000M
  # #SBATCH --cpus-per-task=8
  # #SBATCH --job-name="cMD_ALL_STUDIES_RData"
  # #SBATCH --account=def-castagne
  # module load gcc/9.3.0 r/4.2.2
  # Rscript 2_cMD_FromRData.R

# Load libraries for curatedMetagenomicData Analysis 
library(curatedMetagenomicData)
library(TreeSummarizedExperiment)
library(dplyr)
library(stringr)
library(mia)
library(scater)
library(vegan)
library(DT)
library(ggplot2)
library(microbiome)
library(tibble)
library(stringr)
library(data.table)
print("Libraries loaded")

# Import all studies (as RData files) from the cMG_RData_Studies_Snapshot2023-04-24 (folder)
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies1-3.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies4-6.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies7-11.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies12-14.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies15-18.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies19-23.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies24-26.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies27-28.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies29-35.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies36-39.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies40-41.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies42-46.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies47-50.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies51.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies52-57.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies58-62.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies63-67.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies68.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies69-78.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies79-85.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies86-87.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies88-90.RData")
load("../cMG_RData_Studies_Snapshot2023-04-24/Studies91-93.RData")
print("RData loaded for all studies.")

# Set variables of interest

  # Bacteria of interest (boi) (species name) (separate by | to include all in the same table)
  # The partial strings listed below enable the extraction of data for E. bolteae, E. citroniae, E. asparagiformis, L. pacaense (47FAA)
  boi <- "bolteae|citroniae|asparagiformis|asparagiforme|Enterocloster|47FAA"
  # Genes of interest (goi) in gene family table
  # The codes below correspond to the Molybdopterin-dependent oxidoreductase (ucdO), the 2Fe-2S cluster (ucdF), and the xanthine dehydrogenase (ucdC), respectively
  goi <- "C5EGQ4|G5HFF3|A8RZR5"
  
print("Variables set:")
print(paste("BACTERIA OF INTEREST:", boi))
print(paste("GENES OF INTEREST:", goi))

#==========================================CUSTOM-FUNCTIONS======================================================

# Create functions to fetch data in list according to study#

# Fetches metadata from the relative abundance dataset (same for all other data types in cMG) - inputs:  study = "Study#" (String)
fetch_md <- function(study) {
  temp <- get(study)
  study_item_names <- names(temp)
  
  # Get metadata
  md_num <- grep(".relative_abundance", study_item_names)
  print(paste("Fetching", study_item_names[md_num]))
  study_list_data <- study_item_names[md_num]
  assign("md", as.data.frame(temp[[study_list_data]]@colData))
  md <- tibble::rownames_to_column(md, "id")
  write.csv(md, file = paste(study, "_md.csv", sep = ""))
  # Return a dataframe into the Global Environment
  return(assign(paste(study, "_md", sep = ""), md, envir = .GlobalEnv))
}

# Create functions to query dataframes for strings of interest (relative abundance (ra) and gene families (gf))

# Query relative abundance based on bacteria of interest
query_ra <- function(study, boi) { # study variable should be a dataframe with name "Study#_ra" - works with the output of fetch_ra
  if(class(study) == "character") {
    temp <- get(study)
  }
  if(class(study) == "data.frame") {
    temp <- study
  }
  
  ra_boi <- temp[str_detect(temp$taxa, boi), ]
}

# Query gene families based on genes of interest
query_gf <- function(study, goi) {
  if(class(study) == "character") {
    temp <- get(study)
  }
  if(class(study) == "data.frame") {
    temp <- study
  }
  
  gf_goi <- temp[str_detect(temp$gene_families, goi), ]
}

# Fetch relative abundance - inputs: study = "Study#" (String), e.g. fetch_ra_boi("Study4")
# Saves a csv file in the wd
fetch_ra_boi <- function(study) {
  temp <- get(study)
  study_item_names <- names(temp)
  
  # Get relative abundance values
  ra_num <- grep(".relative_abundance", study_item_names)
  print(paste("Fetching", study_item_names[ra_num]))
  study_list_data <- study_item_names[ra_num]
  assign("ra", as.data.frame(temp[[study_list_data]]@assays@data@listData[["relative_abundance"]])) #This is now a dataframe
  ra <- tibble::rownames_to_column(ra, "taxa") # This gives the first column the names of taxa (move rownames to first column)
  print(boi)
  ra_boi <- query_ra(study = ra, boi = boi)
  write.csv(ra_boi, file = paste(study, "_ra_boi.csv", sep = ""))
  return(assign(paste(study, "_ra_boi", sep = ""), ra_boi, envir = .GlobalEnv))
}

# Fetch gene family abundance (RPK units) - inputs: study = "Study#" (String), e.g. fetch_gf("Study4")
# Saves a csv file in the wd
fetch_gf_goi <- function(study) {
  temp <- get(study)
  study_item_names <- names(temp)
  
  # Get gene family values
  gf_num <- grep(".gene_families", study_item_names)
  print(paste("Fetching", study_item_names[gf_num]))
  study_list_data <- study_item_names[gf_num]
  assign("gf", as.data.frame(as.matrix(temp[[study_list_data]]@assays@data@listData[["gene_families"]]))) #This is now a data frame
  gf <- tibble::rownames_to_column(gf, "gene_families") # This gives the first column the names of gene_families (move rownames to first column)
  gf_goi <- query_gf(study = gf, goi = goi)
  write.csv(gf_goi, file = paste(study, "_gf_goi.csv", sep = ""))
  return(assign(paste(study, "_gf_goi", sep = ""), gf_goi, envir = .GlobalEnv))
}


#============THE FUNCTIONS ARE NOW SET - TIME TO ITERATE THROUGH THE STUDIES======================

# Iterate through the data in the Global Environment and fetch the ra and gf dataframes
for (i in 1:93) {
  # Store the "Study#" into s - this changes with each iteration of the loop
  s <- paste("Study",i,sep="")
  print(paste("Working on", s))
  
  # Fetch the metadata, relative abundance, and gene families (saved as dataframes into Global Environment and as csv files throughout the iterations)
  fetch_md(study = s)
  print("Metadata saved")
  fetch_ra_boi(study = s)
  print("Relative abundance for boi saved")
  fetch_gf_goi(study = s)
  print("Gene families for goi saved")
}  


# Save R environment (.RData)
save.image(file = "cMD_ra_aa_gf.RData")
print("Script done running!")