####################################################
### Script to perform MOFA                       ###
####################################################

# get options
library("optparse")

option_list = list(
  make_option(c("-g", "--group_file"), type="character", default=NULL, help="file with sample groups [default= %default]", metavar="character"),
  make_option(c("-o", "--out"), type="character", default="MOFA_out", help="output directory name [default= %default]", metavar="character"),
  make_option(c("-s", "--suffix"), type="character", default="count.txt", help="Suffix for output names [default= %default]", metavar="character"),
  make_option(c("-R", "--expression"), type="character", default=NULL, help="Expression dataset [default= %default]", metavar="character"),
  make_option(c("-M", "--methylation"), type="character", default=NULL, help="Methylation dataset [default= %default]", metavar="character"),
  make_option(c("-S", "--mutation"), type="character", default=NULL, help="Mutations dataset [default= %default]", metavar="character"),
  make_option(c("-r", "--robustness"), action="store_true", default=FALSE, help="Perform robustness analysis [default= %default]"),
  make_option(c("-i", "--maxiter"), type="numeric", default=10000, help="Maximum number of iterations [default= %default]"),
  make_option(c("-n", "--nconsreps"), type="numeric", default=100, help="Number of iterations for consensus clustering [default= %default]"),
  make_option(c("-m", "--nreps"), type="numeric", default=10, help="Maximum number of iterations for Latent Factor robustness analysis [default= %default]"),
  make_option(c("-d", "--dropLFthres"), type="numeric", default=0, help="Threshold proportion of variance explained to drop latent factors [default= %default]"),
  make_option(c("-t", "--tol"), type="numeric", default=0.01, help="Tolerated delta ELBO to call convergence [default= %default]"),
  make_option(c("-F", "--nLF"), type="numeric", default=10, help="Number of latent factors [default= %default]"),
  make_option(c("-p", "--sparsity"), type="logical", action="store_false", default=TRUE, help="Sparsity of latent factors [default= %default]")
); 

require(MOFAtools)
## add something to specify class of group file columns

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

# read data
data = list()
if(!(is.null(opt$mutation)))    data$Mutation = read.table(opt$mutation,h=T,row.names = 1)
if(!(is.null(opt$expression)))  data$RNA      = read.table(opt$expression,h=T,row.names = 1)
if(!(is.null(opt$methylation))) data$Methyl   = read.table(opt$methylation,h=T,row.names = 1)

if(!is.null(opt$robustness) ){
  print("Perform robustness analysis")
  nrep = as.numeric(opt$nconsreps)
  pItem = 0.8
  for(i in 1:nrep){
    print(i)
    print(c("outfile",paste(opt$out,"/MOFA",opt$suffix,"_sub",i,".hdf5",sep="")))
    rm( MOFAobjecttmp )
    samptmp = sort(unique( unlist(sapply(data,colnames)) ))
    nsamp = length(samptmp)
    samptmp = sort(sample(samptmp,size = round(pItem*nsamp),replace=F) )
    data2 = lapply(data, function(x) x[,unlist(sapply(samptmp, function(xx) which(colnames(x)==xx) )) ] )
    print(sapply(data2,dim))
    V2    = lapply(data2, function(x) apply(x,1,var) )

    zeroV2 = lapply(V2, function(x) which(x==0))
    for(j in 1:length(data2)){
	print(paste("Remove",length(zeroV2[[j]]),"samples with 0 variance in view",j) )
	if( length(zeroV2[[j]])>0 ) data2[[j]] = data2[[j]][-zeroV2[[j]],]
    }
    print(sapply(data2,dim))    
    MOFAobjecttmp <- createMOFAobject(data2)  
    ModelOptions <- getDefaultModelOptions(MOFAobjecttmp)
    ModelOptions$likelihood[names(ModelOptions$likelihood)=="Mutation"] = "bernoulli"
    ModelOptions$numFactors  <- as.numeric(opt$nLF)
    ModelOptions$sparsity    <- as.logical(opt$sparsity)
    TrainOptions <- getDefaultTrainOptions()
    TrainOptions$maxiter = as.numeric(opt$maxiter)
    TrainOptions$DropFactorThreshold = as.numeric(opt$dropLFthres)
    TrainOptions$tolerance = as.numeric(opt$tol)
    DataOptions <- getDefaultDataOptions()
    
    print(ModelOptions)
    print(TrainOptions)
    print(DataOptions)
    MOFAobjecttmp <- prepareMOFA(MOFAobjecttmp, ModelOptions = ModelOptions,TrainOptions = TrainOptions,DataOptions = DataOptions)
  
    # run
    MOFAobjecttmp <- runMOFA(MOFAobjecttmp, paste(opt$out,"/MOFA",opt$suffix,"_sub",i,".hdf5",sep="") )
  }
}

#perform MOFA run
print("Perform multiple runs")
MOFAfinal=c()
bestELBO = -Inf
bestmodel = 1
nrep = as.numeric(opt$nreps)
for(i in 1:nrep){
    print(i)
    print(c("outfile",paste(opt$out,"/MOFA",opt$suffix,"_run",i,".hdf5",sep="")) )
    MOFAobjecttmp <- createMOFAobject(data)
    ModelOptions <- getDefaultModelOptions(MOFAobjecttmp)
    ModelOptions$likelihood[names(ModelOptions$likelihood)=="Mutation"] = "bernoulli"
    ModelOptions$numFactors    <- as.numeric(opt$nLF)
    ModelOptions$sparsity    <- as.logical(opt$sparsity)
    TrainOptions <- getDefaultTrainOptions()
    TrainOptions$maxiter   = as.numeric(opt$maxiter)
    TrainOptions$DropFactorThreshold = as.numeric(opt$dropLFthres)
    TrainOptions$tolerance = as.numeric(opt$tol)
    DataOptions <- getDefaultDataOptions()
    print(ModelOptions)
    print(TrainOptions)
    print(DataOptions)
    MOFAobjecttmp <- prepareMOFA(MOFAobjecttmp, ModelOptions = ModelOptions,TrainOptions = TrainOptions,DataOptions = DataOptions)
    
    # run
    MOFAobjecttmp <- runMOFA(MOFAobjecttmp, paste(opt$out,"/MOFA",opt$suffix,"_run",i,".hdf5",sep="") )
    
    if( rev(MOFAobjecttmp@TrainStats$elbo)[1]> bestELBO){
      bestELBO = rev(MOFAobjecttmp@TrainStats$elbo)[1]
      MOFAfinal = MOFAobjecttmp
      bestmodel = i
    }
}
MOFAobject = MOFAfinal
print(paste("Best model:",bestmodel))

# convergence plot
pdf(paste(opt$out,"/Convergence_ELBO",opt$suffix,".pdf",sep=""),h=4,w=4*2)
par(mfrow=c(1,1),las=1,family="Times")
plot( diff(MOFAobject@TrainStats$elbo),type="l",xlab="Iteration",ylab=expression(Delta~"ELBO"),main="",log="y")
abline(h=0,lty=2)
dev.off()

# analyses      
# correlate other variables
if(!is.null(opt$group_file)){
  meta = read.table(opt$group_file,h=T,sep="\t")

  idsurv = grep("surv", colnames(meta) ,ignore.case = T)
  if( length(idsurv)>0){
    surv = meta[,idsurv]
    meta = meta[,-idsurv]
    issurv = TRUE
  }else{
    issurv = FALSE
  }
  
  MOFAfactors.nona = MOFAfactors[ rowSums( is.na( meta ) )==0 ,]
  
  lml = lapply(1:3, function(i) lm( as.formula(paste("MOFAfactors.nona[,",i+1,"]~", paste(colnames(meta) ,collapse = "+"),"+ Sex:Type + Sex:Smoking_Hx + Age:Smoking_Hx + Smoking_Hx:Type") ),data = na.omit(meta) ) )
  suml = lapply(lml,summary)
  anl = lapply(lml,anova)
  
  require(MASS)
  #lml.dropped = vector("list",3)
  #for(i in 1:3){
  #  lmtmp = lml[[i]]
  #  metatmp = meta
  #  while(1){
  #    dttmp = dropterm(lmtmp)
  #    if(max(dttmp$AIC,na.rm = T)< dttmp$AIC[1]  ) ttmp = rownames(dttmp)[which.min(dttmp$AIC)]
  #    else break
      #if(ttmp=="<none>") break
  #    metatmp = metatmp[,!(colnames(metatmp)==ttmp)]
  #    lmtmp = lm( as.formula(paste("MOFAfactors[,",i+1,"]~", paste(colnames(metatmp),collapse = "+")) ),data = metatmp ) 
      #lmtmp2 = update(lmtmp, paste(".~. - ",ttmp))
      #if( anova(lmtmp,lmtmp2)$`Pr(>F)`[2] >0.05/3) lmtmp = lmtmp2
  #    print(dttmp)
  #    #print(lmtmp)
  #  }
  #  lml.dropped[[i]] = lmtmp
  #}
  #anl.dropped = lapply(lml.dropped,anova,)
  lml.dropped = lapply(1:3, function(i) stepAIC(lml[[i]], direction = "both",test="F") )
  lml.final = lapply(1:3, function(i) lm( as.formula(paste("MOFAfactors[,",i+1,"]~", as.character(lml.dropped[[i]]$call$formula)[3] ) ) ,data = meta) )
  for(i in 1:3) write.table( anl[[i]],file = paste(opt$out,"/Anova",opt$suffix,"_LF",i,".txt",sep=""),quote = F) 
  anl.final = lapply(lml.final,anova)
  
  lsignif = lapply(1:3, function(i) rownames(anl.final[[i]])[which( anl.final[[i]]$`Pr(>F)`<0.05/3 )] )
  #lsignifall = unique( unlist(strsplit(unlist( lsignif ),":")) )
  lsignifall = unique( unlist(lsignif) )
  
  meta2 = cbind(meta, "Type:Sex" = interaction(meta$Type,meta$Sex), "Sex:Smoking_Hx" = interaction(meta$Sex,meta$Smoking_Hx) , "Age:Smoking_Hx" = interaction(meta$Age,meta$Smoking_Hx) )
  
  svg(paste(opt$out,"/Pie",opt$suffix,"_signif.svg",sep=""),h=3*length(lsignifall),w=3*K)
  par(mfrow=c(length(lsignifall),K),family="Times",las=1)
  for(j in 1:length(lsignifall)){
    for(i in 1:K) pie( table(factor(meta2[,lsignifall[j]])[kmb$cluster==i] ) ,radius = sqrt(sum(kmb$cluster==i)/70), init.angle = 90,col= rainbow(length(levels(factor(meta2[,lsignifall[j]])))) )
  }
  dev.off()
  
  fisher.test(table(meta2[,lsignifall[j]],kmb$cluster ))
  
  survS <- Surv( as.numeric(surv[,1]), as.numeric(factor(surv[,2],levels=c("alive","dead") ))-1)
  fitS       <- survfit(survS ~ MOFAfactors[,2] + MOFAfactors[,3] )
  fitS.clust <- survfit(survS ~ factor(kmb$cluster)  )
  fitS.Type  <- survfit(survS ~ factor(meta$Type)  )
  coxS       <- coxph(survS ~ MOFAfactors[,2] + MOFAfactors[,3]) 
  coxS.clust <- coxph(survS ~ factor(kmb$cluster)  )
  coxS.Type   <- coxph(survS ~ factor(meta$Type)  )
  summary(coxS)
  summary(coxS.clust)
  summary(coxS.Type)
  extractAIC(coxS) # better
  extractAIC(coxS.clust)
  extractAIC(coxS.Type)
  #summary(fitCarc.clust2)$table
  
  svg(paste(opt$out,"/Surv",opt$suffix,".svg",sep=""),h=4,w=4)
  par(mfrow=c(1,1),las=1,family="Times")
  plot(fitS.clust,col=1:4)
  legend("bottomright",legend = levels(as.factor(kmb$cluster)),col=1:3 , lty=1)
  legend("bottomleft",legend = paste("logrank P =", format(summary(coxS.clust)$logtest[3],digits = 3,scientific = T))  )
  #plot(fitCarc.Type,col=2:1)
  #legend("bottomright",legend = levels(as.factor(TypeCarc)),col=2:1 , lty=1)
  #legend("bottomleft",legend = paste("logrank P =", format(summary(coxCarc.Type)$logtest[3],digits = 3,scientific = T))  )
  dev.off()
  
  svg(paste(opt$out,"/Surv",opt$suffix,"_Type.svg",sep=""),h=4,w=4)
  par(mfrow=c(1,1),las=1,family="Times")
  plot(fitS.Type,col=1:4)
  legend("bottomright",legend = levels(as.factor(kmb$cluster)),col=1:3 , lty=1)
  legend("bottomleft",legend = paste("logrank P =", format(summary(coxS.Type)$logtest[3],digits = 3,scientific = T))  )
  #plot(fitCarc.Type,col=2:1)
  #legend("bottomright",legend = levels(as.factor(TypeCarc)),col=2:1 , lty=1)
  #legend("bottomleft",legend = paste("logrank P =", format(summary(coxCarc.Type)$logtest[3],digits = 3,scientific = T))  )
  dev.off()
  
  if(sum(colnames(meta)=="Type")>0){
    svg( paste(opt$out,"/ScatterPlot",opt$suffix,"_Type.svg",sep=""),h=4,w=4)
    par(mfrow=c(1,1),family="Times",las=1)
    plot(MOFAfactors[,2:3], xlab = "Latent Factor 1", ylab="Latent Factor 2", col = prettycolors[c(2,3,1,6)][addNA(as.factor(meta$Type))], pch=16 )
    legend("topleft",legend = levels(addNA(as.factor(meta$Type))) , col = prettycolors[c(2,3,1,6)], pch=16 )
    dev.off()
  }
  
  svg(paste(opt$out,"Ki67_AC-TC-LCNEC.svg",sep=""),h=4,w=4*2)
  par(mfrow=c(1,2),las=1,family="Times")
  boxplot( MOFAobject@TrainData$RNA[which( rownames( MOFAobject@TrainData$RNA )=="MKI67" ),] ~ meta$Type  )
  boxplot( MOFAobject@TrainData$RNA[which( rownames( MOFAobject@TrainData$RNA )=="MKI67" ),] ~ kmb$cluster  )
  dev.off()
  
}

save(MOFAobject,MOFAfactors, file = paste(opt$out,"/integration_MOFA",opt$suffix,".RData",sep="") )