#!/usr/bin/env Rscript
args = commandArgs(trailingOnly=TRUE)
print(args)

suppressPackageStartupMessages(library("dplyr"))
library("ggplot2")
library("stringr")
library("reshape2")

reorder_mtx <- function(m,repfile) {
    # Read ordering data and format it
  rep <- read.table(repfile, sep="\t")
  colnames(rep) <- c("family","group")
  rep$order <- 1:length(rep$family)
    # prepare columns and rows
  m_dfc <- as.data.frame(colnames(m))
  colnames(m_dfc) <- c("family")
  m_dfr <- as.data.frame(rownames(m))
  colnames(m_dfr) <- c("family")
    # reorder cols
  tmp <- merge(m_dfc,rep, by="family", all.x=TRUE)
  m <- m[,order(tmp$order)]
    # reorder rows
  tmp <- merge(m_dfr,rep, by="family", all.x=TRUE)
  m <- m[order(tmp$order),]
    # prepare cols
  m_dfc <- as.data.frame(colnames(m))
  colnames(m_dfc) <- c("family")
  m_dfc$order <- 1:length(m_dfc$family)
    # prepare rows
  m_dfr <- as.data.frame(rownames(m))
  colnames(m_dfr) <- c("family")
  m_dfr$order <- 1:length(m_dfr$family)
    # square them
  m_df <- merge(m_dfc,m_dfr, by="family", sort=FALSE)
  m <- m[,m_df$order.x]
  m <- m[m_df$order.y,]
  return(m)
}

#tab <- read.table(args[1], sep="\t", header=T, stringsAsFactors=F)
load(args[1])
if(args[2] == "TAD") {
  tab <- tab[which((tab$chr1 == tab$chr2) & (abs(tab$pos2-tab$pos1) > 5000) & (abs(tab$pos2-tab$pos1) < 1000000)),]
} else {
  if(args[2] == "DIST") {
    tab <- tab[which(((tab$chr1 == tab$chr2) & (abs(tab$pos2-tab$pos1) >= 1000000)) | (tab$chr1 != tab$chr2)),]
  }
}

heatmapFilename <- paste0("heatmap_",args[4],"_label_permutation_norm.",args[2],".png")

#simple basic heatmap of all interactions
if(length(colnames(table(tab$family1,tab$family2))) > 1 & length(colnames(table(tab$family1,tab$family2))) > 1) {

  #heatmap of all interactions; normalization uses an average of num permutations
  tab_sample <- table(sample(tab$family1),sample(tab$family2))
  num <- 12
  for(i in 1:num) {
    tab_sample <- tab_sample + table(sample(tab$family1),sample(tab$family2))
  }
  tab_sample <- tab_sample/num

  mtx_observed <- reorder_mtx(table(tab$family1,tab$family2), args[3])
  mtx_sampled <- reorder_mtx(tab_sample, args[3])

  stopifnot(colnames(mtx_observed)==colnames(mtx_sampled))
  stopifnot(rownames(mtx_observed)==rownames(mtx_sampled))

  nfamilies<-length(rownames(mtx_observed))
  mtx_sampled_normalized <- matrix("", ncol = nfamilies, nrow = nfamilies)
  colnames(mtx_sampled_normalized) <- colnames(mtx_observed)
  rownames(mtx_sampled_normalized) <- rownames(mtx_observed)

  for (i in 1:nfamilies) {
    for (j in 1:nfamilies) {
        
        sampled_value <- as.numeric(mtx_sampled[i,j])
        observed_value <- as.numeric(mtx_observed[i,j])

        if (sampled_value==0 | observed_value==0) { 
            mtx_sampled_normalized[i,j] <- NA #division by 0 not possible; infinity values should be set to NA as well
        } else {
            normalized_value <- observed_value/sampled_value
            mtx_sampled_normalized[i,j] <- normalized_value
        }
    }
  } 
    
  df_normalized <- melt(mtx_sampled_normalized)
  df_normalized$log <- log(as.numeric(as.character(df_normalized$value)))

  png(file = heatmapFilename, units="in", width=6, height=6, res=300)
  g <- ggplot(df_normalized, aes(Var1,Var2,fill=log)) 
  g <- g + geom_tile() + theme(axis.text.x=element_text(angle=90,hjust=1,vjust=0.5)) + scale_fill_distiller(palette = 'Spectral', na.value="gray90") + ggtitle(args[2]) + xlab("Repeat family") + ylab("Repeat family")
  print(g)

  dev.off()
  save(mtx_sampled_normalized, file="matrix_label_permutation_norm.Rdata")

} else {
  #not enough data to plot the heatmap
  png(file = heatmapFilename, units="in", width=6, height=6, res=300)
  par(mar = c(0,0,0,0),cex=0.5)
  plot(c(0, 1), c(0, 1), ann = F, bty = 'n', type = 'n', xaxt = 'n', yaxt = 'n')
  text(x = 0.5, y = 0.5, paste("There was not enough data available to create a heatmap."), cex = 1.6, col = "black")
  dev.off()  
}



