#! /usr/bin/env Rscript
suppressPackageStartupMessages(library("argparse"))
parser <- ArgumentParser(description='Call domains with significant coverage')
parser$add_argument('-p', '--pvalue', type='double', default=0.05,
    help='adjusted p-value threshold for defining significant domains', metavar='NUMBER')
parser$add_argument('-i', '--input-file', type='character', default='-',
    help='input BED file with covearge in the 5th column')
parser$add_argument('-o', '--output-file', type='character', default='-',
    help='output BED file with adjusted p-value in the last column')
parser$add_argument('-b', '--background', type='double', default=0.99,
    help='set coverage quantile below this quantile as background regions')
args <- parser$parse_args()

library(countreg)

ztnb.test <- function(x, mu, theta, max_x=2000) {
    # build a p-value table
    pvalue_table <- dztnbinom(1:max_x, mu=mu, theta=theta)
    pvalue_table <- rev(cumsum(rev(pvalue_table)))
    # calculate pvalues
    pvalues <- rep(pvalue_table[max_x], length(x))
    pvalues[x <= max_x] <- pvalue_table[x[x <= max_x]]
    qvalues <- p.adjust(pvalues, method='BH')
    return(qvalues)
}
message('read coverage file:', args$input_file)
if(args$input_file == '-'){
    bed <- read.table(file('stdin'), sep='\t', header=FALSE)
} else {
    bed <- read.table(args$input_file, sep='\t', header=FALSE)
}
message('number of bins:', nrow(bed))
# remove bins with 1 read
#bed <- bed[bed[, 5] > 1,]
x <- as.integer(bed[, 5])
bg = x[x < quantile(x, args$background)]
if(length(bg) > 0){
    message('fit ZTNB distribution')
    fit <- zerotrunc(bg ~ 1, dist='negbin')
    qvalues <- ztnb.test(x, mu=exp(fit$coefficients), theta=fit$theta)
    bg <- x[qvalues >= args$pvalue]
    message('fit ZTNB distribution')
    fit <- zerotrunc(bg ~ 1, dist='negbin')
    qvalues <- ztnb.test(x, mu=exp(fit$coefficients), theta=fit$theta)
    bed_sig <- bed[qvalues < args$pvalue,]
    bed_sig <- cbind(bed_sig, qvalues[qvalues < args$pvalue])
} else {
    # empty file
    bed_sig <- data.frame()
}
message('number of significant bins:', nrow(bed_sig))
message('write significant domains')
if(args$output_file == '-'){
    write.table(bed_sig, row.names=FALSE, col.names=FALSE, sep='\t', quote=FALSE)
} else {
    write.table(bed_sig, args$output_file, row.names=FALSE, col.names=FALSE, sep='\t', quote=FALSE)
}