library(stringr)
library(ggplot2)
library(reshape2)
library(RColorBrewer)

read.counts <- read.delim("read_counts.txt", sep=" ")
probabilities <- read.delim("probabilities.txt", sep=" ")
bridge.mutations <- c("CHN1", "FHIT", "APC", "ATP7B")
prim <- readLines("prim.txt")
met1 <- readLines("met1.txt")
met2 <- readLines("met2.txt")

matrix.images <- list()
groups <- list(prim=prim, met1=met1, met2=met2)
cube <- function(x) (1/2) * ((2*x-1)^3 + 1)
cube.inv <- function(y) (1/2) * (sign(2*y-1)*abs(2*y-1)^(1/3) + 1)
cube_trans <- scales:::trans_new("cube", cube, cube.inv)
for (i in 1:3)
{
  rc.dat <- melt(cbind(bridge.mutations, read.counts[,groups[[i]]]), id="bridge.mutations")
  p.dat <- melt(cbind(bridge.mutations, probabilities[,groups[[i]]]), id="bridge.mutations")
  rc.dat$value[is.na(rc.dat$value)] <- "0,0"
  rc.dat$value <- apply(str_match(rc.dat$value, "^(\\d*),(\\d*)$")[,-1], 1,
                        function(v) sprintf("%s, %s", v[1], v[2]))
  matrix.images[[names(groups)[i]]] <- ggplot() +
    geom_tile(aes(fill=as.double(value), x=bridge.mutations, y=variable), p.dat) +
    geom_text(aes(label=value, x=bridge.mutations, y=variable), rc.dat, size=3) +
    scale_fill_gradient2(low="#56B4E9", mid="#999999", high="#E69F00", midpoint=.5,
                         na.value="black", limits=c(0,1), name="probability of mutation",
                         trans=cube_trans) +
    theme(axis.text.x = element_text(size=10), axis.text.y=element_text(size=10), legend.text=element_text(size=6)) + 
    theme(axis.title.x=element_blank(), axis.title.y=element_blank())
}

for (name in names(matrix.images))
{
  png(sprintf("%s_vaf.png", name), width=(850/742)*6, height=3.8, units='in', res=600)
  tryCatch(print(matrix.images[[name]] + guides(fill=FALSE)), finally=dev.off())
}
png(sprintf("vaf_legend.png"), width=4, height=4, units='in', res=600)
tryCatch(print(matrix.images[[1]]), finally=dev.off())