# 
# Reads a table and a shuffled table from extract_pairs.pl then uses the shuffled values to normalized
# observed counts and creates the respective matrices (all, TAD, DIST)
# 
library(stringr)
library(reshape)
library(ggplot2)

# read in tables
tab <- read.table("Osativa_323_v7.0_fragments_toHiC_All.tab", sep="\t", header=T)
tab_shuffled <- read.table("Osativa_323_v7.0_fragments_toHiC_All_shuffled.tab", sep="\t", header=T)
# prepare tables
colnames(tab) <- c("read1","pos1","len1","chr1","annot1","read2","pos2","len2","chr2","annot2")
tab$family1 <- unlist(lapply(str_extract(tab$annot1,"annot=[^;]+"),substring,7))
tab$family2 <- unlist(lapply(str_extract(tab$annot2,"annot=[^;]+"),substring,7))
colnames(tab_shuffled) <- c("read1","pos1","len1","chr1","annot1","read2","pos2","len2","chr2","annot2")
tab_shuffled$family1 <- unlist(lapply(str_extract(tab_shuffled$annot1,"annot=[^;]+"),substring,7))
tab_shuffled$family2 <- unlist(lapply(str_extract(tab_shuffled$annot2,"annot=[^;]+"),substring,7))
tab_tad <- tab[which((tab$chr1 == tab$chr2) & (tab$pos2-tab$pos1 > 5000) & (tab$pos2-tab$pos1 < 1000000)),]
tab_dist <- tab[which((tab$chr1 == tab$chr2) & (tab$pos2-tab$pos1 > 1000000) | (tab$chr1 != tab$chr2)),]
tab_tad_shuffled <- tab_shuffled[which((tab_shuffled$chr1 == tab_shuffled$chr2) & (tab_shuffled$pos2-tab_shuffled$pos1 > 5000) & (tab_shuffled$pos2-tab_shuffled$pos1 < 1000000)),]
tab_dist_shuffled <- tab_shuffled[which((tab_shuffled$chr1 == tab_shuffled$chr2) & (tab_shuffled$pos2-tab_shuffled$pos1 > 1000000) | (tab_shuffled$chr1 != tab_shuffled$chr2)),]

# create matrices of correct dimensions (filled with wrong data!)
mtx_normalized <- table(as.character(tab$family1),as.character(tab$family2))
mtx_tad_normalized <- table(as.character(tab_tad$family1),as.character(tab_tad$family2))
mtx_dist_normalized <- table(as.character(tab_dist$family1),as.character(tab_dist$family2))
mtx_shuffled_normalized <- table(as.character(tab_shuffled$family1),as.character(tab_shuffled$family2))
mtx_tad_shuffled_normalized <- table(as.character(tab_tad_shuffled$family1),as.character(tab_tad_shuffled$family2))
mtx_dist_shuffled_normalized <- table(as.character(tab_dist_shuffled$family1),as.character(tab_dist_shuffled$family2))

# calculate normalized values (now filled with correct data)
for (i in 1:length(rownames(mtx_normalized))) {
  for (j in 1:length(colnames(mtx_normalized))) {
    mtx_normalized[i,j] <- mtx_normalized[i,j]/mtx_shuffled_normalized[i,j]
  }
} 
for (i in 1:length(rownames(mtx_tad_normalized))) {
  for (j in 1:length(colnames(mtx_tad_normalized))) {
    mtx_tad_normalized[i,j] <- mtx_tad_normalized[i,j]/mtx_tad_shuffled_normalized[i,j]
  }
}
for (i in 1:length(rownames(mtx_dist_normalized))) {
  for (j in 1:length(colnames(mtx_dist_normalized))) {
    mtx_dist_normalized[i,j] <- mtx_dist_normalized[i,j]/mtx_dist_shuffled_normalized[i,j]
  }
}

df_normalized <- melt(data.frame(mtx_normalized))
df_tad_normalized <- melt(data.frame(mtx_dist_normalized))
df_dist_normalized <- melt(data.frame(mtx_dist_normalized))
df_normalized$log <- log(df_normalized$value)+0.1
df_tad_normalized$log <- log(df_tad_normalized$value)+0.1
df_dist_normalized$log <- log(df_dist_normalized$value)+0.1

ggplot(df_normalized, aes(Var1,Var2,fill=log)) + geom_tile() + theme(axis.text.x=element_text(angle=90,hjust=1,vjust=0.5)) + scale_fill_distiller(palette = 'Spectral') + ggtitle("ALL") + xlab("Repeat Family") + ylab("Repeat family")
ggplot(df_tad_normalized, aes(Var1,Var2,fill=log)) + geom_tile() + theme(axis.text.x=element_text(angle=90,hjust=1,vjust=0.5)) + scale_fill_distiller(palette = 'Spectral') + ggtitle("TAD") + xlab("Repeat Family") + ylab("Repeat family")
ggplot(df_dist_normalized, aes(Var1,Var2,fill=log)) + geom_tile() + theme(axis.text.x=element_text(angle=90,hjust=1,vjust=0.5)) + scale_fill_distiller(palette = 'Spectral') + ggtitle("DIST") + xlab("Repeat Family") + ylab("Repeat family")

