library(tidyverse)
library(optparse)

# Jia Jinbu. 2020.09.15

parser <- OptionParser()
parser <- add_option(parser, c("-i", "--inrelpos"), type="character", 
						help="Intron splicing and Pol II position file generated by prepare_data_for_splice_kinetics.py")
parser <- add_option(parser, c("-a", "--inreadinfo"), type="character",
						help="Input Read info file generated by merge_read_info.R. This file was used to extract the read_core_id which type is elongating")
parser <- add_option(parser, c("-s", "--inintron"), type="character", default="",
						help="Input select intron file. only introns listed in this file were used for calcualte. one line one intron_id, no column name.")
parser <- add_option(parser, c("-o", "--out"), type="character", 
						help="Output File splicing kinetics")
parser <- add_option(parser, c("-p", "--pdf"), type="character", 
						help="Output pdf file for plot splicing kinetics")

opt <- parse_args(parser)
file_rel_pos = opt$inrelpos
file_elongating = opt$inreadinfo
fileout_stat =opt$out
fileout_pdf = opt$pdf
file_select_intron_ids = opt$inintron

stat_intron_splice <- function(data, bin_length = 20, origin_data=T){
  to_bin <- function(x, l=20){
    (x - 1) %/% l * l + (1 + l)/2
  }
  if (origin_data) {
    spliced_stat <- data %>%
    group_by(rel_pos) %>% 
    summarise(unspliced=sum(retention), total=n()) 
    spliced_stat$pos_bin <- spliced_stat$rel_pos
  } else {
    spliced_stat <- data
  }
  if (bin_length != 1){
    spliced_stat$pos_bin <- to_bin(spliced_stat$pos_bin, bin_length)
  }
  spliced_stat <- spliced_stat %>%
    mutate(spliced_ratio = 1 - unspliced/total) %>%
    group_by(pos_bin) %>%
    summarise(mean_spliced_ratio=mean(spliced_ratio),
              unspliced = sum(unspliced),
              total = sum(total),
              spliced = total - unspliced,
              sd=sd(spliced_ratio),
              spliced_ratio = spliced/total)
  return(spliced_stat)
}

plot_intron_splicing_kinetics <- function(stat_data, 
										  one_sample=T, 
										  col="red",
										  min_size=6,
                                          legend.position.x=0.2,
                                          legend.position.y=0.8){
	#pdf("plot_new/elongating.rel_pos.spliced.pdf" , 1.8, 1.6)
	#p <- plot_intron_splicing_kinetics(stat_data)
	#print(p)
	#dev.off()
	if (one_sample) {
		p <- ggplot(stat_data, aes(x=pos_bin, y=spliced_ratio)) + 
				geom_point(size=0.1, col=col) +
				geom_line(size=0.2, col=col) 
	} else {
		p <- ggplot(stat_data, aes(x=pos_bin, y=spliced_ratio, col=sample)) +
				geom_point(size=0.1) + 
				geom_line(size=0.2) 
	}
	p <- p + 
	  coord_cartesian(xlim=c(0,3000), ylim=c(0,1)) + 
	  scale_y_continuous(labels=scales::percent) +
	  xlab("Distance from 3'SS (nt)") + 
	  ylab("Percentage of spliced") +
	  theme_classic() +
	  theme(text=element_text(size=min_size),
	        legend.position = c(legend.position.x, legend.position.y), 
	        legend.key.size = unit(0.7*min_size/6,"line"),
	        legend.title=element_blank(),
	        legend.text = element_text(size=min_size),
	        legend.margin=margin(0,0,0,0),
	        legend.box.margin=margin(-10,0,-10,-10),
	        axis.line.y = element_line(color="black", size = 0.2),
	        axis.line.x = element_line(color="black", size = 0.2),
	        panel.margin = unit(0.1, "lines"),
	        #axis.text.x = element_text(angle = 315, vjust = 0.5, hjust = 0)
	        axis.text.x = element_text(size=min_size, color="black")
	  ) 
	return(p)
}

read_elongating_ids <- function(filein){
	data <- read_tsv(filein, na="NA")
	return(data$read_core_id[data$type=="elongating"])
}

read_select_intron_ids <- function(filein){
	return(read_tsv(filein, col_names = F)$X1)
}

BIN_LENGTH <- 20
PDF_WIDTH <- 1.8
PDF_HEIGHT <- 1.86

#prepare data
rel_pos <- read_tsv(file_rel_pos)
select_read_core_ids <- read_elongating_ids(file_elongating)
select_flag <- rel_pos$read_core_id %in% select_read_core_ids
if (file_select_intron_ids != "") {
	select_intron_ids = read_select_intron_ids(file_select_intron_ids)
	select_flag <- select_flag & (rel_pos$intron_id %in% select_intron_ids)
}
select_rel_pos <- rel_pos[select_flag, ]

#stat and plot
stat_data <- stat_intron_splice(select_rel_pos, BIN_LENGTH)
write_tsv(stat_data, fileout_stat)
pdf(fileout_pdf , PDF_WIDTH, PDF_HEIGHT)
print(plot_intron_splicing_kinetics(stat_data))
dev.off()