library(tidyverse)
library(optparse)

#Version 0.1.1 Jia Jinbu 2021.02.03
#add end3_alignment_scoreend3_alignment_score
#Jia Jinbu, 2020.09.15

#Note: This script will filter out duplicated read_core_id lines

#			exon1    intron1  exon2   intron2   exon3  intron3   exon4
#mRNA     >>>>>>>>>>>.......>>>>>>>>>.........>>>>>>>>>>......>>>>>>>>>>>
#read1:                  >>>>>>>>>>>>.........>>>>>>>>>>>>>
#read2:       >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#read3:                         >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#             first_span_intron_num      last_span_intron_num
#read1:                 1                        2
#read2:                 1                        2
#read3:                 2                        3                

#type:
#	 "elongating"
#	 "splicing_intermediate"
#	 "elongating_3_mapping_low_accuracy"
#	 "elongating_5lost"
#	 "polya"
#	 "polya_3_not_in_last_exon"
#	 "polya_5lost"

#some columns in fileout_detail 
#end5ss_type: 1/0  whether is splicing intermeidate

#end_type:
#    mRNA                >>>>>>>>>>>.......>>>>>>>>>.........>>>>>>>>>>      
#    read  end_type: 0                  >>>>>>>>>>>>>>>
#          end_type: 1     >>>>>>>>>>>>>>>>>>>>>>
#          end_type: 2                  >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#          end_type: 3   >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
#
#        #end5: read overlapped with the the first exon
#        #end3: read overlapped with the the last exon
#        #	     noend5 end5
#        #noend3      0    1
#        #end3        2    3
#polyA_type:
#      0:  no polyA. polyA_length < 15  
# 	   1:  has polyA. polyA_length >= 15
#end_polyA_type:
#	        			end_type=0 end_type=1 end_type=2 end_type=3
#    polyA_type = 0			0			1		  2			3
#    polyA_type = 1			4			5		  6			7

parser <- OptionParser()
parser <- add_option(parser, c("-i", "--inreadinfo"), type="character", 
						help="Read info file generated by extract_read_info.py")
parser <- add_option(parser, c("-a", "--inadapter"), type="character", default="",
						help="Aadapter info file generated by adapterFinder.py. Required only for Nanopore data.")
parser <- add_option(parser, c("-p", "--inpolya"), type="character", 
						help="PolyA info generated by PolyACaller.py (Nanopore) or pacbio_find_polyA.py (PacBio)")
parser <- add_option(parser, c("-t", "--type"), type="character", default="Nanopore",
						help="Nanopore|PacBio, default[Nanopore]")
parser <- add_option(parser, c("-o", "--out"), type="character", 
						help="Output File")
opt <- parse_args(parser)

filein_retention <- opt$inreadinfo
filein_adapter <- opt$inadapter
filein_polyA <- opt$inpolya
data_type <- ifelse(substr(opt$type, 1, 1) %in% c("N", "n"), "Nanopore", "PacBio") 
fileout <- opt$out
#fileout_stat <- args[6]

MIN_POLYA_LENGTH <- 15
MAX_SPLICE_INTERMEDIATE_LEFT_LENGTH <- -10
MAX_SPLICE_INTERMEDIATE_RIGHT_LENGTH <- 10
READ_THROUGH_LENGTH <- 10000
END3_ALIGNMENT_SCORE_LIMIT <- c(-5, 5)

extract_end3_score_value <- function(data){
	if ("r_align_start" %in% colnames(data)){
		# nanopore data
		#for elongating
	    a <- (data$r_align_start - data$r_primer_start) - data$genome_align_end
	    b <- data$genome_align_start - (data$f_align_end + data$f_primer_start)
	    c <- data$rna_strand==data$read_strand
	    v <- b
	    v[c] <- a[c]
		#for polya
	    a1 <- data$polya_start_base - data$genome_align_end - 1
	    b1 <- data$genome_align_start - data$polya_end_base - 1
	    c1 <- data$rna_strand==data$read_strand
	    v1 <- b1
	    v1[c1] <- a1[c1]
		
		c2 <- data$polyA_type==1
		v[c2] <- v1[c2]
	} else {
		# pacbio data
		v <- rep(0, nrow(data))
	}
	return(v)
}

check_3_end_high_mapping <- function(data, limit=c(-5, 5)){
     v <- extract_end3_score_value(data)
     f <- (v >= limit[1]) & (v <= limit[2])
     return(f)
 }
 
#1. merge primer, retention, polyA file to `data`
#filter out duplicated read_core_id lines
if (data_type == "Nanopore"){
	primer <- read_tsv(filein_adapter, na="NA")
	primer <- primer %>% select(-read_align_strand) %>% filter(primer_score<=2)
	duplicated_read_core_id <- primer$read_core_id[duplicated(primer$read_core_id)]
	primer <- primer[!(primer$read_core_id %in% duplicated_read_core_id), ]

	retention <- read_tsv(filein_retention, na="NA")
	polyA <- read_tsv(filein_polyA, na="NA")

	polyA <- select(polyA, -polya_type)
	duplicated_read_core_id <- polyA$read_core_id[duplicated(polyA$read_core_id)]
	polyA <- polyA[!(polyA$read_core_id %in% duplicated_read_core_id), ]

	data <- inner_join(primer, polyA, by="read_core_id")
	gene_read_num <- nrow(retention)
	data <- inner_join(retention, data, by="read_core_id")
	gene_read_has_primerR_num <- nrow(data)
	data <- data %>% filter(mRNA_strand==rna_strand) 
} else { #"PacBio"
	retention <- read_tsv(filein_retention, na="NA")
	polyA <- read_tsv(filein_polyA, na="NA")
	#filter out duplicated read_core_id lines
	duplicated_read_core_id <- polyA$read_core_id[duplicated(polyA$read_core_id)]
	polyA <- polyA[!(polyA$read_core_id %in% duplicated_read_core_id), ]
	gene_read_num <- nrow(retention)
	data <- inner_join(retention, polyA, by="read_core_id")
	gene_read_has_primerR_num <- nrow(data)
	#for pacbio
	data <- data %>% filter(mRNA_strand==read_strand)
	gene_read_strand_right <- nrow(data)
}

#2. calculate end_polyA_type
f <- data$polya_length >= MIN_POLYA_LENGTH
data$polyA_type <- 0
data$polyA_type[f] <- 1
data$end_polyA_type <- data$end_type
data$end_polyA_type[f] <- data$end_type[f] + 4

#3. find splicing intermeidate
data$end5ss_type <- 0
f0 <- data$l_feature_type == "exon"
f1 <- (f0) & (data$l_pos3 > MAX_SPLICE_INTERMEDIATE_LEFT_LENGTH) & (data$l_feature_num != (data$mRNA_intron_num+1))
f2 <- (!f0) & (data$l_feature_length + data$l_pos3 <= MAX_SPLICE_INTERMEDIATE_RIGHT_LENGTH)
data$end5ss_type[f1 | f2] <- 1

#4. type
if (data_type == "Nanopore"){	
	data$end3_high_mapping <- check_3_end_high_mapping(data)
	data$low_accuracy_3end_mapped <- !data$end3_high_mapping
	data$type <- ""
	data$type[(data$end_polyA_type %in% c(1, 3)) & (data$end3_high_mapping) & (data$end5ss_type == 0)] <- "elongating"
	data$type[(data$end_polyA_type %in% c(1, 3)) & (data$end3_high_mapping) & (data$end5ss_type == 1)] <- "splicing_intermediate"
	data$type[(data$end_polyA_type %in% c(1, 3)) & (!data$end3_high_mapping) ] <- "elongating_3_mapping_low_accuracy"
	data$type[(data$end_polyA_type %in% c(0, 2)) ] <- "elongating_5lost"
	data$type[data$end_polyA_type == 7] <- "polya"
	data$type[data$end_polyA_type == 5] <- "polya_3_not_in_last_exon"
	data$type[data$end_polyA_type %in% c(4, 6)] <- "polya_5lost"
	data$end3_high_mapping <- NULL
} else { #"PacBio"
	data$type <- ""
	data$type[(data$end_polyA_type %in% c(1, 3)) & (data$end5ss_type == 0)] <- "elongating"
	data$type[(data$end_polyA_type %in% c(1, 3)) & (data$end5ss_type == 1)] <- "splicing_intermediate"
	data$type[(data$end_polyA_type %in% c(0, 2)) ] <- "elongating_5lost"
	data$type[data$end_polyA_type == 7] <- "polya"
	data$type[data$end_polyA_type == 5] <- "polya_3_not_in_last_exon"
	data$type[data$end_polyA_type %in% c(4, 6)] <- "polya_5lost"
}

#5. filter out reads which mRNA_pos and rel_mRNA_pos is not consistent
tmp_flag = (data$mRNA_pos5==data$rel_mRNA_pos5) & (data$mRNA_pos3 == data$rel_mRNA_pos3)
gene_read_strand_right <- sum(tmp_flag)
gene_read_not_consistent_rel_mRNA_pos <- nrow(data) - gene_read_strand_right
data <- data[tmp_flag,]
		
write_tsv(data, fileout)

#data_simple <- data %>% select(read_core_id, mRNA, mRNA_intron_num, type, 
#					    rel_mRNA_pos5, rel_mRNA_pos3, retention_introns, 
#						polya_length, first_span_intron_num, last_span_intron_num)
#write_tsv(data_simple, fileout_simple)
#stat <- tibble(gene_read_num=gene_read_num, 
#				  gene_read_has_primerR_num=gene_read_has_primerR_num, 
#				  gene_read_not_consistent_rel_mRNA_pos=gene_read_not_consistent_rel_mRNA_pos,
#				  gene_read_strand_right=gene_read_strand_right
#				  )
#print(stat)
#write_tsv(stat, fileout_stat)