setwd("/nfs/nfs_wooi/Botas/EA12048_botas/work")
source("/nfs/nfs_wooi/WTools/wrs/plotMultiplePCA.R")
source("/nfs/nfs_wooi/WTools/wrs/PCA.SampleNames.Gen.R")
library(DESeq2)

# Read in raw counts
cnt.orig <- read.delim("../Batch.2016-07-12_24089_Supp047/rsem_gcounts_matrix.txt", sep="\t", check.names=FALSE, stringsAsFactors=F)

# Read in sample information
sampleInfo <- read.delim("SampleInfo.Checks.txt", sep="\t", stringsAsFactors=TRUE)
rownames(sampleInfo) <- sampleInfo[,1]
sampleInfo <- sampleInfo[colnames(cnt.orig)[-1], ]

# Generic
# Extract data (g1/g2)
g1 = "RN5"
g2 = "RW5"

# -- 
GenericRun(g1="RN5", g2="RW5")
GenericRun(g1="RN7", g2="RW7")
GenericRun(g1="RN8", g2="RW8")
GenericRun(g1="RF18", g2="RW18")
GenericRun(g1="RF20", g2="RW20")
GenericRun(g1="RF22", g2="RW22")

GenericRun(g1="EN7", g2="EW7")
GenericRun(g1="EN9", g2="EW9")
GenericRun(g1="EN11", g2="EW11")
GenericRun(g1="EF18", g2="EW18")
GenericRun(g1="EF20", g2="EW20")
GenericRun(g1="EF22", g2="EW22")


GenericRun <- function(g1, g2){
  dds.outf = paste("DESeq_out/", g1, "vs", g2, "_", "dds.saved", sep="")
  dds.rld.outf = paste("DESeq_out/", g1, "vs", g2, "_", "dds.rld.saved", sep="")
  dds.res.outf = paste("DESeq_out/", g1, "vs", g2, "_", "dds.res.txt", sep="")
  dds.PCA.outf = paste("DESeq_out/", g1, "vs", g2, "_", "PCA.pdf", sep="")
  dds.PCA2.outf = paste("DESeq_out/", g1, "vs", g2, "_", "PCA_rlog.pdf", sep="")
  
  
  sInfo <- sampleInfo[c(grep(g1, sampleInfo$Sample), grep(g2, sampleInfo$Sample)), ]
  scnt <- cnt.orig[, rownames(sInfo)]
  rownames(scnt) <- cnt.orig[ ,1]
  
  scond <- data.frame(cond = factor(gsub("(A|B|C)$", "", sInfo$Sample), levels=c(g2, g1)))
  dds <- DESeqDataSetFromMatrix(countData = round(scnt), 
                                colData = scond,
                                design = ~cond)
  dds.rlog <- rlogTransformation(dds, blind = TRUE)
  dds.deseq <- DESeq(dds)
  dds.res <- results(dds.deseq, contrast=c("cond", g1, g2))
  dds.res$grp.mean <- sapply(levels(dds.deseq$cond), function(lvl) rowMeans(counts(dds.deseq,normalized=TRUE)[,dds.deseq$cond == lvl] ) )
  dds.res$norm.counts <- counts(dds.deseq, normalized=TRUE)
  colnames(dds.res$norm.counts) <- colnames(scnt)
  
  save(dds, file=dds.outf)
  save(dds.rlog, file=dds.rld.outf)
  write.table(as.data.frame(dds.res), file=dds.res.outf, sep="\t")
  
  ##pdf(file=dds.PCA.outf, width=6, height=5)
  PCA.SampleNames.Gen.Simple(as.matrix(dds.res$norm.counts), groups=scond$cond, ntop=nrow(dds.res$norm.counts), 
                             title=paste("PCA ", g1, "_", g2, sep=""), tsize=3, lpos="right", bsize=13)
  ggsave(file=dds.PCA.outf, width=6, height=5)
  
  pdat <- assay(dds.rlog)
  colnames(pdat) <- colnames(scnt)
  PCA.SampleNames.Gen.Simple(pdat, groups=scond$cond, ntop=nrow(pdat), 
                             title=paste("PCA rlog ", g1, "_", g2, sep=""), tsize=3, lpos="right", bsize=13)
  ggsave(file=dds.PCA2.outf, width=6, height=5)
  
}


# --- Now put all lists together
library(plyr)

cmps <- read.delim("ComparisonsList.txt", sep="\t", header=F, stringsAsFactors=F)

# Get neuron comparisons
cmps.neuron <- cmps[7:12, ]
cmps.glia <- cmps[1:6, ]

all.neuron <- c()
suff <- paste(cmps.neuron[1,2], "vs", cmps.neuron[1,1], sep="")
temp <- read.delim(paste("DESeq_out/", suff, "_dds.res.txt", sep=""), sep="\t", stringsAsFactors=F)[c(2, 5:8)]
colnames(temp) <- paste(suff, "_", colnames(temp), sep="")
all.neuron <- temp
all.neuron$gene <- rownames(all.neuron)

for(i in 2:nrow(cmps.neuron)){
  suff <- paste(cmps.neuron[i,2], "vs", cmps.neuron[i,1], sep="")
  temp <- read.delim(paste("DESeq_out/", suff, "_dds.res.txt", sep=""), sep="\t", stringsAsFactors=F)[c(2, 5:8)]
  colnames(temp) <- paste(suff, "_", colnames(temp), sep="")
  temp$gene <- rownames(temp)
  
  all.neuron <- join_all(list(all.neuron,temp), by = 'gene', type = 'full')  
}

write.table(all.neuron, file="DESeq_all.neuron.txt", sep="\t")

# --
all.glia <- c()
suff <- paste(cmps.glia[1,2], "vs", cmps.glia[1,1], sep="")
temp <- read.delim(paste("DESeq_out/", suff, "_dds.res.txt", sep=""), sep="\t", stringsAsFactors=F)[c(2, 5:8)]
colnames(temp) <- paste(suff, "_", colnames(temp), sep="")
all.glia <- temp
all.glia$gene <- rownames(all.glia)

for(i in 2:nrow(cmps.glia)){
  suff <- paste(cmps.glia[i,2], "vs", cmps.glia[i,1], sep="")
  temp <- read.delim(paste("DESeq_out/", suff, "_dds.res.txt", sep=""), sep="\t", stringsAsFactors=F)[c(2, 5:8)]
  colnames(temp) <- paste(suff, "_", colnames(temp), sep="")
  temp$gene <- rownames(temp)
  
  all.glia <- join_all(list(all.glia,temp), by = 'gene', type = 'full')  
}

write.table(all.glia, file="DESeq_all.glia.txt", sep="\t")


# -------- Plot heatmap
all.neuron <- read.delim("DESeq_all.neuron.txt", sep="\t", stringsAsFactors=F)
all.glia <- read.delim("DESeq_all.glia.txt", sep="\t", stringsAsFactors=F)

th = 0.01

rownames(all.neuron) <- all.neuron$gene
rownames(all.glia) <- all.glia$gene

# ~~~~~~~~~~~~~~  Neuron
padj.neuron <- apply(all.neuron[, grep("_padj$", colnames(all.neuron))], 2, function(x) !is.na(x) & x < 0.01)
apply(padj.neuron, 2, sum)

or.neuron <- apply(padj.neuron, 1, sum)
sum(or.neuron>0)

# make data for plotting
plot.dat <- all.neuron[or.neuron>0, grep("_log2FoldChange$", colnames(all.neuron))]
colnames(plot.dat) <- gsub("_log2FoldChange$", "", colnames(plot.dat))

sfac <- factor(c("Sig", "NotSig"), levels=c("Sig", "NotSig"))

###annR <- apply(padj.neuron[or.neuron>0, ], 2, function(x) factor(ifelse(x, "Sig", "NotSig")))
temp = padj.neuron[or.neuron>0, ]

annR <- data.frame(sfac[as.numeric(!temp[,1])+1])

for(i in 2:ncol(temp)){
  annR <- data.frame(annR, sfac[as.numeric(!temp[,i])+1])
}
colnames(annR) <- colnames(temp)
rownames(annR) <- rownames(temp)


temp = unlist(strsplit(colnames(plot.dat), "vs"))[seq(1, 12, by=2)]
temp = factor(temp, levels=temp)
annC <- data.frame(genotype=temp)
rownames(annC) <- colnames(plot.dat)

ccnt = 20
cols <- colorRampPalette(c("navy", "white", "firebrick3"))(ccnt)
colbk <- seq(-2, 2, length.out=ccnt+1)

png(file="test.png", width=8 height=8,units="in",res=600)
pheatmap(as.matrix(plot.dat), 
         color = cols, breaks=colbk,
         annotation_col = annC,
         annotation_row = annR,
         annotation_colors = list(EN7vsEW7_padj=c(Sig="red3", NotSig="gray90"),
                                  EN9vsEW9_padj =c(Sig="red3", NotSig="gray90"),
                                  EN11vsEW11_padj =c(Sig="red3", NotSig="gray90"),
                                  EF18vsEW18_padj =c(Sig="red3", NotSig="gray90"),
                                  EF20vsEW20_padj =c(Sig="red3", NotSig="gray90"),
                                  EF22vsEW22_padj=c(Sig="red3", NotSig="gray90")),
         fontsize_row=2,
         cluster_col=FALSE
)
dev.off()



# ~~~~~~~~~~~~~~  glia
padj.glia <- apply(all.glia[, grep("_padj$", colnames(all.glia))], 2, function(x) !is.na(x) & x < 0.01)
apply(padj.glia, 2, sum)

or.glia <- apply(padj.glia, 1, sum)
sum(or.glia>0)

# make data for plotting
plot.dat <- all.glia[or.glia>0, grep("_log2FoldChange$", colnames(all.glia))]
colnames(plot.dat) <- gsub("_log2FoldChange$", "", colnames(plot.dat))

sfac <- factor(c("Sig", "NotSig"), levels=c("Sig", "NotSig"))

###annR <- apply(padj.glia[or.glia>0, ], 2, function(x) factor(ifelse(x, "Sig", "NotSig")))
temp = padj.glia[or.glia>0, ]

annR <- data.frame(sfac[as.numeric(!temp[,1])+1])

for(i in 2:ncol(temp)){
  annR <- data.frame(annR, sfac[as.numeric(!temp[,i])+1])
}
colnames(annR) <- colnames(temp)
rownames(annR) <- rownames(temp)


temp = unlist(strsplit(colnames(plot.dat), "vs"))[seq(1, 12, by=2)]
temp = factor(temp, levels=temp)
annC <- data.frame(genotype=temp)
rownames(annC) <- colnames(plot.dat)

ccnt = 20
cols <- colorRampPalette(c("navy", "white", "firebrick3"))(ccnt)
colbk <- seq(-2, 2, length.out=ccnt+1)

png(file="test.png", width=8 height=8,units="in",res=600)
pheatmap(as.matrix(plot.dat), 
         color = cols, breaks=colbk,
         annotation_col = annC,
         annotation_row = annR,
         annotation_colors = list(RN5vsRW5_padj=c(Sig="red3", NotSig="gray90"),
                                  RN7vsRW7_padj =c(Sig="red3", NotSig="gray90"),
                                  RN8vsRW8_padj =c(Sig="red3", NotSig="gray90"),
                                  RF18vsRW18_padj =c(Sig="red3", NotSig="gray90"),
                                  RF20vsRW20_padj =c(Sig="red3", NotSig="gray90"),
                                  RF22vsRW22_padj=c(Sig="red3", NotSig="gray90")),
         fontsize_row=2,
         cluster_col=FALSE
)
dev.off()


# AND
or.glia <- apply(padj.glia, 1, sum)
sum(or.glia == 6)

# make data for plotting
plot.dat <- all.glia[or.glia == 6, grep("_log2FoldChange$", colnames(all.glia))]
colnames(plot.dat) <- gsub("_log2FoldChange$", "", colnames(plot.dat))

sfac <- factor(c("Sig", "NotSig"), levels=c("Sig", "NotSig"))

###annR <- apply(padj.glia[or.glia>0, ], 2, function(x) factor(ifelse(x, "Sig", "NotSig")))
temp = padj.glia[or.glia == 6, ]

annR <- data.frame(sfac[as.numeric(!temp[,1])+1])

for(i in 2:ncol(temp)){
  annR <- data.frame(annR, sfac[as.numeric(!temp[,i])+1])
}
colnames(annR) <- colnames(temp)
rownames(annR) <- rownames(temp)


temp = unlist(strsplit(colnames(plot.dat), "vs"))[seq(1, 12, by=2)]
temp = factor(temp, levels=temp)
annC <- data.frame(genotype=temp)
rownames(annC) <- colnames(plot.dat)

ccnt = 20
cols <- colorRampPalette(c("navy", "white", "firebrick3"))(ccnt)
colbk <- seq(-6, 7, length.out=ccnt+1)

png(file="test.png", width=8 height=8,units="in",res=600)
pheatmap(as.matrix(plot.dat), 
         color = cols, breaks=colbk,
         annotation_col = annC,
         annotation_row = annR,
         annotation_colors = list(RN5vsRW5_padj=c(Sig="red3", NotSig="gray90"),
                                  RN7vsRW7_padj =c(Sig="red3", NotSig="gray90"),
                                  RN8vsRW8_padj =c(Sig="red3", NotSig="gray90"),
                                  RF18vsRW18_padj =c(Sig="red3", NotSig="gray90"),
                                  RF20vsRW20_padj =c(Sig="red3", NotSig="gray90"),
                                  RF22vsRW22_padj=c(Sig="red3", NotSig="gray90")),
         fontsize_row=4,
         cluster_col=FALSE
)
dev.off()


# ---






  
# 








