#########################################################
#                                                       #
# Script for running Zero inflated Poisson Log-Normal   #
#                                                       #
# By: Viktor Jonsson,   2018-03-17                      #
#                                                       #
#########################################################


#Initialize required packages
library("rjags")
library("coda")


#Source required functions

source("ZiGLMMfunctions.R")


#Prepare data

#LoadData from file


Y <- read.table("yatsunenko2012.txt",h=T,sep="\t",quote="",row.names = 1) 
Y <- Y[1:200,1:10] #Note this is just an example that uses a small part
# of the Yatsunenko 2012 dataset to get a rapidly runnable test case.

groups <- c(5,5) #Indicates the size of each group samples


group=c(rep(0,groups[1]),rep(1,groups[2])) 
# indicator vector indicating which samples
# belong to which group. Note this implementation requires
# the two groups to be in consecutive order

#add pseudo counts to any gene with zeros in one group
pseudoCounts <- 0
for (i in 1:dim(Y)[1]){
  
  if(sum(Y[i,1:groups[1]])==0) {
    Y[i,sample(1:groups[1],size=1)] <- 1
    pseudoCounts = pseudoCounts + 1
  }
  if(sum(Y[i,(1+groups[1]):(groups[2]+groups[1])])==0) {
    Y[i,sample((1+groups[1]):(groups[2]+groups[1]),size=1)] <- 1
    
    pseudoCounts = pseudoCounts +1
  }
}





normalizingConstants <- colSums(Y) #Uses total counts per sample
                                   #as normalization

#MCMC run settings
nrAdapt=1000
nrBurnIn=1000
nrSamples=10000
nrThin = 1

#runs the mcmc for the selected settings and returns the significance score,
#fold change and zero inflation parameter for each feature
resultList <- ZiMGLMMLoan(Y=Y,group = group, N = normalizingConstants,
                          nrAdapt=nrAdapt,nrBurnIn=nrBurnIn,nrSamples=nrSamples,
                          nrThin=nrThin,returnBeta=TRUE)
sigScore <- resultList[[1]]
foldChange <- resultList[[2]]
zeroInf <- resultList[[3]]
resOrder <- order(sigScore)


#export results
resultMat<-cbind(rownames(Y),sigScore,foldChange,zeroInf) 
resultMatOrdered<-cbind(rownames(Y)[resOrder],sigScore[resOrder],foldChange[resOrder],zeroInf[resOrder]) 
write.table(resultMatOrdered,file="resultTable.txt",sep="\t")