##########################################################################

### Script to perform integrative Clustering of several omics datasets ###

##########################################################################

# get options

library("optparse")

option_list = list(

  make_option(c("-d", "--DNA"), type="character", default=NULL, help="file with mutation data [default= %default]", metavar="character"),

  make_option(c("-r", "--RNA"), type="character", default=NULL, help="file with expression data [default= %default]", metavar="character"),

  make_option(c("-k", "--Kmin"), type="numeric", default=2, help="minimum number of clusters [default= %default]", metavar="numeric"),

  make_option(c("-K", "--Kmax"), type="numeric", default=6, help="maximum number of clusters [default= %default]", metavar="numeric"),

  make_option(c("-m", "--methylation"), type="character", default=NULL, help="file with methylation data [default= %default]", metavar="character"),

  make_option(c("-C", "--CNV"), type="character", default=NULL, help="file with copy number data [default= %default]", metavar="character"),

  make_option(c("-o", "--out"), type="character", default="out", help="output directory name [default= %default]", metavar="character"),

  make_option(c("-c", "--cores"), type="numeric", default=1, help="number of cores for statistical computation [default= %default]", metavar="numeric")

); 

opt_parser = OptionParser(option_list=option_list);

opt = parse_args(opt_parser);

Kmin  = as.numeric(opt$Kmin)

Kmax  = as.numeric(opt$Kmax)

cores = as.numeric(opt$cores)

require(iClusterPlus)

require(gplots)

require(lattice)

# read data

Ddata = NULL

Cdata = NULL

Rdata = NULL

Mdata = NULL

dtl = vector("list",4)

ntypes = 0

type = c()

row.order = c()

plot.chr  = c()

sparse    = c()

cap       = c()

if(!is.null(opt$DNA)){

  ntypes = ntypes +1

  Ddata = read.table(opt$DNA,h=T,row.names = 1)

  dtl[[ntypes]] = t(Ddata)

  type = c(type,"binomial")

  row.order = c(row.order,T)

  plot.chr  = c(plot.chr,F)

  sparse    = c(sparse,T)

  cap       = c(cap,T)

}

if(!is.null(opt$CNV)){

  ntypes = ntypes +1

  dtl[[ntypes]] = t(Cdata)

  type = c(type,"gaussian")

  row.order = c(row.order,T)

  plot.chr  = c(plot.chr,T)

  sparse    = c(sparse,T)

  cap       = c(cap,T)

}

if(!is.null(opt$RNA)){

  ntypes = ntypes +1

  Rdata = read.table(opt$RNA,h=T,row.names = 1,check.names = T)

  dtl[[ntypes]] = t(Rdata)

  type = c(type,"gaussian")

  row.order = c(row.order,T)

  plot.chr  = c(plot.chr,F)

  sparse    = c(sparse,T)

  cap       = c(cap,T)

}

if(!is.null(opt$methylation)){

  ntypes = ntypes +1

  Mdata = read.table(opt$methylation,h=T,row.names = 1)

  dtl[[ntypes]] = t(Mdata)

  type = c(type,"gaussian")

  row.order = c(row.order,T)

  plot.chr  = c(plot.chr,F)

  sparse    = c(sparse,T)

  cap       = c(cap,T)

}

print(dtl)

print(type)

# tune lambda

cv.fitl = vector("list",Kmax-Kmin+1)

for(k in (Kmin:Kmax-1) ){

  cv.fitl[[k]] = tune.iClusterPlus(cpus=cores,dt1=dtl[[1]],dt2=dtl[[2]],dt3=dtl[[3]],dt4=dtl[[4]], type=type, n.lambda=NULL,K=k,maxiter=20)

  cv.fittmp = cv.fitl[[k]]

  save(cv.fittmp, file=paste(opt$out,"cv.fit.K",k+1,".Rdata",sep=""))

}

warnings()

nLambda = nrow(cv.fitl[[1]]$lambda)

nK = length(cv.fitl)

BIC = getBIC(cv.fitl)

devR = getDevR(cv.fitl)

minBICid = apply(BIC,2,which.min)

devRatMinBIC = rep(NA,nK)

for(i in 1:nK){

  devRatMinBIC[i] = devR[minBICid[i],i]

}

#best.fitl=sapply( Kmin:Kmax-1 , function(k) cv.fitl[[k]]$fit[[which.min(BIC[,k])]] )

clusters=getClusters(cv.fitl)

colnames(clusters)=paste("K=",Kmin:Kmax,sep="")

clusters2 = clusters

k = Kmin

while(k<Kmax){

  print(k)

  clusters2 = clusters

  for(i in 1:max(clusters[,k-1]) ){

    print(i)

    tatmp = table(clusters[clusters[,k-1]==i,k])

    newid = as.numeric(names(tatmp)[which.max(tatmp)])

    clusters2[clusters[,k]==newid,k] = i

    clusters2[clusters[,k]==i,k] = newid

    print(clusters2)

    clusters = clusters2

  }

  k = k+1

}

svg(paste(opt$out,"iClusters.svg",sep=""),h=6,w=8*2)

par(family="Times",mfrow=c(1,2))

plot(Kmin:Kmax,devRatMinBIC,type="b",xlab=expression(italic(K)),ylab=expression("Pseudo "~italic(R^2)),las=1)

image(1:nrow(clusters),Kmin:Kmax,clusters,col=rainbow(Kmax),axes=F,xlab="",ylab=expression(italic(K)),las=2 )

axis(1,at = 1:nrow(clusters),labels = rownames(dtl[[1]]) ,las=2)

axis(2,at = Kmin:Kmax,las=1)

dev.off()

# plot heatmap and write features

features = lapply(dtl,colnames)

for(k in (Kmin:Kmax-1) ){

  print(k)

  best.cluster=clusters[,k]

  print(best.cluster)

  best.fit=cv.fitl[[k]]$fit[[which.min(BIC[,k])]]

  print(best.fit)

  col.scheme = alist()

  col.scheme[[1]] = bluered(256)

  col.scheme[[2]] = bluered(256)

  col.scheme[[3]] = bluered(256)

  col.scheme[[4]] = bluered(256)

  png(paste(opt$out,"Heatmap_K",k+1,".png",sep=""),h=4*200,w=4*300,bg=0)

  plotHeatmap(fit=best.fit,datasets=dtl[1:ntypes], type=type, col.scheme = col.scheme[1:ntypes], row.order=row.order,plot.chr=plot.chr,sparse=sparse,cap=cap)

  dev.off()

  # get best features

  sigfeatures=alist()

  idsigfeatures=alist()

  print("yop")

  print(sum(!sapply(features,is.null)))

  for(i in 1:sum(!sapply(features,is.null)) ){

    print(i)

    rowsum=apply(abs(best.fit$beta[[i]]),1, sum)

    upper=quantile(rowsum,prob=0.75)

    sigfeatures[[i]]=(features[[i]])[which(rowsum>upper)]

    idsigfeatures[[i]]=which(rowsum>upper)

    write.table(sigfeatures[[i]],file = paste(opt$out,"_K",k+1,"_features",i,".txt",sep=""),col.names = F,quote = F)

  }

  #names(sigfeatures)=c("expression","methylation")

}