#!/usr/bin/env Rscript
# Jane A. Pascar
# 2021-09-29

# Usage: nohup Rscript --vanilla aedes-aegypti-microbiome.R args[1] args[2] args[3] args[4] args[5] > aedes-aegypti-microbiome.out &

#args[1] = [~/aedes_aegypti_microbiome/raw_data/bracken-output/] path to bracken directory
#args[2] = [~/aedes_aegypti_microbiome/raw_data/kreports/] path to kreport directory
#args[3] = [~/aedes_aegypti_microbiome/raw_data/] path to raw data files
#args[4] = [~/aedes_aegypti_microbiome/filtered_data/] path to file output directory
#args[5] = [~/aedes_aegypti_microbiome/figures/] path to figure output directory
#args[6] = [~/aedes_aegypti_microbiome/supplemental_tables/] path to table output directory

#args <- c("~/Documents/GitHub/aedes_aegypti_microbiome/raw_data/bracken-output/", 
#          "~/Documents/GitHub/aedes_aegypti_microbiome/raw_data/kreports/", 
#          "~/Documents/GitHub/aedes_aegypti_microbiome/raw_data/",
#          "~/Documents/GitHub/aedes_aegypti_microbiome/filtered_data/", 
#          "~/Documents/GitHub/aedes_aegypti_microbiome/figures/",
#          "~/Documents/GitHub/aedes_aegypti_microbiome/supplemental_tables/")

args = commandArgs(trailingOnly=TRUE)

### Check to make sure that arguments are supplied -----
# Script will stop running if 5 arguments are not correctly supplied.
rlang::inform('\U0001F440 Check that the correct number of arguments was supplied... \n')
if (length(args) < 5) {
  stop('\U0001F6A8 Missing arguements! See script for more detailed instructions \U0001F6A8', call.=FALSE)
} else {
  rlang::inform('All arguments supplied. \U0002705 \n')
}

### Load necessary libraries ----
library(readr); library(dplyr); library(tidyr)
library(stringr); library(ggplot2); library(cowplot)
library(ggpubr); library(vegan); library(ape)
library(forcats); library(emmeans); library(multcomp)
library(multcompView)

### Check to make sure the files you need exist ----
rlang::inform('\U0001F440 Check that output directory contains necessary files... \n')
if (length(file.exists(c(paste(args[3], "ictv_viral_taxonomy.csv", sep = "")),
                       (paste(args[3], "asc_strain_filtered.csv", sep = "")),
                       (paste(args[3], "PRJNA412140_metadata.txt", sep = "")),
                       (paste(args[3], "missing_taxa2.csv", sep = "")),
                       (paste(args[3], "rankedlineage.dmp", sep = "")))) == 5) {
  rlang::inform('Necessary files all exist for downstream analysis. \U0002705 \n')
} else {
  stop('\U0001F6A8 Files are missing that are nessary. See script for more info. \U0001F6A8', call.=FALSE)
}

### 1. Concatenate bracken files into a single csv ----
# Purpose: The output from Bracken is a .tsv file for each sample accession ID that is run through the pipeline. 
# This function will take each of those files and add a column to include the sample accession ID and the taxonomic level of which the IDs are made.
# It will merge all samples into a single dataframe.
# Input: You must supply the absolute path to the directory that contains all of the Bracken reports and you must supply an output directory -- this corresponds to args[1] and args[4]
# Output: 
# 1. The directory that contains all files being concatenated 
# 2. Number of files matching character pattern [acc]_S.bracken 
# 3. concatenated_Species_bracken.csv file written to args[4]
rlang::inform('\U0001F99F 1. Starting to concatenate bracken files. \U0001F99F \n')

ranks <- c("Species") # change this if you want a different taxonomic level

# Check to see if args[4]/concatenated_Species_bracken.csv already exists.
# If it does, this step will be skipped. 
if (file.exists(paste(args[4], "concatenated_", ranks, "_bracken.csv", sep = ""))) {
  rlang::inform('1. File exists: concatenated_Species_bracken.csv \n Moving on to next step. \U0002705 \n')
} else {
  # Output filename: concatenated_Species_bracken.csv
  for (r in ranks) {
    # Make the list of files in the directory a vector
    files <- list.files(path = paste(args[1], r, sep = ""), 
                        pattern = "*.bracken", 
                        full.names = FALSE, 
                        recursive = TRUE)
    print(paste(args[1], r, sep = ""))
    print(length(files))
    # Create a vector of all accession numbers
    acc_num <- gsub("_.*$", "", files)
    filename <- paste(args[4], "concatenated_", r,"_bracken.csv", sep = "")
    # Initiate the output file with new column names
    handle <- file(filename)
    writeLines("name,taxonomy_id,taxonomy_lvl,kraken_assigned_reads,added_reads,new_est_reads,fraction_total_reads,run_accession,rank", handle)
    close(handle)
    for (acc in acc_num) {
      # Import a single run accession bracken output
      file_data <- read_tsv(paste(args[1], r, "/", acc, "_", substr(r, 0, 1), ".bracken", sep = ""), col_names = T, col_types = cols())
      # Add column for accession number
      file_data$run_accession <- acc
      # Add column for taxonomic rank
      file_data$rank <- r
      # Write the concatenated dataframe to the directory specified
      write_csv(file_data, filename, col_names = F, append = T)
    }
  }
  
  # check to make sure all of the accessions were added:
  cat_bracken <- read_csv(paste(args[4], 
                                "concatenated_", 
                                ranks, 
                                "_bracken.csv", 
                                sep = ""), 
                          guess_max = 1e4, 
                          col_names = T, 
                          col_types = cols())
  if (length(acc_num) == length(unique(cat_bracken$run_accession))) {
    rlang::inform('\U0001F31F 1. Completed! Bracken files concatenated successfully. \nFile written: concatenated_Species_bracken.csv \U0001F31F \n')
  } else {
    stop('\U0001F6A8 There are bracken files that were not added correctly. \U0001F6A8 \n', call.=FALSE)
  }
}

### 2. Clean rankedlineage.dmp ----
# Check to see if args[4]/cleaned_rankedlineage.csv already exists.
# If it does, this step will be skipped. 
rlang::inform('\U0001F99F 2. Starting to clean up rankedlineage.dmp \U0001F99F \n')
if (file.exists(paste(args[4], "cleaned_rankedlineage.csv", sep = ""))) {
  rlang::inform('2. File exists: cleaned_rankedlineage.csv \n Moving on to next step. \U0002705 \n')
} else {
  # Purpose: Bracken only outputs the taxa that the read matches but does not include the full taxonomy. 
  # NCBI has a file called rankedlineage.dmp that contains all of this information.
  # This function will clean up that file before it is merged
  # Available: https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/new_taxdump/
  # Output: cleaned_rankedlineage.csv - cleaned up file containing full taxonomies
  
  # the file is deliminated by "\t|\t" but I can't figure out how to have that read automatically so this is really messy
  # there are a bunch of hidden tabs since with the read_delim command I think I can only read one deliminator, "|"
  
  lineage <- read_delim(file = paste(args[3], "rankedlineage.dmp", sep = ""), 
                        delim = "|", 
                        col_names = F, 
                        col_types = cols(),
                        guess_max = 1e4)
  lineage <- lineage[, 1:10] # the last column is filled with NAs and is useless
  header <- c("tax_id", "tax_name", "species", "genus", "family", "order", "class", "phylum", "kingdom", "superkingdom")
  colnames(lineage) <- header
  
  # Currently missing values are indicated by two tabs
  # this replaces all of the missing values with NA
  for (h in header) {
    lineage[[h]] <- gsub("\t\t", "NA", fixed = FALSE, x = lineage[[h]])
  }
  
  # Need to remove the "\t" before and after the strings in each column
  # the first column only has "\t" at the end of the string
  for (h in header) {
    lineage[[h]] <- gsub("\t$", "", x = lineage[[h]])
  }
  names <- c("tax_name", "species", "genus", "family", "order", "class", "phylum", "kingdom", "superkingdom")
  
  for (n in names) {
    lineage[[n]] <- gsub("^\t", "", x = lineage[[n]])
  }
  
  # change appropriate columns to factors
  col.name <- c("species", "genus", "family", "order", "class", "phylum", "kingdom", "superkingdom")
  lineage[col.name] <- lapply(lineage[col.name], factor)  # as.factor() could also be used
  
  write_csv(lineage, paste(args[4], "cleaned_rankedlineage.csv", sep = ""), col_names = T)
  
  rlang::inform('\U0001F31F 2. Completed! File written: cleaned_rankedlineage.csv \U0001F31F \n')
}

### 3. Determine species IDs that are missing phylum information ----
# This function will merge cleaned_rankedlineage.csv with concatenated_Species_bracken.csv to fill in full taxonomic information for each species ID. Then it will identify any species that are missing phylum information (sometimes NCBI databases are missing data) and build a manually curated library to fill in that missing data.
# Input: 
# 1. concatenated_Species_bracken.csv (contains all species level Bracken data)
# 2. cleaned_rankedlineage.csv
# Output:
# 1. concatenated_species_with_incomplete_lineage.csv - the NCBI taxonomy file merged with the Bracken data (may still have missing phyla info)
# 2. phylum_key.csv - a manually curated library of missing phyla
rlang::inform('\U0001F99F 3. Starting to identify missing taxonomic information. \U0001F99F \n')

# Check to see if args[4]/concatenated_Species_bracken.csv and args[4]/phylum_key.csv already exists.
# If it does, this step will be skipped. 
if (file.exists(paste(args[4], "concatenated_species_with_incomplete_lineage.csv", sep = "")) & file.exists(paste(args[4], "phylum_key.csv", sep = ""))) {
  rlang::inform('3. Files exist: concatenated_species_with_incomplete_lineage.csv and phylum_key.csv \n Moving on to next step. \U0002705 \n')
} else {
  # Merge the lineage and bracken outputs.
  cat_species <- read_csv(paste(args[4], "concatenated_Species_bracken.csv", sep = ""),
                          col_types = cols(), na = c("NA", "NULL"), guess_max = 1e6)
  cat_species$run_accession <- as.factor(cat_species$run_accession)
  
  lineage <- read_csv(paste(args[4], "cleaned_rankedlineage.csv", sep = ""),
                      col_types = cols(), na = c("NA", "NULL"),
                      guess_max = 1e6)
  
  # Change column name taxonomy_id to tax_id so it can be merged properly.
  colnames(cat_species)[2] <- "tax_id"
  lineage_samples <- merge(x = cat_species, 
                           y = lineage, 
                           by = "tax_id", 
                           all.x = TRUE) %>%
    dplyr::arrange(desc(run_accession))
  
  # Fill in the genera for species that are still missing genera information.
  species_missing <- lineage_samples %>% 
    dplyr::filter(is.na(genus) | is.na(phylum)) %>% 
    dplyr::filter(superkingdom != "Viruses" | is.na(superkingdom))
  
  # Make a data frame with the genus info for these species.
  merge_genus <- read_csv(paste(args[3], "missing_taxa2.csv", sep = ""))
  
  species_missing2 <- species_missing %>% dplyr::select(-c(genus, superkingdom))
  filled_in_species <- merge(species_missing2, merge_genus, by = "tax_id")
  
  # Replace just these rows in the main dataframe.
  lineage_samples2 <- lineage_samples %>% 
    dplyr::filter(!tax_id %in% merge_genus$tax_id)
  
  final_lineage_samples <- rbind(lineage_samples2, filled_in_species)
  
  # Check to make sure there are no NAs in the genus column besides viruses.
  if (sum(is.na((final_lineage_samples %>% dplyr::filter(!superkingdom == "Viruses"))$genus)) > 1) {
    stop("Error: There are tax_ids missing taxonomic information. See missing_taxa2.csv")
  }
  # Check to make sure it merged correctly.
  # cat_species should have the same number of rows as lineage_samples
  if (nrow(cat_species) != nrow(final_lineage_samples)) {
    stop("Error: cleaned_rankedlineage.csv did not merge correctly with concatenated_Species_bracken.csv")
  } else {
    # Lineage_samples now contains all of the bracken data from concatenated_Species_bracken.csv plus the added taxonomy provided by the NCBI rankedlineage.dmp file.
    # note: there is quite a lot of missing info, especially in regards kingdoms and superkingdoms.
    write_csv(final_lineage_samples, paste(args[4], "concatenated_species_with_incomplete_lineage.csv", sep = ""))
    rlang::inform('File written: concatenated_species_with_incomplete_lineage.csv \n')
  }
  # Get a list of the genera that are missing phylum info.
  # Generate phylum_key.csv:
  no_phyla <- final_lineage_samples %>% 
    dplyr::filter(is.na(phylum)) %>% 
    dplyr::filter(superkingdom != "Viruses" | is.na(superkingdom))
  curate_these <- no_phyla %>% 
    dplyr::count(genus) %>% 
    dplyr::filter(!is.na(genus))
  genera <- curate_these %>% 
    dplyr::pull(genus) # these are the genera that need to be manually curated
  curated_taxa <- data.frame("genus" = c("Allopseudarcicella", "Azoarcus", "Bacillus","Bartonella", "Bradyrhizobium", "Burkholderia", "Campylobacter", "Candidatus Babela", "Candidatus Vampirococcus", "Cohnella", "Cryptomonas", "Desulfovibrio", "Duncaniella", "Frankia", "Guillardia", "Hemiselmis", "Lysobacter", "Paenibacillus", "Serratia", "Streptomyces", "Thalassococcus", "Thermobaculum", "Thielavia", "Unclassified Bacteria", "Uncultured Bacteria", "Vampirococcus", "Xanthomonas"),
                             "phylum" = c("Bacteroidetes", "Proteobacteria", "Firmicutes", "Proteobacteria", "Proteobacteria", "Proteobacteria", "Proteobacteria", "Bacteria Candidate Phyla", "Candidatus Omnitrophica", "Firmicutes", "Cryptophyta", "Proteobacteria", "Bacteroidetes", "Actinobacteria", "Cryptophyta", "Cryptophyta", "Proteobacteria", "Firmicutes", "Proteobacteria", "Actinobacteria", "Proteobacteria", "Unclassified Terrabacteria Group", "Ascomycota", "Unclassified Bacteria", "Uncultured Bacteria", "Bacteria Candidate Phyla", "Proteobacteria"))
  
  check_curated_list <- curated_taxa$genus
  if (nrow(curate_these %>% dplyr::filter(!genus %in% check_curated_list)) == 0) {
    write.csv(curated_taxa, paste(args[4], "phylum_key.csv", sep = ""), row.names = F)
    rlang::inform('\U0001F31F 3. Completed! File written: phylum_key.csv \U0001F31F \n')
  } else {
    print("There are phyla in the data set that have not been added to the key")
    print("These are missing from the dictionary:")
    print(curate_these %>% dplyr::filter(!genus %in% check_curated_list))
    stop("Manual curation was not successful!", call.=FALSE)
  }
}

### 4. Fill in missing phylum information ----
# This function takes the phylum_key.csv generated in the last step and merges that information into the dataframe that is still missing phylum level taxonomic information. 
# Input: 
# 1. concatenated_species_with_incomplete_lineage.csv 
# 2. phylum_key.csv
# Output:
# 1. concatenated_species_with_complete_lineage.csv - all species (except viral IDs) contain taxonomic info at both the species and phylum level. 

rlang::inform('\U0001F99F 4. Starting to fill in missing phylum info. \U0001F99F \n')

# Check to see if args[4]/concatenated_species_with_complete_lineage.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "concatenated_species_with_complete_lineage.csv", sep = ""))) {
  rlang::inform('4. File exists: concatenated_species_with_complete_lineage.csv \n Moving on to next step. \U0002705 \n')
} else {
  curated_taxa <- read_csv(paste(args[4], "phylum_key.csv", sep = ""), 
                           col_names = T, col_types = cols())
  genera <- unique(curated_taxa$genus)
  
  lineage_samples <- read_csv(paste(args[4], "concatenated_species_with_incomplete_lineage.csv", sep = ""), 
                              col_names = T, col_types = cols())
  # These rows have the phyla filled in
  complete <- lineage_samples %>% 
    dplyr::filter(!genus %in% genera) 
  # Incomplete also filters out entries with no genus info
  incomplete <- lineage_samples %>% 
    dplyr::filter(genus %in% genera) %>% 
    dplyr::filter(!is.na(genus)) # These need to have their phyla filled in
  # Drop the phylum column so it can be repopulated
  incomplete <- subset(incomplete, select = -c(phylum))
  post_curation <- merge(x = incomplete, 
                         y = curated_taxa, 
                         by = "genus", 
                         all.x = TRUE)
  
  # Merge the manually curated samples back with the completed rows
  master <- bind_rows(complete, post_curation)
  
  # Check to make sure the only rows with NA phyla are viral.
  # Doesn't check for entries that have no genus, because no phylum can be easily added for them.
  if ((nrow(master %>% dplyr::filter(is.na(phylum)) %>% dplyr::filter(!is.na(genus)) %>% dplyr::count(superkingdom))) != 1) {
    stop("There are still genera that need manual curation.", call.=FALSE)
  }
  # Check to make sure that the number of rows pre- and post-curation are equivalent.
  if (nrow(master) != nrow(lineage_samples)) {
    stop("There is a problem with the manual curation.", call.=FALSE)
  } else {
    write_csv(master, paste(args[4], "concatenated_species_with_complete_lineage.csv", sep = ""))
    rlang::inform('\U0001F31F 4. Completed! File written: concatenated_species_with_complete_lineage.csv \U0001F31F \n')
  }
}

### 5. Merge NCBI informal ranks ----
# Purpose: NCBI does not differentiate between eukaryotes (inverts v. protozoans v. fungi) which is important for us to distinguish, also sometimes kingdom level info is missing. This is determined by which NCBI genomes refseq directory the sequences are stored. This function makes a list of all the unique phyla identified and then manually assigns one of the informal ranks - Bacteria, Fungi, Invertebrate, Protozoa, or Archaea.
# Input: 
# 1. concatenated_species_with_complete_lineage.csv
# 2. na_superkingdoms (generated manually)
# Output: 
# 1. informal_rank_key.csv
# 2. cleaned_informal_rank_with_lineage.csv - all samples will now have a column that listed its NCBI informal rank, and there are checks to make sure there is no

rlang::inform('\U0001F99F 5. Starting to add NCBI informal ranks. \U0001F99F \n')

# Check to see if args[4]/cleaned_informal_rank_with_lineage.csv and args[4]/informal_rank_key.csv exist.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "cleaned_informal_rank_with_lineage.csv", sep = "")) & file.exists(paste(args[4], "informal_rank_key.csv", sep = ""))) {
  rlang::inform('5. File exists: cleaned_informal_rank_with_lineage.csv \n Moving on to next step. \U0002705 \n')
} else {
  master <- read_csv(paste(args[4], "concatenated_species_with_complete_lineage.csv", sep = ""), 
                     col_names = T, col_types = cols())
  informal_viral <- master %>% 
    dplyr::filter(superkingdom == "Viruses")
  informal_master <- master %>% 
    dplyr::filter(superkingdom != "Viruses")
  
  if (nrow(informal_viral) + nrow(informal_master) != nrow(master)) {
    stop("There is an error in the superkingdom column, likely needs to be factored.", call. = FALSE)
  } else {
    # This is how the key is built:
    informal_master$phylum <- as.factor(as.character(informal_master$phylum))
    rank_these <- informal_master %>% dplyr::count(phylum)
    
    informal_key <- data.frame("phylum" = c("Proteobacteria", "Dictyoglomi", "Bacteroidetes", "Planctomycetes", "Spirochaetes", "Nitrospirae", "Deinococcus-Thermus", "Chlamydiae", "Firmicutes", "Fusobacteria", "Chlorobi", "Chloroflexi", "Cyanobacteria", "Actinobacteria", "Thermodesulfobacteria", "Tenericutes", "Euryarchaeota", "Crenarchaeota", "Thermotogae", "Verrucomicrobia", "Bacillariophyta", "Ascomycota", "Basidiomycota", "Apicomplexa", "Microsporidia", "Platyhelminthes", "Nematoda", "Mollusca", "Arthropoda", "Chordata", "Acidobacteria", "Deferribacteres", "Aquificae", "Coprothermobacterota", "Synergistetes", "Gemmatimonadetes", "Calditrichaeota", "Thaumarchaeota", "Chrysiogenetes", "Elusimicrobia", "Candidatus Cloacimonetes", "Ignavibacteriae", "Armatimonadetes", "Kiritimatiellaeota", "Balneolaeota", "Candidatus Korarchaeota", "Fibrobacteres", "Caldiserica", "Candidatus Bipolaricaulota", "Candidatus Micrarchaeota", "Cercozoan", "Bacteria Candidate Phyla", "Cryptophyta", "Amoebozoa", "Sarcomastigophora", "Ciliophora", "Candidatus Saccharibacteria", "Candidatus Gracilibacteria", "Lentisphaerae", "candidate division Zixibacteria", "Cercozoa", "Euglenozoa", "Evosea", "Unclassified Bacteria", "Unclassified Terrabacteria Group", "Uncultured Bacteria", "Candidatus Omnitrophica"),
                               "informal_rank" = c("Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Archaea", "Archaea", "Bacteria", "Bacteria", "Protozoa", "Fungi", "Fungi", "Protozoa", "Protozoa", "Invertebrate", "Invertebrate", "Invertebrate", "Invertebrate", "Invertebrate", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Archaea", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Archaea", "Bacteria", "Bacteria", "Bacteria", "Archaea", "Protozoa", "Bacteria", "Protozoa", "Protozoa", "Protozoa", "Protozoa", "Bacteria", "Bacteria", "Bacteria", "Bacteria", "Protozoa", "Protozoa", "Protozoa", "Bacteria", "Bacteria", "Bacteria", "Bacteria"))
    check_curated_ncbi <- informal_key$phylum
    
    # Again, not adding phyla for entries with no genus, so some will have NA for phylum
    if (nrow(rank_these %>% dplyr::filter(!phylum %in% check_curated_ncbi) %>% dplyr::filter(!is.na(phylum))) != 0) {
      print(rank_these %>% dplyr::filter(!phylum %in% check_curated_ncbi))
      stop("There is a problem with the manual curation.", call.=FALSE)
    } else {
      write_csv(informal_key, paste(args[4], "informal_rank_key.csv", sep = ""))
      print("informal_rank_key.csv written.")
      # Merge the key with the master dataframe without viruses
      informal_complete <- merge(x = informal_master, y = informal_key, by = "phylum", all.x = TRUE)
      # Add a column in viruses for their informal rank
      informal_viral$informal_rank <- "Viruses"
      # Bind them back together for the master
      rank_master <- NULL
      rank_master <- bind_rows(informal_complete, informal_viral) %>% arrange(desc(run_accession))
      rank_master <- rank_master %>% dplyr::select("run_accession", "tax_id", "name", "taxonomy_lvl", "kraken_assigned_reads", "added_reads", "new_est_reads", "fraction_total_reads", "tax_name", "genus", "family", "order", "class", "phylum", "informal_rank", "kingdom", "superkingdom")
      if (nrow(rank_master) != nrow(master)) {
        stop("There is a problem with adding the NCBI informal ranks.", call.=FALSE)
      } else {
        write_csv(rank_master, paste(args[4], "cleaned_informal_rank_with_lineage.csv", sep = ""), col_names = T)
        rlang::inform('\U0001F31F 5. Completed! File written: cleaned_informal_rank_with_lineage.csv \U0001F31F \n')
      }
    }
  }
}

### 6.  Merge kreports for the accessions included in analyses ----
# Purpose: The Bracken files do not contain information about the number of reads that were unable to be classified. However, the Kraken reports do, so those need to be cleaned up and merged with the master dataframe.
# Input: 
# 1. path to the directory storing all of the kreports
# 2. cleaned_informal_rank_with_lineage.csv
# Output: 
# 1. bracken-kreport.csv - master dataframe including rows with the number of unclassified reads for each sample

rlang::inform('\U0001F99F 6. Starting to merge kreports for the accessions included in analyses. \U0001F99F \n')

# Check to see if args[4]/bracken-kreport.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "bracken-kreport.csv", sep = ""))) {
  rlang::inform('6. File exists: bracken-kreport.csv \n Moving on to next step. \U0002705 \n')
} else{
  # Make the list of files in the directory a vector.
  files <- list.files(path = args[2], pattern = "SRR[0-9]*.kreport2", full.names = FALSE, recursive = FALSE)
  acc_num <- gsub("\\.kreport2*$", "", files) # Keep only the accession numbers
  ### Output: table with all samples and all identified genera
  tmp <- NULL
  for (acc in acc_num) {
    raw <- read_tsv(paste(args[2], acc, ".kreport2", sep = ""), col_names = F, col_types = cols())
    colnames(raw) <- c("percent_covered", "number_of_reads", "num_assigned_to_taxa", "rank_code", "tax_id", "name")
    raw$run_accession <- acc  # Add column for accession number
    name <- paste(acc, "_kreport", sep = "")
    # assign(name, raw) uncomment if you want indiv. df for each accession num
    tmp <- rbind(tmp, raw)
  }
  # Keep only rows either classified at the species level or unclassified
  cat_kreport <- tmp %>% 
    dplyr::filter(rank_code == "S" | rank_code == "U")
  
  rank_master <- read_csv(paste(args[4], "cleaned_informal_rank_with_lineage.csv", sep = ""), 
                          col_names = T, col_types = cols())
  # Merge the bracken_kreport df with the master df
  master <- merge(x = rank_master, y = cat_kreport, by = c("run_accession", "name"), all = T) %>%
    arrange(desc(run_accession))
  master[master == "NA"] <- NA
  
  # Keep tax_id.y not .x because it includes 0 for unclassified reads
  # Rename columns to be more descriptive
  colnames(master) <- c("run_accession", "species", "tax_id.x", "taxonomy_lvl", "kraken_assigned_reads", "bracken_added_reads", "new_est_reads", "fraction_classified_reads", "tax_name", "genus", "family", "order", "class", "phylum", "ncbi_informal_rank", "kingdom", "superkingdom", "fraction_total_reads", "bracken_assigned_reads", "reads_assigned_to_taxa", "rank_code", "tax_id")
  master <- master %>% 
    dplyr::select("run_accession", "species", "tax_id", "kraken_assigned_reads", "bracken_added_reads", "new_est_reads", "fraction_classified_reads", "bracken_assigned_reads", "fraction_total_reads", "reads_assigned_to_taxa", "taxonomy_lvl", "rank_code", "tax_name", "genus", "family", "order", "class", "phylum", "kingdom", "superkingdom","ncbi_informal_rank")
  # Change the NAs from the unclassified reads to "unclassified"
  master$ncbi_informal_rank[master$rank_code == "U"] <- "Unclassified"
  
  if (nrow(master) != nrow(rank_master) + length(unique(rank_master$run_accession))) {
    stop("Error: kreports did not merge correctly.")
  } else {
    write_csv(master, paste(args[4], "bracken-kreport.csv", sep = ""))
    rlang::inform('\U0001F31F 6. Completed! File written: bracken-kreport.csv \U0001F31F \n')
  }
}

### 7. Merge sample metadata from NCBI ----
# PRJNA412140_metadata.txt contains the metadata about each sample, sequencing depth, collection site, sex, etc. This will merge the metadata with the cleaned kraken/bracken output based of run accession number.
# Input: 
# 1. bracken-kreport.csv
# 2. PRJNA412140_metadata.txt
# Output: 
# 1. master_added_metadata_species.csv - master dataframe including sample metadata from NCBI.

rlang::inform('\U0001F99F 7. Starting to merge sample metadata from NCBI. \U0001F99F \n')

# Check to see if args[4]/master_added_metadata_species.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "master_added_metadata_species.csv", sep = ""))) {
  rlang::inform('7. File exists: master_added_metadata_species.csv \n Moving on to next step. \U0002705 \n')
} else {
  # Make a list of all of the accession numbers that are included in the study.
  dat <- readr::read_csv(file = paste(args[4], "bracken-kreport.csv", sep = ""), col_names = T, col_types = cols())
  acc_num <- unique(dat$run_accession)
  
  # Take the metadata from NCBI and keep only those rows with accessions in the kraken/bracken dataframe.
  meta_ncbi <- readr::read_csv(file = paste(args[3], "PRJNA412140_metadata.txt", sep = ""), col_names = T, col_types = cols()) # Import the NCBI metadata
  meta_ncbi <- meta_ncbi %>% dplyr::select(c("Run", "Bases", "BioProject",
                                             "BioSample", "Bytes", "Experiment",
                                             "lat_lon", "Sample Name")) # Select the relevant columns.
  colnames(meta_ncbi)[colnames(meta_ncbi) == "Run"] <- "run_accession" # Change the column names for consistency when merging.
  colnames(meta_ncbi)[colnames(meta_ncbi) == "SRA_Sample"] <- "sra_sample_accession"
  meta_ncbi <- meta_ncbi %>% 
    dplyr::filter(run_accession %in% acc_num) %>% 
    dplyr::arrange(desc(run_accession)) # Keep only accessions that are applicable for our subset data.
  
  master_meta <- merge(x = dat, y = meta_ncbi, by = "run_accession") 
  
  # Check to make sure that every run accession you have in your data has metadata available for it
  if (isTRUE(all.equal(acc_num, meta_ncbi[["run_accession"]])) != TRUE) {
    stop("There is an accession number in the experimental data that does not have corressponding metadata from NCBI.
         Try searching NCBI for the BioProject that all of these are stored under and redownloading the summary.txt file from the run selector", call.=FALSE)
  } else {
    write_csv(master_meta, paste(args[4], "master_added_metadata_species.csv", sep = ""), col_names = T)
    rlang::inform('\U0001F31F 7. Completed! File written: master_added_metadata_species.csv \U0001F31F \n')
  }
}

### 8. Merge better viral taxonomy (not really necessary) ----
# NCBI has incomplete taxonomy for most viral species, a more complete viral taxonomy is available through the ICTV website. (see ictv_bettervirustax.csv). This will merge the ICTV viral taxonomy with the data that has been cleaned up to this point. 
# Note: this step actually isn't necessary since we don't end up analyzing viral species for a number of reasons... but removing it breaks some of the downstream code so I've kept it as a step. It also does not affect any of the other data, so there is no issues in including this step. 
# Input: 
# 1. master_added_metadata_species.csv
# 2. ictv_bettervirustax.csv
# Output: 
# 1. better_metadata_species.csv

rlang::inform('\U0001F99F *. Starting to merge better viral taxonomy from ICTV. \U0001F99F \n')

# Check to see if args[4]/better_metadata_species.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "better_metadata_species.csv", sep = ""))) {
  rlang::inform('8. File exists: better_metadata_species.csv \n Moving on to next step. \U0002705 \n')
} else {
  original_data <- read_csv(file = paste(args[4], "master_added_metadata_species.csv", sep = ""),
                            col_names = T, col_types = cols())
  # Make a separate data frame for just viral IDs and one for everything else
  non_viruses <- original_data %>% 
    dplyr::filter(ncbi_informal_rank != "Viruses")
  viruses <- original_data %>% 
    dplyr::filter(ncbi_informal_rank == "Viruses") %>%
    dplyr::select(-family, -order, -class, -phylum)
  
  # Check to make sure the sum of the rows in the two data frames equals the row number from the original file.
  if (nrow(non_viruses) + nrow(viruses) != nrow(original_data)) {
    stop("Error: Dataframe is not filtering correctly. See lines prior to this error message.")
  } else {
    new_tax <- read_csv(file = paste(args[3], "ictv_bettervirustax.csv", sep = ""), 
                        col_names = T, col_types = cols()) %>% 
      dplyr::select(Phylum, Class, Order, Family, Genus) %>%
      dplyr::distinct() %>%
      dplyr::filter(!is.na(Genus))
    colnames(new_tax) <- c("phylum", "class", "order", "family", "genus")
    
    # Merge the viral taxonomy
    new_tax2 <- merge(x = viruses, y = new_tax, by = "genus", all.y = FALSE, all.x = T)
    
    # Combine the two split data frames back together.
    final_tax <- rbind(non_viruses, new_tax2)
    
    final_tax2 <<- transform(final_tax, superkingdom = ifelse(ncbi_informal_rank == "Unclassified", "Unclassified", superkingdom))
    
    # Check to make sure the output row number is equal to the number of rows in the input file (master_added_metadata_species.csv)
    if (nrow(final_tax2) != nrow(original_data)) {
      stop("Error! There was an issue with merging viral taxonomy. Output row number is not equal to input row number.")
    } else {
      write_csv(final_tax2, paste(args[4], "better_metadata_species.csv", sep = ""))
      rlang::inform('\U0001F31F 8. Completed! File written: better_metadata_species.csv \U0001F31F \n')
    }
  }
}

### 9. Merge restriction enzyme count information ----
# Because this sequencing data was collected using ddRADseq, we need to account for the number of restriction cut sites in each of the species that is present in our kraken database. 

# Input: 
# 1. better_metadata_species.csv
# 2. asc_strain_filtered.csv
# Output: 
# 1. added_restriction_site.csv

rlang::inform('\U0001F99F 9. Starting to merge restriction enzyme count information. \U0001F99F \n')

# Check to see if args[4]/added_restriction_site.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "added_restriction_site.csv", sep = ""))) {
  rlang::inform('9. File exists: added_restriction_site.csv \n Moving on to next step. \U0002705 \n')
} else {
  # Import cleaned kraken/bracken data.
  dat <- read_csv(paste(args[4], "better_metadata_species.csv", sep = ""), 
                  col_names = T, col_types = cols())
  # Import site count information generated on the VM -- Henry
  sites <- read_csv(paste(args[3], "asc_strain_filtered.csv", sep = ""), 
                    col_names = T, col_types = cols()) %>% 
    dplyr::select(-c(species_name)) %>% 
    dplyr::group_by(tax_id) %>%
    dplyr::slice(1) %>%
    dplyr::ungroup()
  
  # Merge the two data frames based on NCBI TaxID number
  merged_sites <- merge(x = dat, y = sites, by = "tax_id", all.x = T)
  
  if (nrow(dat) != nrow(merged_sites)) {
    stop("Error: There was an error merging in the restriction site counts.")
  }
  
  unclassified <- merged_sites %>% 
    dplyr::filter(ncbi_informal_rank == "Unclassified")
  classified <- merged_sites %>% 
    dplyr::filter(ncbi_informal_rank != "Unclassified")
  
  remove_these <- classified %>% 
    dplyr::filter(is.na(total_sites)) %>% 
    dplyr::group_by(tax_id) %>% 
    dplyr::arrange(new_est_reads) %>% 
    dplyr::slice(1) %>%
    dplyr::pull(tax_id)
  
  classified_filter <- classified %>% 
    dplyr::filter(!tax_id %in% remove_these)
  
  total_filtered <<- rbind(unclassified, classified_filter)
  
  if (sum(is.na((merged_sites %>% dplyr::filter(ncbi_informal_rank == "Unclassified"))$total_sites)) != length(unique(merged_sites$run_accession))) {
    stop("Error: There was an error merging in the restriction site counts.")
  } else {
    write_csv(total_filtered, paste(args[4], "added_restriction_site.csv", sep = ""))
    rlang::inform('\U0001F31F 9. Completed! File written: added_restriction_site.csv \U0001F31F \n')
  }
}

### 10. Calculate percent composition of the microbiome ----
# Purpose: Summarize some basic stats about the composition of each sample -
# ex. how much of the sample is host DNA, how much is bacteria, etc.
# While running the function it will ask you some questions about stats you may want to include, answer yes to both to get the most amount of data.
# Input: 
# 1. added_restriction_site.csv
# Output: 
# 1. master_merged_summ_species.csv - file with all of the summary stats

rlang::inform('\U0001F99F 10. Starting to calculate percent composition of the microbiome. \U0001F99F \n')

# Check to see if args[4]/master_merged_summ_species.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "master_merged_summ_species.csv", sep = ""))) {
  rlang::inform('10. File exists: master_merged_summ_species.csv \n Moving on to next step. \U0002705 \n')
} else {
  dat <- read_csv(file = paste(args[4], "added_restriction_site.csv", sep = ""), 
                  col_names = T, col_types = cols(), guess_max = 1e6)
  dat[dat=="Virus"] <- "Viruses"
  
  # Filling in the number of unclassified reads.
  dat2 <- dat %>%
    dplyr::mutate(new_est_reads = ifelse(ncbi_informal_rank == "Unclassified", dat$bracken_assigned_reads, new_est_reads)) %>%
    dplyr::select(run_accession, ncbi_informal_rank, new_est_reads)
  
  # Calculate the total number of reads for each rank for each accession number.
  sample_summary <- dat2 %>% 
    dplyr::group_by(ncbi_informal_rank, run_accession) %>%
    dplyr::summarise(total = sum(new_est_reads)) %>%
    dplyr::ungroup()
  
  # Calculate the total number of reads in each sample
  total_read_count <- sample_summary %>% 
    dplyr::group_by(run_accession) %>% 
    dplyr::summarise(total_read_count = sum(total, na.rm = T)) %>%
    dplyr::ungroup()
  
  # Reformat the data so each taxonomic group is a column
  sample_summary <- sample_summary %>% 
    tidyr::spread(key = ncbi_informal_rank, value = total) #%>% select(-c(`<NA>`))
  # map kingdom NA counts to 0
  sample_summary[is.na(sample_summary)] <- 0
  
  # calculate the total number of classified reads
  sample_summary <- sample_summary %>%
    dplyr::mutate(total_classified_reads = Archaea + Bacteria + Fungi + Invertebrate + Protozoa + Viruses)
  
  # merge the two working dataframes together
  merged_summ <- merge(x = sample_summary, y = total_read_count, by = "run_accession")
  colnames(merged_summ) <- c("run_accession", "total_archaeal_reads", 
                             "total_bacterial_reads", "total_fungal_reads", 
                             "total_invert_reads", "total_protozoan_reads", 
                             "total_unclassified_reads", "total_viral_reads", 
                             "total_classified_reads", "total_read_count")
  
  # Calculate the proportion of reads classified and unclassified
  # Need to normalize the classified reads to the single sample with the greatest number of total reads.
  # Use factor_lvl to multiply all of the assigned reads by to normalize them to the sample with the greatest number of reads.
  merged_summ <- merged_summ %>%
    dplyr::mutate(percent_classified = total_classified_reads / total_read_count,
                  percent_unclassified = total_unclassified_reads / total_read_count, 
                  max_total_reads = max(total_read_count),
                  factor_lvl = max_total_reads / total_read_count)
  
  for (name in c("archaeal", "bacterial", "fungal", "invert", "protozoan", "unclassified", "viral", "classified")) {
    new_name <- paste("norm", name, "reads", sep = "_")
    current_name <- paste("total", name, "reads", sep = "_")
    merged_summ[,new_name] <- merged_summ[,current_name] * merged_summ$factor_lvl
  }
  
  # Calculate some more useful stats
  merged_summ <- merged_summ %>%
    dplyr::mutate(norm_microbiome_read_count = norm_archaeal_reads + norm_bacterial_reads + norm_fungal_reads + norm_protozoan_reads + norm_viral_reads,
                  archaea_in_microbiome = norm_archaeal_reads / norm_microbiome_read_count,
                  bacteria_in_microbiome = norm_bacterial_reads / norm_microbiome_read_count, 
                  fungi_in_microbiome =  norm_fungal_reads / norm_microbiome_read_count,
                  protozoa_in_microbiome = norm_protozoan_reads / norm_microbiome_read_count, 
                  virus_in_microbiome = norm_viral_reads / norm_microbiome_read_count)
  
  write_csv(merged_summ, paste(args[4], "master_merged_summ_species.csv", sep = ""))
  rlang::inform('\U0001F31F 10. Completed! File written: master_merged_summ_species.csv \U0001F31F \n')
}

### 11. Remove rare species and normalize just bacterial read counts ----
# Purpose: Filter for just bacterial species, normalize based on both the total number of reads per sample and the number of restriction sites. Remove any species that does not occur in at least one sample at > 0.01% of the total bacterial microbiome.
# Input: 
# 1. master_merged_summ_species.csv
# 2. added_restriction_site.csv
# Output: 
# 1. bacteria_rare_removed.csv

rlang::inform('\U0001F99F 11. Starting to remove rare species and normalize. \U0001F99F \n')

# Check to see if args[4]/bacteria_rare_removed.csv exists.
# If it does, this step will be skipped.
if (file.exists(paste(args[4], "bacteria_rare_removed.csv", sep = ""))) {
  rlang::inform('11. File exists: bacteria_rare_removed.csv \n Moving on to next step. \U0002705 \n')
} else {
  merged_summ <- read_csv(file = paste(args[4], "master_merged_summ_species.csv", sep = ""),
                          col_names = T, col_types = cols())
  total_filtered <- read_csv(file = paste(args[4], "added_restriction_site.csv", sep = ""),
                             col_names = T, col_types = cols(), guess_max = 1e6)
  
  bacteria <- merge(x = total_filtered, y = merged_summ, by = "run_accession") %>% 
    dplyr::filter(ncbi_informal_rank == "Bacteria") %>%
    dplyr::mutate(norm_new_est_reads = new_est_reads * factor_lvl,
                  prop_bacteria = (norm_new_est_reads / norm_bacterial_reads)*100)
  
  check_max_amount <- bacteria %>% 
    dplyr::group_by(species) %>% 
    dplyr::summarise(max_amount = max(prop_bacteria)) %>% 
    dplyr::ungroup() %>% 
    dplyr::filter(max_amount > .1) %>%
    dplyr::pull(species)
  
  rare_removed <- bacteria %>% 
    dplyr::filter(species %in% check_max_amount) %>%
    dplyr::mutate(max_sites = max(total_sites),
                  site_factor = max_sites / total_sites,
                  double_norm_reads = norm_new_est_reads * site_factor)
  
  write_csv(rare_removed, paste(args[4], "bacteria_rare_removed.csv", sep = ""))
  rlang::inform('\U0001F31F 11. Completed! File written: bacteria_rare_removed.csv \U0001F31F \n')
}
