args <- commandArgs(trailingOnly = TRUE)

#types files
gene2typeFile  = "input/gene2typ.csv"

########################################################
###################FIGURE 4a############################
#heatmap for each type of interests overall time points

#data files
tp02File = args[1] 
tp04File = args[2] 
tp08File = args[3] 
tp16File = args[4] 
tp24File = args[5] 

#output dir
dir = args[6]

#read mapping
gene2type = read.table(gene2typeFile, header=T, stringsAsFactors=F)

# read data
tp02 = read.table(tp02File,header=T,stringsAsFactors=F,row.names=1)
tp04 = read.table(tp04File,header=T,stringsAsFactors=F,row.names=1)
tp08 = read.table(tp08File,header=T,stringsAsFactors=F,row.names=1)
tp16 = read.table(tp16File,header=T,stringsAsFactors=F,row.names=1)
tp24 = read.table(tp24File,header=T,stringsAsFactors=F,row.names=1)

row.names(tp02) = gsub("ID=","",gsub("gene_id=","",gsub(";.*","",row.names(tp02))))
row.names(tp04) = gsub("ID=","",gsub("gene_id=","",gsub(";.*","",row.names(tp04))))
row.names(tp08) = gsub("ID=","",gsub("gene_id=","",gsub(";.*","",row.names(tp08))))
row.names(tp16) = gsub("ID=","",gsub("gene_id=","",gsub(";.*","",row.names(tp16))))
row.names(tp24) = gsub("ID=","",gsub("gene_id=","",gsub(";.*","",row.names(tp24))))




 # combine fold changes as one table and set column names and row names


genenames = unique(c(row.names(tp02),row.names(tp04),row.names(tp08),row.names(tp16), row.names(tp24)))

allLogFoldChanges = matrix(nrow = length(genenames), ncol= 5)
row.names(allLogFoldChanges) = genenames
for(i in genenames){
  if(any(row.names(tp02) == i)){
    allLogFoldChanges[i,1] = tp02[i,1]
  }
  if(any(row.names(tp04) == i)){
    allLogFoldChanges[i,2] = tp04[i,1]
  }
  if(any(row.names(tp08) == i)){
    allLogFoldChanges[i,3] = tp08[i,1]
  }
  if(any(row.names(tp16) == i)){
    allLogFoldChanges[i,4] = tp16[i,1]
  }
  if(any(row.names(tp24) == i)){
    allLogFoldChanges[i,5] = tp24[i,1]
  }
}

row.names(allLogFoldChanges) = gsub("\\.[0-9]+","",genenames)
colnames(allLogFoldChanges) = c("2 h", "4 h","8 h","16 h","24 h")


#make the same for the pvalues
allPVals =  matrix(nrow = length(genenames), ncol= 5)
row.names(allPVals) = genenames
for(i in genenames){
  if(any(row.names(tp02) == i)){
    allPVals[i,1] = tp02[i,5]
  }
  if(any(row.names(tp04) == i)){
    allPVals[i,2] = tp04[i,5]
  }
  if(any(row.names(tp08) == i)){
    allPVals[i,3] = tp08[i,5]
  }
  if(any(row.names(tp16) == i)){
    allPVals[i,4] = tp16[i,5]
  }
  if(any(row.names(tp24) == i)){
    allPVals[i,5] = tp24[i,5]
  }
}

colnames(allPVals) = c("2 h", "4 h","8 h","16 h","24 h")
row.names(allPVals) = gsub("\\.[0-9]+","",genenames)



sigcount= apply(allPVals,1,function(x){
  return(length(which(x<=0.05)))
})

orgdata = allLogFoldChanges
allLogFoldChanges = allLogFoldChanges[sigcount>0,]



#filter for sets
mRNAs = gene2type[gene2type[,2] == "mRNA",1]
lncRNAs = gene2type[gene2type[,2] == "lncRNA",1]
miRNAs = gene2type[gene2type[,2] == "miRNA",1]
snRNAs = gene2type[gene2type[,2] == "snRNA",1]
snoRNAs = gene2type[gene2type[,2] == "snoRNA",1]

mRNA_fcs = allLogFoldChanges[row.names(allLogFoldChanges) %in% mRNAs,]
lncRNA_fcs = allLogFoldChanges[row.names(allLogFoldChanges) %in% lncRNAs,]
miRNA_fcs = allLogFoldChanges[row.names(allLogFoldChanges) %in% miRNAs,]
snRNA_fcs = allLogFoldChanges[row.names(allLogFoldChanges) %in% snRNAs,]
snoRNA_fcs = allLogFoldChanges[row.names(allLogFoldChanges) %in% snoRNAs,]
mRNA_pvals = allPVals[row.names(allPVals) %in% mRNAs,]
lncRNA_pvals = allPVals[row.names(allPVals) %in% lncRNAs,]
miRNA_pvals = allPVals[row.names(allPVals) %in% miRNAs,]
snRNA_pvals = allPVals[row.names(allPVals) %in% snRNAs,]
snoRNA_pvals = allPVals[row.names(allPVals) %in% snoRNAs,]

mRNA_org =   orgdata[row.names(orgdata) %in% mRNAs,]
lncRNA_org = orgdata[row.names(orgdata) %in% lncRNAs,]
miRNA_org =  orgdata[row.names(orgdata) %in% miRNAs,]
snRNA_org =  orgdata[row.names(orgdata) %in% snRNAs,]
snoRNA_org = orgdata[row.names(orgdata) %in% snoRNAs,]


#make tables with log fold changes and p-values for each cell type 
n = 5
s <- rep(1:n, each = 2) + (0:1) * n
colnames = c("2h log2FC", "2h adj. p-value","4h log2FC", "4h adj. p-value","8h log2FC", "8h adj. p-value","16h log2FC", "16h adj. p-value","24h log2FC", "24h adj. p-value")

#mRNA_org
mRNA_pvals = allPVals[rownames(mRNA_fcs),]
mRNAs_table = cbind.data.frame(mRNA_fcs,mRNA_pvals)[,s]
colnames(mRNAs_table) = colnames
write.table(mRNAs_table,file="mRNA.csv",quote=F,sep="\t",row.names=T,col.names=T)


#lncRNA_org
lncRNA_pvals = allPVals[rownames(lncRNA_fcs),]
lncRNAs_table = cbind.data.frame(lncRNA_fcs,lncRNA_pvals)[,s]
colnames(lncRNAs_table) = colnames
write.table(lncRNAs_table,file="lncRNA.csv",quote=F,sep="\t",row.names=T,col.names=T)

#miRNA_fcs
miRNA_pvals = allPVals[rownames(miRNA_fcs),]
miRNAs_table = cbind.data.frame(miRNA_fcs,miRNA_pvals)[,s]
colnames(miRNAs_table) = colnames
write.table(miRNAs_table,file="miRNA.csv",quote=F,sep="\t",row.names=T,col.names=T)


#snRNA_fcs
snRNA_pvals = allPVals[rownames(snRNA_fcs),]
snRNAs_table = cbind.data.frame(snRNA_fcs,snRNA_pvals)[,s]
colnames(snRNAs_table) = colnames
write.table(snRNAs_table,file="snRNA.csv",quote=F,sep="\t",row.names=T,col.names=T)


#snoRNA_fcs
snoRNA_pvals = allPVals[rownames(snoRNA_fcs),]
snoRNAs_table = cbind.data.frame(snoRNA_fcs,snoRNA_pvals)[,s]
colnames(snoRNAs_table) = colnames
write.table(snoRNAs_table,file="snoRNA.csv",quote=F,sep="\t",row.names=T,col.names=T)


replaceNA = function(x){
  for(i in 1:length(x)){
    if(is.na(x[i])){
      x[i] = 0
    }
  }
  return(x)
}

mRNA_fcs = t(apply(mRNA_fcs,1,replaceNA))
lncRNA_fcs = t(apply(lncRNA_fcs,1,replaceNA))
miRNA_fcs = t(apply(miRNA_fcs,1,replaceNA))
snRNA_fcs = t(apply(snRNA_fcs,1,replaceNA))
snoRNA_fcs = t(apply(snoRNA_fcs,1,replaceNA))

mRNA_org = t(apply(mRNA_org,1,replaceNA))
lncRNA_org = t(apply(lncRNA_org,1,replaceNA))
miRNA_org = t(apply(miRNA_org,1,replaceNA))
snRNA_org = t(apply(snRNA_org,1,replaceNA))
snoRNA_org = t(apply(snoRNA_org,1,replaceNA))




#load gplots for heatmap.2
library(gplots)



#color scale
colorRamp = colorRampPalette(c("blue","white","red"))



goiM = c("ENSG00000198786","ENSG00000198886","ENSG00000198727","ENSG00000198888","ENSG00000256618","ENSG00000198763","ENSG00000198840","ENSG00000198938","ENSG00000198899","ENSG00000198712","ENSG00000198695","ENSG00000228253","ENSG00000198804","ENSG00000269028","ENSG00000169429","ENSG00000184557","ENSG00000120833")


goiL = c("ENSG00000213468","ENSG00000245532")

mRNADendro = heatmap.2(as.matrix(mRNA_fcs),Colv=F,dendrogram="row",col=colorRamp(40),trace="none",labRow=F,labCol=F,key=F)$rowDendrogram

weights = numeric(length(row.names(mRNA_fcs)))
names(weights) = row.names(mRNA_fcs)
weights = rep(5,length(weights))
weights[goiM] = rep(-50,length(goiM))

newDendro = reorder(mRNADendro,weights)
mRNAOrder = order.dendrogram(newDendro)
mRNA_fcs = mRNA_fcs[mRNAOrder,]

lncRNADendro = heatmap.2(as.matrix(lncRNA_fcs),Colv=F,dendrogram="row",col=colorRamp(40),trace="none",labRow=F,labCol=F,key=F)$rowDendrogram
weights = numeric(length(row.names(lncRNA_fcs)))
names(weights) = row.names(lncRNA_fcs)

weights[goiL] = rep(-5,length(goiL))

newDendro = reorder(lncRNADendro,weights)
lncRNAOrder = order.dendrogram(newDendro)
lncRNA_fcs = lncRNA_fcs[lncRNAOrder,]

miRNADendro = heatmap.2(as.matrix(miRNA_fcs),Colv=F,dendrogram="row",col=colorRamp(40),trace="none",labRow=F,labCol=F,key=F)$rowDendrogram
miRNAOrder = order.dendrogram(miRNADendro)
miRNA_fcs = miRNA_fcs[miRNAOrder,]


snRNADendro = heatmap.2(as.matrix(snRNA_fcs),Colv=F,dendrogram="row",col=colorRamp(40),trace="none",labRow=F,labCol=F,key=F)$rowDendrogram
snRNAOrder = order.dendrogram(snRNADendro)
snRNA_fcs = snRNA_fcs[snRNAOrder,]


snoRNADendro = heatmap.2(as.matrix(snoRNA_fcs),Colv=F,dendrogram="row",col=colorRamp(40),trace="none",labRow=F,labCol=F,key=F)$rowDendrogram
snoRNAOrder = order.dendrogram(snoRNADendro)
snoRNA_fcs = snoRNA_fcs[snoRNAOrder,]





#make heatmaps
svg(paste(dir,"mRNA_weightedvarnorm.svg", sep="/"))
heatmap.2(as.matrix(mRNA_fcs),Colv=F,Rowv=F,dendrogram="none",col=colorRamp(40),trace="none",labRow=F,labCol=F,breaks=seq(-2,2,0.1),key=F,margin=c(0.1,0.1),lhei=c(0.01,0.999),lwid=c(0.01,0.999))
dev.off()

svg(paste(dir,"lncRNA_weightedvarnorm.svg",sep="/"))
heatmap.2(as.matrix(lncRNA_fcs),Colv=F,Rowv=F,dendrogram="none",col=colorRamp(40),trace="none",labRow=F,labCol=F,breaks=seq(-2,2,0.1),key=F,margin=c(0.1,0.1),lhei=c(0.01,0.999),lwid=c(0.01,0.999))
dev.off()

svg(paste(dir,"miRNA_weightedvarnorm.svg",sep="/"))
heatmap.2(as.matrix(miRNA_fcs),Colv=F,Rowv=F,dendrogram="none",col=colorRamp(40),trace="none",labRow=F,labCol=F,breaks=seq(-2,2,0.1),key=F,margin=c(0.1,0.1),lhei=c(0.01,0.999),lwid=c(0.01,0.999))
dev.off()


svg(paste(dir,"snRNA_weightedvarnorm.svg",sep="/"))
heatmap.2(as.matrix(snRNA_fcs),Colv=F,Rowv=F,dendrogram="none",col=colorRamp(40),trace="none",labRow=F,labCol=F,breaks=seq(-2,2,0.1),key=F,margin=c(0.1,0.1),lhei=c(0.01,0.999),lwid=c(0.01,0.999))
dev.off()

svg(paste(dir,"snoRNA_weightedvarnorm.svg",sep="/"))
heatmap.2(as.matrix(snoRNA_fcs),Colv=F,Rowv=F,dendrogram="none",col=colorRamp(40),trace="none",labRow=F,labCol=F,breaks=seq(-2,2,0.1),key=F,margin=c(0.1,0.1),lhei=c(0.01,0.999),lwid=c(0.01,0.999))
dev.off()


#count regulated
nrow(mRNA_fcs)
nrow(lncRNA_fcs)
nrow(miRNA_fcs)
nrow(snRNA_fcs)
nrow(snoRNA_fcs)
#count detected
nrow(mRNA_org)
nrow(lncRNA_org)
nrow(miRNA_org)
nrow(snRNA_org)
nrow(snoRNA_org)

#######do manually afterwards############
#put together the different heatmaps and annotated them with the correspoding data
#######end manually do##################
