library(stringr)
source('heatmap.3.R')

co5 <- read.delim('seg_data/CO5.3ormore.long.txt')
co8 <- read.delim('seg_data/CO8combined.3ormore.long.txt')


co5.filtered.cell.names <- readLines("CO5_filtered_cell_names_normals.txt")
co8.filtered.cell.names <- readLines("CO8_filtered_cell_names_normals.txt")

co5.prim.ploidy <- 2.5
co5.met.ploidy <- 2.9
co8.prim.ploidy <- 3.3
co8.met.ploidy <- 3.0

co5.ratios <- 2^co5[,-(1:2)]
co8.ratios <- 2^co8[,-(1:2)]
co5.ratios <- co5.ratios[,str_match(colnames(co5.ratios), '([A-Z]A?\\.\\d+)[._].*$')[,2] %in% co5.filtered.cell.names]
co8.ratios <- co8.ratios[,str_match(colnames(co8.ratios), '([A-Z]A?\\.\\d+)[._].*$')[,2] %in% co8.filtered.cell.names]

co5.normal.numbers <- c(19, 1)
co8.normal.numbers <- c(68, 58, 60, 69, 70, 12)

co5.is.normal <- as.integer(str_match(colnames(co5.ratios), "^C05[MP]\\.(\\d+)\\.bl$")[,2]) %in% co5.normal.numbers
co8.is.normal <- as.integer(str_match(colnames(co8.ratios), "^[MP]A\\.(\\d+)_$")[,2]) %in% co8.normal.numbers

colnames(co5.ratios)[co5.is.normal] <- str_replace(colnames(co5.ratios)[co5.is.normal], "[MP]", "N")
colnames(co8.ratios)[co8.is.normal] <- str_replace(colnames(co8.ratios)[co8.is.normal], "[MP]", "N")

co5.ploidy <- sapply(str_extract(colnames(co5.ratios), '[MPN]'), switch, P=co5.prim.ploidy, M=co5.met.ploidy, N=2)
co8.ploidy <- sapply(str_extract(colnames(co8.ratios), '[MPN]'), switch, P=co8.prim.ploidy, M=co8.met.ploidy, N=2)

co5.ploidyscaled <- sweep(co5.ratios, 2, co5.ploidy, '*')
co8.ploidyscaled <- sweep(co8.ratios, 2, co8.ploidy, '*')


#colsep <- which(c(0,diff(co5$chrom))!=0)
colsep <- NULL

plot_options <- list(
                     list(suffix="upto6", my_palette=c("dodgerblue4", "dodgerblue2", "white", "#FFCD2E", "#FF7701", "#CC2200", "#702F34")),
                     list(suffix="upto8", my_palette=c("dodgerblue4", "dodgerblue2", "white", "#FFCD2E", "#FF7701", "#CC2200", "#702F34", "#5A5E5A", "#517068"))
                )
for (these.options in plot_options)
{
attach(these.options)

col_breaks <- 0:length(my_palette) - .001

mat <- round(t(co5.ploidyscaled))
colnames(mat) <- co5$chrom
rownames(mat) <- apply(str_match(colnames(co5.ploidyscaled), "^C05([MPN])\\.(\\d+)\\.bl$")[,-1], 1, function(x) sprintf("%s-%s", x[1], x[2]))
species <- as.integer(as.factor(str_extract(colnames(co5.ploidyscaled), '[MPN]')))
png(sprintf("C05_heatmap_%s.png",suffix), width=2*2353-22,height=2*1093, units="px")
tryCatch(heatmap.3(as.matrix(mat), dendrogram="row", distfun = dist, hclustfun = function(x) hclust(x, method='ward'),
          RowSideColors=matrix(c("#DE9395", "#94BFFA", "#69D46D")[as.numeric(species)], nrow=1),
#          ColSideColors=matrix(c("firebrick","gold", "grey", "khaki1","lightblue","lightcoral","blue","red", "orange","magenta","burlywood4","brown","blue", "forestgreen","lightcoral", "deepskyblue","pink","green","palegreen4","saddlebrown","salmon","tan1","orchid")[as.numeric(colnames(mat))],ncol=1),
          ColSideColors=matrix(rep(c("#BFBFBF","#D4D4D4"),12)[as.numeric(colnames(mat))],ncol=1),
          rowsep=1:nrow(mat), sepcolor='white', sepwidth=c(30,.1),
          colsep=colsep,
          Colv=NA, Rowv=TRUE,notecol="black",col=my_palette,breaks=col_breaks,key=FALSE, keysize=0.5, density.info="none", trace="none", cexRow=6,cexCol=1.5,cex.main=10,cex.lab=10, symm=F,symkey=F,symbreaks=T, margins=c(25,15)),
finally=dev.off())

mat <- round(t(co8.ploidyscaled))
colnames(mat) <- co8$chrom
rownames(mat) <- apply(str_match(colnames(co8.ploidyscaled), "^([MPN])A\\.(\\d+)_$")[,-1], 1, function(x) sprintf("%s-%s", x[1], x[2]))
species <- as.integer(as.factor(str_extract(colnames(co8.ploidyscaled), '[MPN]')))
#my_palette <- colorRampPalette(c("blue", "white", "red"))(n = 999)
#col_breaks = c(seq(0,1,length=20),seq(1,2,length=320),seq(2,3,length=230),seq(4,5,length=380),seq(5,6,length=50))
col_breaks <- 0:length(my_palette) - .001
png(sprintf("C08_heatmap_%s.png",suffix), width=2*2353-21,height=2*1093, units="px")
tryCatch(heatmap.3(as.matrix(mat), dendrogram="row", distfun = dist, hclustfun = function(x) hclust(x, method='ward'),
          RowSideColors=matrix(c("#DE9395", "#94BFFA", "#69D46D")[as.numeric(species)], nrow=1),
#          ColSideColors=matrix(c("firebrick","gold", "grey", "khaki1","lightblue","lightcoral","blue","red", "orange","magenta","burlywood4","brown","blue", "forestgreen","lightcoral", "deepskyblue","pink","green","palegreen4","saddlebrown","salmon","tan1","orchid")[as.numeric(colnames(mat))],ncol=1),
          ColSideColors=matrix(rep(c("#BFBFBF","#D4D4D4"),12)[as.numeric(colnames(mat))],ncol=1),
          rowsep=1:nrow(mat), sepcolor='white', sepwidth=c(30,.1),
          colsep = colsep,
          Colv=NA, Rowv=TRUE,notecol="black",col=my_palette,breaks=col_breaks,key=FALSE, keysize=0.5, density.info="none", trace="none", cexRow=6,cexCol=1.5,cex.main=10,cex.lab=10, symm=F,symkey=F,symbreaks=T, margins=c(25,15)),
finally=dev.off())
detach(these.options)
}