#! /usr/bin/Rscript

# Plot tRNA coverage by isoacceptor
# Accepts coverage files generated by getCoverage.py

suppressMessages(library(ggplot2))
suppressMessages(library(RColorBrewer))
suppressMessages(library(plyr))
suppressMessages(library(gridExtra))

args = commandArgs(trailingOnly = TRUE)

aa_groups = data.frame(aa = c("Gly","Ala","Val","Leu","Met","iMet","Ile","Ser","Thr","Cys","SeC","Pro","Asn","Gln","Phe","Tyr","Trp","Lys","Arg","His","Asp","Glu","Und","Undet","Sup"), 
	group = c(rep("Nonpolar, aliphatic",7), rep("Polar, uncharged",7), rep("Aromatic",3), rep("Positively charged",3), rep("Negatively charged",2), rep("Other",3)))

if (length(args)==0) {
	stop("At least one argument must be supplied (input file).n", call.=FALSE)

} else if (length(args)>0) {
	
	out_dir = args[3]
	mito_trnas = args[5]
	if (mito_trnas == ''){
		mito_trnas = NA
	}
	sorted_aa = args[4]
	sorted_aa = unlist(strsplit(sorted_aa, "_"))

	## Aggregated coverage plots for all Isoacceptors
	cov_byaa = read.table(args[2], header=TRUE, sep = "\t")
	cov_byaa = cov_byaa[complete.cases(cov_byaa),]
	cov_byaa$bam = gsub(".unpaired_uniq.bam","",cov_byaa$bam)
	cov_byaa$bam = gsub("(.*/).*?","\\2",cov_byaa$bam)
	cov_byaa = cov_byaa[!grepl("eColi", cov_byaa$aa),]
	cov_byaa$aa = factor(cov_byaa$aa, levels = sorted_aa)
	cyt_colourCount = length(unique(subset(cov_byaa$aa, !grepl("mito", cov_byaa$aa) & !grepl("plastid", cov_byaa$aa) & !grepl("nmt", cov_byaa$aa))))
	mit_colourCount = length(unique(subset(cov_byaa$aa, grepl("mito", cov_byaa$aa) | grepl("plastid", cov_byaa$aa) | grepl("nmt", cov_byaa$aa))))
	facetCount = length(unique(cov_byaa$bam))
	getPalette = colorRampPalette(brewer.pal(10, 'Paired'))
	
	cov_byaa_norm = ggplot(subset(cov_byaa, !grepl("mito", cov_byaa$aa) & !grepl("plastid", cov_byaa$aa) & !grepl("nmt", cov_byaa$aa)), aes(x = bin, y = cov_norm, fill = aa, group = aa)) + 
		geom_bar(stat = "identity", alpha = 0.8) + 
		facet_wrap(~bam, ncol = 4) + 
		xlab("Gene (%)") + 
		ylab("Normalised coverage (coverage/uniquely mapped reads)") + 
		labs(fill = "Isotype") + 
		guides(fill=guide_legend(ncol=2)) + 
		scale_fill_manual(values = getPalette(cyt_colourCount)) + 
		theme_bw()
	ggsave(paste(out_dir, "coverage_byaa_norm.pdf", sep = ''), cov_byaa_norm, height = ceiling(facetCount/4) * 4, width = 14)

	if (!is.na(mito_trnas)){
		mito_cov_byaa = subset(cov_byaa, grepl("mito", cov_byaa$aa) | grepl("plastid", cov_byaa$aa) | grepl("nmt", cov_byaa$aa))
		if (nrow(mito_cov_byaa) != 0 ) {
			mitocov_byaa_norm = ggplot(mito_cov_byaa, aes(x = bin, y = cov_norm, fill = aa, group = aa)) +
				geom_bar(stat = "identity", alpha = 0.8) + 
				facet_wrap(~bam, ncol = 4) + 
				xlab("Gene (%)") + ylab("Normalised coverage (coverage/uniquely mapped reads)") + 
				labs(fill = "Isotype") + 
				guides(fill=guide_legend(ncol=2)) +
				scale_fill_manual(values = getPalette(mit_colourCount)) + 
				theme_bw()
			ggsave(paste(out_dir, "Organelle_coverage_byaa_norm.pdf", sep = ''), mitocov_byaa_norm, height = ceiling(facetCount/4) * 4, width = 14)	
		}
		
	}

	cyto_cov_byaa = subset(cov_byaa, !grepl("mito", cov_byaa$aa) & !grepl("plastid", cov_byaa$aa) & !grepl("nmt", cov_byaa$aa))
	cyto_cov_byaa_sum = aggregate(x = cyto_cov_byaa$cov_norm, by = list(bin = cyto_cov_byaa$bin, bam = cyto_cov_byaa$bam), FUN = sum)
	cyto_scale_factors = cyto_cov_byaa_sum[which(cyto_cov_byaa_sum$bin == 96),] # 96 is second last bin of 4%
	cyto_cov_byaa$cov_norm_scaled = NA
	for (i in 1:nrow(cyto_cov_byaa)){
		cyto_cov_byaa[i,'cov_norm_scaled'] = cyto_cov_byaa[i,'cov_norm']/cyto_scale_factors[which(cyto_cov_byaa[i,'bam'] == cyto_scale_factors$bam),3]
	}

	if (!is.na(mito_trnas) && nrow(mito_cov_byaa) != 0){
		mito_cov_byaa_sum = aggregate(x = mito_cov_byaa$cov_norm, by = list(bin = mito_cov_byaa$bin, bam = mito_cov_byaa$bam), FUN = sum)
		mito_scale_factors = mito_cov_byaa_sum[which(mito_cov_byaa_sum$bin == 96),] # 96 is second last bin of 4%
		mito_cov_byaa$cov_norm_scaled = NA
		for (i in 1:nrow(mito_cov_byaa)){
			mito_cov_byaa[i,'cov_norm_scaled'] = mito_cov_byaa[i,'cov_norm']/mito_scale_factors[which(mito_cov_byaa[i,'bam'] == mito_scale_factors$bam),3]
		}
		cov_byaa_scaled = rbind(cyto_cov_byaa, mito_cov_byaa)
	}
	else {
		cov_byaa_scaled = cyto_cov_byaa
	}
	
	cov_byaa_norm_scaled = ggplot(subset(cov_byaa_scaled, !grepl("mito", cov_byaa_scaled$aa) & !grepl("plastid", cov_byaa_scaled$aa) & !grepl("nmt", cov_byaa_scaled$aa)), aes(x = bin, y = cov_norm_scaled, fill = aa, group = aa)) + 
		geom_bar(stat = "identity", alpha = 0.8) + facet_wrap(~bam, ncol = 4) + 
		xlab("Gene (%)") + 
		ylab("Scaled normalised coverage") + 
		labs(fill = "Isotype") + 
		guides(fill=guide_legend(ncol=2)) + 
		scale_y_continuous(breaks = seq(0,1,0.25)) + 
		scale_fill_manual(values = getPalette(cyt_colourCount)) + 
		theme_bw()
	ggsave(paste(out_dir, "coverage_byaa_norm_scaled.pdf", sep = ''), cov_byaa_norm_scaled, height = ceiling(facetCount/4) * 4, width = 14)

	if (!is.na(mito_trnas)){
		mitocov_byaa_scaled = subset(cov_byaa_scaled, grepl("mito", cov_byaa_scaled$aa) | grepl("plastid", cov_byaa_scaled$aa) | grepl("nmt", cov_byaa_scaled$aa))
		if (nrow(mitocov_byaa_scaled) != 0) {
			mitocov_byaa_norm_scaled = ggplot(mitocov_byaa_scaled, aes(x = bin, y = cov_norm_scaled, fill = aa, group = aa)) + 
				geom_bar(stat = "identity", alpha = 0.8) + facet_wrap(~bam, ncol = 4) + 
				xlab("Gene (%)") + 
				ylab("Scaled normalised coverage") + 
				labs(fill = "Isotype") +
				guides(fill=guide_legend(ncol=2)) + 
				scale_y_continuous(breaks = seq(0,1,0.25)) + 
				scale_fill_manual(values = getPalette(mit_colourCount)) + 
				theme_bw()
			ggsave(paste(out_dir, "Organelle_coverage_byaa_norm_scaled.pdf", sep = ''), mitocov_byaa_norm_scaled, height = ceiling(facetCount/4) * 4, width = 14)
		}
	}

	cov_byaa$aa_groups = aa_groups[match(cov_byaa$aa, aa_groups$aa),2]

	cov_byaa_agg = aggregate(x = cov_byaa$cov_norm, by = list(bam = cov_byaa$bam, bin = cov_byaa$bin, aa_groups = cov_byaa$aa_groups), FUN = mean)
	cov_byaa_line = ggplot(cov_byaa_agg, aes(x = bin, y = x, color = aa_groups, group = aa_groups)) + 
		geom_line() + 
		facet_wrap(~bam, ncol = 4) +
		xlab("Gene (%)") + 
		ylab("Normalised coverage (coverage/uniquely mapped reads)") + 
		labs(color = "Amino acid group") + 
		theme_bw()
	ggsave(paste(out_dir, "coverage_byaa_line.pdf", sep = ''), cov_byaa_line, height = ceiling(facetCount/4) * 4, width = 14)

	## Coverage plots per cluster multipage
	cov_bygene = read.table(args[1], header = TRUE, sep="\t")
	cov_bygene$Cluster = sub(".*_mito_tRNA-", "mito", cov_bygene$Cluster)
	cov_bygene$Cluster = sub(".*_plastid_tRNA-", "plastid", cov_bygene$Cluster)
	cov_bygene$Cluster = sub(".*_nmt_tRNA-", "nmt", cov_bygene$Cluster)
	cov_bygene$Cluster = sub(".*_tRNA-", "", cov_bygene$Cluster)
	cov_bygene$Cluster = sub(".*_tRX-", "tRX-", cov_bygene$Cluster)
	cov_bygene$bam = gsub(".unpaired_uniq.bam","",cov_bygene$bam)
	cov_bygene$bam = gsub("(.*/).*?","\\2",cov_bygene$bam)
	plot_func = ggplot(cov_bygene, aes(x = bin, y = cov_norm)) + geom_bar(stat = 'identity', fill = "#6BA7BB") + 
  		facet_grid(Cluster~bam, scales='free') + xlab("Gene (%)") + ylab("Normalised coverage") + 
  		theme(axis.title.x=element_blank(), 
  			axis.text.y=element_text(colour="black",size=8),
            axis.text.x=element_blank(), strip.text.y = element_text(angle=0, size = 6),
            strip.text.x = element_text(angle=0, size = 4)) +
  		theme_bw()

  	plots = dlply(cov_bygene , "Cluster", `%+%`, e1 = plot_func)
  	multi = marrangeGrob(plots, nrow = 3, ncol = 1)
  	ggsave(paste(out_dir, "coverage_byUniquetRNA_norm.pdf", sep = ''), multi, height = 12, width = 14)


}
