#!/usr/bin/env Rscript

library("optparse")
library("reshape2")
library("ggplot2")
library("cowplot")
library("ggtree")

option_list = list(
  make_option(c("-i","--input_dir"), type="character", default=NULL,
              help="Input dir", metavar="character"),
    make_option(c("-o","--out"), type="character", default="out",
              help="output prefix [default= %default]", metavar="character")
);

opt_parser = OptionParser(option_list=option_list);
opt = parse_args(opt_parser);

if (is.null(opt$input_dir)){
  print_help(opt_parser)
  stop("At least one argument must be supplied (input_dir)", call.=FALSE)
}

input_dir = opt$input_dir
## fun...


## program...

print(input_dir)

files = paste0(list.files(input_dir))

print(files)

files_split = strsplit(files, "@", fixed = T)
files_df = as.data.frame(do.call(rbind, files_split))
files_df_ok = data.frame(files= paste0(input_dir,"/",files), tree_prefix = files_df$V1, hyp = gsub(".nw","", files_df$V2), stringsAsFactors = FALSE)

print(files_df_ok)




plot_tree_prefix = function(x) {
    b_tree_plot = function(name) {
        if (name %in% x$hyp) {
            file = x$files[x$hyp == name]
            tree <- read.tree(file=file)
            plot = ggtree(tree) + geom_treescale() + geom_tiplab(align=F, size=3, linesize=.3)
            plot = plot + geom_text2(aes(x=branch, label=round(branch.length,3)), size = 3, vjust = -0.5)
            plot = plot + ggtitle(name)
            return(plot)
        } else {
            return(NULL)
        }
    }
  
  plotlist = lapply(list("input_tree", "H0",
                         "HaPC", "HaPCOC",
                         "H0_NeBig","HaPC_NeBig",
                         "H0_NeSmall","HaPC_NeSmall",
                         "H0_BigNeInSmallNe","HaPC_BigNeInSmallNe",
                         "H0_SmallNeInBigNe","HaPC_SmallNeInBigNe"
                    ), b_tree_plot)
  plotlist = plyr::compact(plotlist)
  plot = plot_grid(plotlist = plotlist, labels=names(plotlist), ncol = 2)
  
  tree_name = as.character(unique(x$tree_prefix))
  
  plot = plot_grid(NULL, plot, labels=c(NA,tree_name), label_size = 24,
                   label_x = 0, label_y = 1, hjust = - 0.1 , vjust = - 1,
                   ncol = 1, rel_heights = c(0.05,1))
  return(plot)
}

plotlist = lapply(split(files_df_ok,paste0(files_df_ok$tree_prefix)), plot_tree_prefix)


plot_all <- plot_grid(plotlist = plotlist,
                      labels=NA, nrow =  length(unique(files_df_ok$tree_prefix)))

output_pdf = paste0(opt$out,".pdf")
save_plot(output_pdf,
          plot_all,
          ncol = 2 ,
          nrow = 10 * length(unique(files_df_ok$tree_prefix)),
          base_aspect_ratio = 2,
          limitsize = FALSE
          )