#!/usr/bin/env Rscript

#File Name    : bestDirichletModelParse.R
#Author       : Alaric D'Souza, alaric.dsouza@wustl.edu
#Created On   : Thu Oct  3 21:43:15 CDT 2019
#Last Modified: Thu Oct  3 21:43:21 CDT 2019
#Description  : Takes in community data matrix and directory with Dirichlet Multinomial model objects and outputs data from the best fit model.
#Usage: bestDirichletModelParse.R -i <inputfile (community data matrix)> -r <directory where .rds files are stored> -o <output directory> -m <model evaluation dataframe>

#install and load required package
library("optparse")
library("DirichletMultinomial")
library("lattice")

#take in arguments
option_list = list(
  make_option(c("-i", "--inputfile"), action="store", default=NA, type='character',
              help="community matrix file"),
  make_option(c("-r", "--rdsdirectory"), action="store", default=NA, type='character',
              help="directory with model rds files"),
  make_option(c("-k", "--clustvalue"), action="store", default="50", type='character',
              help="number of initial clusters set"),
  make_option(c("-o", "--outputdirectory"), action="store", default=NA, type='character',
              help="output directory"),
  make_option(c("-m", "--modeldf"), action="store", default=NA, type='character',
              help="dataframe of model values")
)
#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$inputfile) & !is.na(opt$outputdirectory) & !is.na(opt$rdsdirectory) & !is.na(opt$modeldf)) {
#create output directory
ifelse(!dir.exists(opt$outputdirectory), dir.create(opt$outputdirectory), FALSE)
#dir.create(opt$outputdirectory)

renormMatrix <- function(inmatrix){
    return(inmatrix/(min(inmatrix[inmatrix>0])))
    }

#read in community matrix
count=renormMatrix(as.matrix(read.table(file=opt$inputfile,sep="\t",row.names=1,header=T)))

#read in model evaluation file file
df <- read.table(
  file=opt$modeldf,
  sep = "\t",
  stringsAsFactors = F,
  header = T,
  quote = "",
  comment.char = "",
  colClasses  = c("character","numeric","numeric","integer")
)

#store best seed
bestSeed<-as.character(df$seed[which.min(df$Laplace)])
#store best k for best seed
bestK<-df$k[which.min(df$Laplace)]

#read in best seed model
best.dmn <- 
  readRDS(
    file.path(opt$rdsdirectory,
              paste0("DMM_",opt$clustvalue,"k_",bestSeed,".rds")
              )
    )

#save mixture data frame
mixtureDF <- mixture(best.dmn[[bestK]])
#rename columns for mixture DF
colnames(mixtureDF) <- paste0("m",1:bestK)

#write out the sample enterotype table
write.table(
    cbind(Sample=row.names(mixtureDF),mixtureDF),
    file=file.path(opt$outputdirectory, paste0("dmn_mixtureSamples_",bestK,".txt")),
    sep="\t",
    row.names=F,
    quote=F)


#store the fitted first model
firstModel <- fitted(best.dmn[[1]], scale=TRUE)
#store the fitted best model
bestModel <- fitted(best.dmn[[bestK]], scale=TRUE)

#set colnames of the best model
colnames(bestModel) <- paste0("m",1:bestK)
#save the meandifferences of the best model
meandiff <- colSums(abs(bestModel - as.vector(firstModel)))
#store the differences of the best model
diff <- rowSums(abs(bestModel - as.vector(firstModel)))
#store the order of the maximum differences
o <- order(diff, decreasing=TRUE)
#store the ratio of the differnces
cdiff <- cumsum(diff[o]) / sum(diff)

#save taxonomic contribution table
taxDF <- cbind(Mean=firstModel[o], bestModel[o,], diff=diff[o], cdiff)

#write out the taxonomic contribution table
write.table(
    cbind(taxonomy=row.names(taxDF),taxDF),
    file=file.path(opt$outputdirectory, paste0("dmn_taxonomicContribution_",bestK,".txt")),
    sep="\t",
    row.names=F,
    quote=F)

#save the PDF file for the samples
pdf(file.path(opt$outputdirectory,paste0("dmn_dirichletSamplesHeatmap_",bestK,".pdf")))
heatmapdmn(count, best.dmn[[1]], best.dmn[[bestK]], 30)
dev.off()

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