#!/usr/bin/env Rscript

# This script generates dendrograms from StrainPhlAn outputs.

# load the optparse library
library(optparse)

# Create command line argument parser
help_description <- "

This script generates dendrograms from StrainPhlAn output files.

The following positional arguments are required:

RAxML_bestTree.species.tree (Input file): This tree is generated by StrainPhlAn
metadata.txt (Input file): This is the metadata file for the tree
species.fasta (Input file): This fasta file is generated by StrainPhlAn
output_tree_1.png (Output file): This is the first dendrogram image written
output_tree_2.png (Output file): This is the second dendrogram image written"

args <- OptionParser(usage = "%prog RAxML_bestTree.species.tree metadata.txt species.fasta output_tree_1.png output_tree_2.png", 
                      add_help_option = TRUE, prog='strainphlan_ggtree.R',
                      description=help_description )
args_list <- parse_args(args, positional_arguments=TRUE)

# load other libraries after optparse to not load them on help
library(ggplot2)
library(ggtree)
library(RColorBrewer)

# read newick tree
e.sir.tre <- read.tree( args_list$args[1] )
# create ggtree object
e.sir.gg <- ggtree( e.sir.tre )
# read in metadata file
e.sir.meta <- read.delim( args_list$args[2], header = T, sep = "\t" )
# add metadata to dendrogram plot
e.sir.gg <- e.sir.gg %<+% e.sir.meta

# plot with color terminal edges and sample names with Country and add tiplabels
png( args_list$args[4], width = 5, height = 6, res = 300, units = "in" )
# strainphlan_tree_1.pdf
e.sir.gg +  
  geom_tippoint( size = 3, aes( color = Country) ) +
  aes( branch.length = 'length' ) +
  theme_tree2() + theme(legend.position="right")
temp <- dev.off()
# visualize tree with multiple sequence alignment (MSA)

# path to alignment file
e.sir.fasta <- (args_list$args[3])
# plot tree with slice of MSA
png( args_list$args[5], width = 7, height = 6, res = 300, units = "in" )
# strainphlan_tree_2.pdf
msaplot( e.sir.gg + geom_tippoint( size = 3, aes( color = Country) ), 
         e.sir.fasta, window = c( 490,540 ), 
         color = brewer.pal(6, "Set3") ) + 
  theme( legend.position = 'right' )
temp <- dev.off()


