#!/usr/bin/env Rscript

#
# Script: generate_heatmap_reference_joint_probability.R
#
# Author: Matej Lexa 
# Year: 2021
#
# This script is part of the HiC-TE Analysis pipeline for Nextflow
#
# It reads a table of interactions at family resolution, creates 
# a matrix with interaction counts from it. After normalization
# against joint probability estimates it sources reorder_matrix
# to arrange rows and columns using an external ordering file 
# and plots a heatmap.
#

args = commandArgs(trailingOnly=TRUE)
print(args)

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

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)
}

# Read in the interaction table and filter it from distance depending on tye in input
load(args[1]) #args[1] is family table, args[2] is the type of heatmap to make (ALL, TAD or DIST)
if(args[2] == "TAD") {
  tab <- tab[which((tab$chr1 == tab$chr2) & (tab$pos2-tab$pos1 > 5000) & (tab$pos2-tab$pos1 < 1000000)),]
} else {
  if(args[2] == "DIST") {
    tab <- tab[which((tab$chr1 == tab$chr2) & (tab$pos2-tab$pos1 >= 1000000) | (tab$chr1 != tab$chr2)),]
  }
}

# A function to normalize matrix
norm1 <-function(tab) {
  mtx <- table(tab$family1,tab$family2)
  vec1 <- table(tab$family1)
  vec2 <- table(tab$family2)
  mtx_sum <- sum(mtx)
  
  for(i in 1:length(vec1)) {
    for(j in 1:length(vec2)) {
      if((vec1[i] > 0) && (vec2[j] > 0) && (mtx[i,j] > 0)) {
        both_families_expected<-(vec1[i]/mtx_sum)*(vec2[j]/mtx_sum)*mtx_sum #counts, not probabilities
        both_families_observed<-mtx[i,j] #counts, not probabilities
        mtx[i,j] <- both_families_observed/both_families_expected
      } else {
        mtx[i,j] <- NA
      }
    }
  }
  return(mtx)
}


#simple basic heatmap of all interactions normalized by comparison to joint probability for two families calculated from family counts
if(length(colnames(table(tab$family1,tab$family2))) > 1 & length(colnames(table(tab$family1,tab$family2))) > 1) {
  
  png(file = paste0("heatmap_reference_joint_probability_norm.",args[2],".png"), units="in", width=6, height=6, res=300)
  mtx_all<-reorder_mtx(norm1(tab), args[3])
  save(mtx_all, file="matrix_joint_probability.Rdata")
  df_normalized <- melt(data.frame(mtx_all))
  df_normalized$log <- log(df_normalized$value)
  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)

} else {
  #not enough data to plot the heatmap
  png(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()  
}
dev.off()





