#!/usr/bin/env Rscript

#File Name    : dmn_plotmaker.R 
#Author       : Alaric D'Souza, alaric.dsouza@wustl.edu
#Created On   : Fri Oct  4 14:11:15 CDT 2019
#Last Modified: Fri Oct 18 14:29:19 CDT 2019
#Description  : Outputs a data table for AIC values of DMN models and makes aggregated AIC and Laplace plots.
#Usage: dmn_plotmaker.R -i <inputdirectory (contains RDS files of DMN models)> -o <output directory>

#install and load required package
library(plyr)
library(ggplot2)
library(DirichletMultinomial)
library(optparse)

#take in arguments
option_list = list(
  make_option(c("-i", "--inputdirectory"), action="store", default=NA, type='character',
              help="directory of DMN models"),
  make_option(c("-o", "--outputdirectory"), action="store", default=NA, type='character',
              help="output directory")
  )
#parse arguments to list
opt = parse_args(OptionParser(option_list=option_list))

# main point of program is here, do this whether or not "verbose" is set
if(!is.na(opt$inputdirectory) & !is.na(opt$outputdirectory)) {
#create output directory
ifelse(!dir.exists(opt$outputdirectory), dir.create(opt$outputdirectory), FALSE)
#dir.create(opt$outputdirectory)

#list RDS files in the directory
filelist<-list.files(path=opt$inputdirectory,pattern="*.rds")

#store AIC values in a list of lists
dmnList<-
setNames(
lapply(
    filelist,
    function(datafile){
        ldply(
            lapply(
                readRDS(file.path(opt$inputdirectory,datafile)),
                function(myob){
                    c(myob@goodnessOfFit[3],myob@goodnessOfFit[5])
                    })
        )
}
),
gsub("DMM_\\d+k_(\\d+).rds","\\1",filelist)
)

#make data frame from list of AIC values
df<-ldply(
        lapply(dmnList,function(df){
            df$k<-as.integer(row.names(df)); return(df)
            }),
        .id="seed")

#write out AIC and Laplace table
write.table(
    df,
    file=file.path(opt$outputdirectory, "dmn_AIC_Laplace.txt"),
    sep="\t",
    row.names=F,
    quote=F)


#save function for custom theme
customTheme <- function(myplot) {
  return(
    myplot +
      theme(
        panel.background = element_blank(),
        panel.border = element_rect(
          fill = NA,
          color = "black",
          linetype = 1,
          size = 1
        ),
        panel.grid = element_blank(),
        axis.text = element_text(color = "black")
      )
  )
}

#save AIC plot with points and boxplots by cluster
AIC_plot<-
    customTheme(
        ggplot(
          df,
          aes(x=k,y=AIC)
        )+
          geom_boxplot(aes(x=as.factor(k)),color="red",outlier.shape = NA,fill=NA)+
          #geom_point(data=aggregate(AIC~k,df,FUN = mean),color="red",shape=15)+
          geom_point(alpha=0.1,color="black")
    )
#save plot
ggsave(filename=file.path(opt$outputdirectory, "dmn_AIC_plot.pdf"),plot=AIC_plot,width=11,height=8.5)

#save Laplace plot with points and boxplots by cluster
Laplace_plot<-
    customTheme(
        ggplot(
          df,
          aes(x=k,y=Laplace)
        )+
          geom_boxplot(aes(x=as.factor(k)),color="red",outlier.shape = NA,fill=NA)+
          #geom_point(data=aggregate(Laplace~k,df,FUN = mean),color="red",shape=15)+
          geom_point(alpha=0.1,color="black")
    )
#save plot
ggsave(filename=file.path(opt$outputdirectory, "dmn_Laplace_plot.pdf"),plot=Laplace_plot,width=11,height=8.5)

} else {
    cat("you didn't specify both input file and output directory\n", file=stderr()) # print error messages to stderr
}