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

library("stringr") 
suppressPackageStartupMessages(library(gplots))
suppressPackageStartupMessages(library(dplyr))

#EXAMPLE: SRR6765292_5part00_galaxy.inClusters.tab
clusters_labelled <- read.table(args[1],sep="\t") #formatted cluster annotation

#EXAMPLE: CLUSTER_TABLE.csv
cluster <-read.table(args[2], sep="\t", skip=5, header=TRUE) #output from the RepeatExplorer

#EXAMPLE: blast_formatted.out
blast <- as.data.frame(read.table(args[3], sep="\t")) #output from the blast


blast_table <- table(blast[,c(1,3)])
blast_count <- as.data.frame(blast %>% count(V1)) #count number of hits per cluster
colnames(blast_count) <- c("cluster","max")

get_max_families <- function() {
   result <- colnames(blast_table)[apply(blast_table,1,which.max)]
   return(result)
}

blast_count$family <- get_max_families()
blast_count$cluster_num <- as.numeric(unlist(lapply(blast_count$cluster,
str_extract, "[0-9]+")))

# expand Repeat Explorer annotations to unannotated clusters
add_max_families <- function() {
   # prepare an empty new column to add to cluster data frame
   result <- rep("",length(cluster$Cluster))
   # for each row in cluster (cluster reported in CLUSTER_TABLE.tsv)
   for(i in 1:length(cluster$Cluster)) {
     # if the cluster seems to be annotated
     if((as.character(cluster$Automatic_annotation[i]) != "All") &
(as.character(cluster$Automatic_annotation[i]) !=
"All/repeat/mobile_element/Class_I/LTR")) {
       # use the last word after the "/" separator as annotation
       result[i] <-
unlist(str_split(as.character(cluster$Automatic_annotation[i]),"/"))[length(unlist(str_split(as.character(cluster$Automatic_annotation[i]),"/")))]
     # otherwise try to annotate the cluster
     } else {
       # using blast results obtained against classification database for the given cluster
       # if they exist
       if(length(which(blast_count$cluster_num == i)) > 0) {
         # and there are at least three different protein domains that were top hit
         if(blast_count$max[which(blast_count$cluster_num == i)] > 2) {
           # annotate the cluster with this family
           result[i] <- blast_count$family[which(blast_count$cluster_num
== i)]
         }
       }
     }
   }
   return(result)
}

# Add Final_annotation column to clusters
cluster$Final_annotation <- add_max_families()
#cluster$Final_annotation<-sapply(cluster$Automatic_annotation,function(x) {basename(as.character(x))})

colnames(clusters_labelled) <-
c("read1","cluster1","read2","cluster2")

#ALL THE READ PAIRS ARE ASSIGNED TO CLUSTERS (OR NOT)
print(dim(clusters_labelled))

#WE NEED TO DROP READ PAIRS IN WHICH BOTH READS ARE ASSIGNED TO NO CLUSTER
clusters_labelled <- clusters_labelled[!with(clusters_labelled,is.na(cluster1) & is.na(cluster2)),]
print(dim(clusters_labelled))

#WE ONLY WANT TO KEEP READ PAIRS IN WHICH AT LEAST ONE READ MAPS TO A CLUSTER OF INTEREST
clusters_to_keep<-paste0("CL",cluster$Cluster)
my_regex <- regex(paste("\\b", clusters_to_keep, "\\b", sep = "", collapse = "|"))
clusters_labelled<-clusters_labelled %>% filter(str_detect(cluster1, my_regex) | str_detect(cluster2, my_regex))

cluster$CL<-paste0("CL",cluster$Cluster)
summaryTable<-merge(clusters_labelled,cluster[,c("CL","Final_annotation")],by.x=c("cluster1"),by.y = c("CL"),all.x=TRUE)
names(summaryTable)[names(summaryTable) == "Final_annotation"] <- "family1"
summaryTable<-merge(summaryTable,cluster[,c("CL","Final_annotation")],by.x=c("cluster2"),by.y = c("CL"),all.x=TRUE)
names(summaryTable)[names(summaryTable) == "Final_annotation"] <- "family2"
summaryTable[summaryTable == ""] <- NA 

print(colnames(summaryTable))
print(head(summaryTable))


emptyColumn<-rep("NA",nrow(summaryTable))
#format into the .tab format that is used for the reference-based branch
tab<-as.data.frame(cbind(gsub("\\/1", "", summaryTable$read1),emptyColumn,emptyColumn,summaryTable$family1,emptyColumn,emptyColumn,summaryTable$family2))

colnames(tab)<-c("read_id","pos1","chr1","family1","pos2","chr2","family2")

#DROP THE READ PAIR IF ONE OF THE FAMILIES IS UNDEFINED
tab<-tab[!is.na(tab$family1), ]
tab<-tab[!is.na(tab$family2), ]

write.table(tab, file = paste0(args[1],".formatted.tab"), sep = "\t", row.names = FALSE, col.names = TRUE, quote = FALSE) #write the output



