#!/usr/bin/Rscript --vanilla
library(getopt)
spec <- matrix(c("help", "h", 0, "logical", "view this help", "segfile1", "s", 1, "character", "seg file 1", "segfile2", "t", 1, "character", "seg file 2", "output", "o", 1, "character", "output file", "fdr", "f", 2, "character", paste("fdr method [", paste(p.adjust.methods, collapse="|"), "]", sep=""),
        "reference", "r", 2, "character", "reference to use [b37|hg19|GRCh37|mm9|NCBIM37|mm10|GRCm38|dm3|BDGP5]", "annot", "a", 2, "character", "annotation to add [both|gene|cpg]", "processes", "p", 2, "integer", "number of cluster instances to open [1]"), ncol=5, byrow=T)
opt <- getopt(spec)

# set default options
if(is.null(opt$fdr)) opt$fdr  <- "BH"
if(is.null(opt$reference)) opt$reference  <- "hg19"
if(is.null(opt$annot)) opt$annot  <- "both"
if(is.null(opt$processes)) opt$processes  <- 1

# check if any invalid options
if(! opt$annot %in% c("both", "gene", "cpg") || !opt$fdr %in% p.adjust.methods || !opt$reference %in% c("b37", "hg19", "GRCh37", "mm9", "NCBIM37", "mm10", "GRCm38", "dm3", "BDGP5")) {
        opt$help <- 1
}

# print help file if any incorrect
if(!is.null(opt$help) || is.null(opt$segfile1) || is.null(opt$output) || is.null(opt$output) || is.na(opt$processes)) {
        self = commandArgs()[1];
        cat(getopt(spec, usage=T, command="differential_methylation.R"))
        q(status=1)
}

library(snow)
cl <- makeCluster(opt$p, type = "MPI")

segfile1 <- read.table(opt$segfile1, sep="\t", as.is=T, head=T, quote = "")
segfile2 <- read.table(opt$segfile2, sep="\t", as.is=T, head=T, quote = "")
	
rownames(segfile1) <- paste(segfile1[,2], ":", segfile1[,3], "-", segfile1[,4], sep="")
rownames(segfile2) <- paste(segfile2[,2], ":", segfile2[,3], "-", segfile2[,4], sep="")
	
common_reg <- intersect(rownames(segfile1), rownames(segfile2))	
prop_test <- function(x, n, fdr) {
	library(abind)
	ESTIMATE <- x/n
	DELTA <- ESTIMATE[,1L] - ESTIMATE[,2L]
	YATES <- pmin(abs(DELTA)/rowSums(1/n), 0.5)
	p <- rowSums(x)/rowSums(n)
	df <- 1
	x <- abind(x, n - x, along=3)
	E <- abind(n*p, n*(1-p), along=3)
	STATISTIC <- rowSums((abs(x - E) - YATES)^2/E)
	STATISTIC[is.na(STATISTIC)] <- 0
	PVAL <- pchisq(STATISTIC, df, lower.tail = FALSE)
	return(data.frame(X_Squared = round(STATISTIC, 3), P.value= round(PVAL, 6), P.adjusted = round(p.adjust(PVAL, method=fdr), 6)))
}
results <- prop_test(cbind(segfile1[common_reg, "Methylated"], segfile2[common_reg, "Methylated"]), cbind(segfile1[common_reg, "Total"], segfile2[common_reg, "Total"]), opt$fdr)
	
sample1 <- segfile1[1, 1]
sample2 <- segfile2[1, 1]
tab_out <- data.frame(ID = paste(sample1, "vs", sample2, sep="."), segfile1[common_reg, c(2:4)], segfile1[common_reg, c("Methylated", "Total", "FractionMethylated")],
		segfile2[common_reg, c("Methylated", "Total", "FractionMethylated")], results, DiffProp = round(segfile1[common_reg, "FractionMethylated"] - segfile2[common_reg, "FractionMethylated"], 6), stringsAsFactors=F)
colnames(tab_out)[c(5:10)] <- c(paste(sample1, "Methylated", sep="."), paste(sample1, "Total", sep="."), paste(sample1, "Proportion", sep="."), paste(sample2, "Methylated", sep="."),
		paste(sample2, "Total", sep="."), paste(sample2, "Proportion", sep="."))

rm(segfile1)
rm(segfile2)
gc()

add_annot <- function(tab, annotation, genome) {
        # find closest feature
	if(genome == "hg19" || genome == "GRCh37" || genome == "b37") {
		genome <- "hg19"
		dataset <- "hsapiens_gene_ensembl"
		biomart <- "ensembl"
		host <- "www.biomart.org"
		attributes <- c("ensembl_gene_id", "hgnc_symbol", "refseq_mrna", "start_position", "end_position")
	} else if(genome == "mm9" || genome == "NCBIM37") {
		genome <- "mm9"
		dataset <- "mmusculus_gene_ensembl"
		biomart <- "ENSEMBL_MART_ENSEMBL"
		host <- "may2012.archive.ensembl.org"
		attributes <- c("ensembl_gene_id", "mgi_symbol", "refseq_mrna", "start_position", "end_position")
	} else if(genome == "mm10" || genome == "GRCm38") {
		genome <- "mm10"
		dataset <- "mmusculus_gene_ensembl"
		biomart <- "ensembl"
		host <- "www.biomart.org"
		attributes <- c("ensembl_gene_id", "mgi_symbol", "refseq_mrna", "start_position", "end_position")
	} else if(genome == "dm3" || genome == "BDGP5") {
		genome <- "dm3"
		dataset <- "dmelanogaster_gene_ensembl"
		biomart <- "ensembl"
		host <- "www.biomart.org"
		attributes <- c("ensembl_gene_id", "flybasecgid_gene", "refseq_mrna", "start_position", "end_position")
	}

        e_to_U = c("GL000191.1" = "1_gl000191_random", "GL000192.1" = "1_gl000192_random", "GL000193.1" = "4_gl000193_random", "GL000194.1" = "4_gl000194_random",
		"GL000195.1" = "7_gl000195_random", "GL000196.1" = "8_gl000196_random", "GL000197.1" = "8_gl000197_random", "GL000198.1" = "9_gl000198_random",
		"GL000199.1" = "9_gl000199_random", "GL000200.1" = "9_gl000200_random", "GL000201.1" = "9_gl000201_random", "GL000202.1" = "11_gl000202_random",
		"GL000203.1" = "17_gl000203_random", "GL000204.1" = "17_gl000204_random", "GL000205.1" = "17_gl000205_random", "GL000206.1" = "17_gl000206_random",
		"GL000207.1" = "18_gl000207_random", "GL000208.1" = "19_gl000208_random", "GL000209.1" = "19_gl000209_random", "GL000210.1" = "21_gl000210_random",
		"GL000211.1" = "Un_gl000211", "GL000212.1" = "Un_gl000212", "GL000213.1" = "Un_gl000213", "GL000214.1" = "Un_gl000214",
		"GL000215.1" = "Un_gl000215", "GL000216.1" = "Un_gl000216", "GL000217.1" = "Un_gl000217", "GL000218.1" = "Un_gl000218",
		"GL000219.1" = "Un_gl000219", "GL000220.1" = "Un_gl000220", "GL000221.1" = "Un_gl000221", "GL000222.1" = "Un_gl000222",
		"GL000223.1" = "Un_gl000223", "GL000224.1" = "Un_gl000224", "GL000225.1" = "Un_gl000225", "GL000226.1" = "Un_gl000226",
		"GL000227.1" = "Un_gl000227", "GL000228.1" = "Un_gl000228", "GL000229.1" = "Un_gl000229", "GL000230.1" = "Un_gl000230",
		"GL000231.1" = "Un_gl000231", "GL000232.1" = "Un_gl000232", "GL000233.1" = "Un_gl000233", "GL000234.1" = "Un_gl000234",
		"GL000235.1" = "Un_gl000235", "GL000236.1" = "Un_gl000236", "GL000237.1" = "Un_gl000237", "GL000238.1" = "Un_gl000238",
		"GL000239.1" = "Un_gl000239", "GL000240.1" = "Un_gl000240", "GL000241.1" = "Un_gl000241", "GL000242.1" = "Un_gl000242",
		"GL000243.1" = "Un_gl000243", "GL000244.1" = "Un_gl000244", "GL000245.1" = "Un_gl000245", "GL000246.1" = "Un_gl000246",
		"GL000247.1" = "Un_gl000247", "GL000248.1" = "Un_gl000248", "GL000249.1" = "Un_gl000249")

	# function to conver ucsc chroms to ensembl
	ensembl <- function(chr, genome) {
		# ensembl does not us chr and M is MT
		chr <- sub("^chr", "", chr)
		if(genome == "dm3" || genome == "BDGP5") {
			chr <- sub("^M$", "dmel_mitochondrion_genome", chr)
		} else {
			chr <- sub("^M$", "MT", chr)
		}
		chr[which(chr %in% e_to_U)] <- names(e_to_U)[match(chr[which(chr %in% e_to_U)], e_to_U)]
		return(chr)
	}

	# function to conver ensembl chroms to ucsc
	ucsc <- function(chr) {
		chr <- sub("^chr", "", chr)
		tmp <- which(chr %in% names(c(e_to_U, dmel_mitochondrion_genome = "M", MT = "M")))
		chr[tmp] <- c(e_to_U, dmel_mitochondrion_genome = "M", MT = "M")[chr[tmp]]
		paste("chr", chr, sep="")
	}

	# function to get the gene annotation by chromosome
        get_genes <- function(chr) {
		library(biomaRt)
		mart <- useMart(biomart=biomart, dataset=dataset, host=host)
		tab <- getBM(attributes = attributes, filters = "chromosome_name", values = chr, mart = mart)
		if(any(is.na(tab))) for(i in 1:ncol(tab)) tab[is.na(tab[,i]), i] <- ""
		mult_ids <- names(which(table(tab$ensembl_gene_id) > 1))
		rem <- NULL
		for(i in mult_ids) {
			index <- which(tab$ensembl_gene_id == i)
			refseq <- which(tab$ensembl_gene_id == i & tab$refseq_mrna != "")
			if(length(refseq) > 0) {
				rem <- c(rem, setdiff(index, refseq[1]))
			} else {
				rem <- c(rem, index[-1])
			}
		}
		if(length(rem) > 0) tab <- tab[-rem,]
		colnames(tab)[4:5] <- c("feature_start", "feature_end")
		return(tab)
	}
	
	get_cpg <- function(chr) {
		library(rtracklayer)
		options(stringsAsFactors=F)
		session <- browserSession()
		genome(session) <- genome
		query <- ucscTableQuery(session, "CpG Islands", GRangesForUCSCGenome(genome, chr))
		cpg <- getTable(query)
		cpg <- cpg[, c("name", "perCpg", "perGc", "chromStart", "chromEnd")]
		colnames(cpg) <- c("cpg", "cpg_perCpg", "cpg_perGc", "cpg_start", "cpg_end")
		return(cpg)
	}
	chr <- tab[1, 2]

	# get the gene info for this chrom
	if(annotation == "gene") {
		tab <- cbind(tab, ensembl_gene_id = "", id = "", refseq_mrna = "", feature_start = "", feature_end = "", Distance_To_Feature = "", stringsAsFactors=F)
		annotData <- get_genes(ensembl(chr, genome))
		colnames(tab)[16] <- attributes[2]
	} else {
		tab <- cbind(tab, cpg = "", cpg_perCpg = "", cpg_perGc = "", cpg_start = "", cpg_end = "", Distance_To_cpg = "", stringsAsFactors=F)
		annotData <- get_cpg(ucsc(chr))
	}
	if(nrow(annotData) == 0) return(tab)
	
	starts <- t(matrix(cbind(-1, as.numeric(annotData[, 4])), ncol=2, byrow=F))
	ends <- t(matrix(cbind(-1, as.numeric(annotData[, 5])), ncol=2, byrow=F))

	# calculate the distances from the features to the regions
	dist_start_start <- matrix(cbind(tab$loc.start, 1), ncol=2, byrow=F) %*% starts
	dist_start_end <- matrix(cbind(tab$loc.start, 1), ncol=2, byrow=F) %*% ends
	dist_end_start <- matrix(cbind(tab$loc.end, 1), ncol=2, byrow=F) %*% starts
	dist_end_end <- matrix(cbind(tab$loc.end, 1), ncol=2, byrow=F) %*% ends
		
	# determine which regions overlap at least 1 feature
	sum_signs <- abs(sign(dist_start_start) + sign(dist_start_end) + sign(dist_end_start) + sign(dist_end_end))
	regions <- which(sum_signs != 4, arr.ind=TRUE)
	if(length(regions) > 0) {
		overlap <- sort(unique(regions[,1]))
		non_overlap <- c(1:nrow(tab))[-overlap]
	} else {
		overlap <- NULL
		non_overlap <- c(1:nrow(tab))
	}
		
	# reduce to regions with no overlaping feqature
	if(length(overlap) > 0) {
		dist_start_start <- matrix(dist_start_start[non_overlap,], ncol=ncol(dist_start_start))
		dist_start_end <- matrix(dist_start_end[non_overlap,], ncol=ncol(dist_start_end))
		dist_end_start <- matrix(dist_end_start[non_overlap,], ncol=ncol(dist_end_start))
		dist_end_end <- matrix(dist_end_end[non_overlap,], ncol=ncol(dist_end_end))
	}
		
	rm(sum_signs)
	gc()
		
	# extract the annot for the regions with overlaping features
	if(length(overlap) > 0) {
		annot <- sapply(overlap, function(x, y) {
					x <- regions[which(regions[,1] == x), 2]
					sub("^//$", "", c(paste(annotData[x,1], collapse="//"), paste(annotData[x,2], collapse="//"), paste(annotData[x,3], collapse="//"), paste(annotData[x,4], collapse="//"), paste(annotData[x,5], collapse="//")))
				}, regions)
		annot <- as.data.frame(t(annot), stringsAsFactors=F)
		annot <- cbind(annot, "Overlap", stringsAsFactors=F)
		colnames(annot) <- c(colnames(annotData), tail(colnames(tab),1))
	}
	rm(regions)
	gc()
		
	# for non the non-overlaps the distance of the closest features to each region
	if(length(non_overlap) > 0) {
		clst_pts <- matrix(0, ncol=4, nrow=length(non_overlap))
		clst_pts[,1] <- max.col(-abs(dist_start_start), "last")
		clst_pts[,2] <- max.col(-abs(dist_start_end), "last")
		clst_pts[,3] <- max.col(-abs(dist_end_start), "last")
		clst_pts[,4] <- max.col(-abs(dist_end_end), "last")
		dist <- matrix(0, ncol=4, nrow=length(non_overlap))
		if(length(clst_pts[,1]) == 1) {
			dist[,1] <- dist_start_start[1, clst_pts[,1]]
			dist[,2] <- dist_start_end[1, clst_pts[,2]]
			dist[,3] <- dist_end_start[1, clst_pts[,3]]
			dist[,4] <- dist_end_end[1, clst_pts[,4]]

			# extract the annot for the regions with non-overlaping features
			clst_all <- max.col(-abs(dist))
			dist_to_feat <- dist[1, clst_all]	
			clst <- clst_pts[1, clst_all]
		} else {
			dist[,1] <- dist_start_start[cbind(seq(clst_pts[,1]), clst_pts[,1])]
			dist[,2] <- dist_start_end[cbind(seq(clst_pts[,2]), clst_pts[,2])]
			dist[,3] <- dist_end_start[cbind(seq(clst_pts[,3]), clst_pts[,3])]
			dist[,4] <- dist_end_end[cbind(seq(clst_pts[,4]), clst_pts[,4])]

			# extract the annot for the regions with non-overlaping features
			clst_all <- max.col(-abs(dist))
			dist_to_feat <- dist[cbind(seq(clst_all), clst_all)]	
			clst <- clst_pts[cbind(seq(clst_all), clst_all)]
		}
		annot_non_overlap <- cbind(annotData[clst, ], dist_to_feat)
		colnames(annot_non_overlap) <- c(colnames(annotData), tail(colnames(tab),1))
	}
		
	rm(dist_start_start)
	rm(dist_start_end)
	rm(dist_end_start)
	rm(dist_end_end)
	gc()
		
	if(length(overlap) > 0) tab[overlap, colnames(annot)] <- annot
	if(length(non_overlap) > 0) tab[non_overlap, colnames(annot_non_overlap)] <- annot_non_overlap
	return(tab)
}

tab_list <- list()
for(i in unique(tab_out[,2])) tab_list[[i]] <- tab_out[which(tab_out$chrom == i),]
rm(tab_out)
gc()
if(opt$annot == "both" || opt$annot == "gene") {
	annotation <- "gene"
	tab_list <- clusterApplyLB(cl, tab_list, add_annot, annotation, opt$reference)
}

if(opt$annot == "both" || opt$annot == "cpg") {
	annotation <- "cpg"
	tab_list <- clusterApplyLB(cl, tab_list, add_annot, annotation, opt$reference)
}

tab_out <- NULL
for(i in 1:length(tab_list)) tab_out <- rbind(tab_out, tab_list[[i]])

cat("'", file=opt$output)
write.table(tab_out[,c(1:13, 15:26, 14)], opt$output, row.names=F, col.names=T, quote=F, sep="\t", append=T)
stopCluster(cl)
q(status=0)