#!/usr/bin/env Rscript
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)
}

load(args[1]) #args[1] is family table, args[2] name of the heatmap type to make (ALL, TAD, DIST)

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)),]
    #print(nrow(tab))
  }
}


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

mtx_normalized <- table(as.character(tab$family1),as.character(tab$family2))
mtx_shuffled_normalized <- table(as.character(tab_shuffled$family1),as.character(tab_shuffled$family2))

#find shared columns and rows
shared_fam_cols<-intersect(colnames(mtx_normalized),colnames(mtx_shuffled_normalized))
shared_fam_rows<-intersect(rownames(mtx_normalized),rownames(mtx_shuffled_normalized))
shared_fams<-intersect(shared_fam_cols,shared_fam_rows)

print(length(shared_fams))

  heatmapFilename <- paste0("heatmap_reference_genome_reshuffling_norm.",args[2],".png")
  print(heatmapFilename)
  
if (length(shared_fams)>=2) {
  #subset to shared columns and rows
  mtx_normalized <- mtx_normalized[shared_fams,shared_fams]
  mtx_shuffled_normalized <- mtx_shuffled_normalized[shared_fams,shared_fams]

  print(mtx_normalized)
  print(mtx_shuffled_normalized)
  #list_of_samples_subsetted <- lapply(list_of_samples,function(x) {x[shared_fams,shared_fams]})

  #the dimensions between the data and control must match
  stopifnot(colnames(mtx_normalized)==colnames(mtx_shuffled_normalized))
  stopifnot(rownames(mtx_normalized)==rownames(mtx_shuffled_normalized))

  for (i in 1:length(rownames(mtx_normalized))) {
    for (j in 1:length(colnames(mtx_normalized))) {
        
        value <- mtx_normalized[i,j]
        shuffled_value <- mtx_shuffled_normalized[i,j]

        if (shuffled_value==0 | value==0) {
            mtx_normalized[i,j] <- NA #division by 0 not possible, infinity should also become NA
        } else {
            normalized_value <- mtx_normalized[i,j]/shuffled_value
            mtx_normalized[i,j] <- normalized_value
        }
    }

  } 

  mtx_normalized<-reorder_mtx(mtx_normalized, args[4]) #reorder the names based on the external file
  df_normalized <- melt(as.data.frame(mtx_normalized))
  df_normalized$log <- log(df_normalized$value)
  save(df_normalized, file = "stuff.RData")
  print(df_normalized)

  
  #heatmap of all interactions; normalized by shuffled annotation
  png(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() 

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


