library(lattice)

#
# Reads depth at position information
#
# depths_file = output from GATK DepthOfCoverage
#
read.target.depths.gatk <- function(depths_file)
{
    x = read.table(depths_file, header=F, skip=1, sep="\t", col.names=c("locus", "depth", "avg", "sample"), colClasses=c("character", "numeric", "NULL", "NULL"))

    if (length(x$locus) == 0)
    {
        print(paste("ERROR: File empty:", depths_file))
        return(data.frame(chr = character(0), pos = numeric(0), depth = numeric(0), index = numeric(0), genotype = character(0)))
    }
    
    y = strsplit(x$locus, ":")
    z = do.call(rbind, y)
    w = data.frame(z, x$depth)
    names(w) = c("chr", "pos", "depth")
    w$index = seq(length(w$chr))

    w$genotype = "Heterozygous"
    x = w
    w$genotype = "Homozygous"
    y = w
    depths = rbind(x, y)
    depths$genotype = as.factor(depths$genotype)

    return(depths)
}

#
# Reads depth at position information and adds recall at depth columns
#
# depths_file = output from GATK DepthOfCoverage or depth_across_targets.pl
# recall_file = recall.tsv included
#
read.target.depths.and.recall <- function(depths_file, recall_file)
{
    depths = read.target.depths.gatk(depths_file)
    recall = read.table(recall_file, header=T)
    
    for (z in levels(recall$genotype))
    {
        for (i in seq(0, 100))
        {
            depths$recall_low[depths$genotype == z & depths$depth == i]  = recall$low[recall$genotype == z & recall$depth == i]
            depths$recall_mean[depths$genotype == z & depths$depth == i] = recall$mean[recall$genotype == z & recall$depth == i]
            depths$recall_high[depths$genotype == z & depths$depth == i] = recall$high[recall$genotype == z & recall$depth == i]
        }
        
        depths$recall_low[depths$genotype == z & depths$depth >= 101]  = recall$low[recall$genotype == z & recall$depth == 101]
        depths$recall_mean[depths$genotype == z & depths$depth >= 101] = recall$mean[recall$genotype == z & recall$depth == 101]
        depths$recall_high[depths$genotype == z & depths$depth >= 101] = recall$high[recall$genotype == z & recall$depth == 101]
    }

    return(depths)
}

#
# Calculates target region boundaries for use in plot.recall.over.target
#
# depths = output from read.target.depths.and.recall
#
calculate.breaks <- function(depths)
{
    breaks = c()
    i = 1
    for (index in seq(2,length(depths$chr)/2))
    {
        # check for gaps in the coordinates
        if (depths$chr[index] != depths$chr[index-1] | depths$pos[index] > depths$pos[index-1] + 1)
        {
            breaks[i] = index
            i = i + 1
        }
    }

    return(breaks)
}

#
# Plots recall over a given set of target regions, with vertical lines delimiting individual targets
#
# depths = output from read.target.depths.and.recall
# breaks = output from calculate.breaks
#
plot.recall.over.target <- function(depths, breaks, title=NULL)
{
    my.panel = function(...) { panel.abline(..., v=c(breaks), col="grey"); panel.xyplot(...) }
    xyplot(recall_mean ~ index | genotype, depths, type="l", lty=1, panel=my.panel, layout=c(1,2), xlab="Indexed position", ylab="Recall", main=title)
}

#
# Reads depth distribution
#
# depths_file = sample_cumulative_coverage_counts output from GATK DepthOfCoverage
#
read.depth.dist.gatk <- function(depths_file)
{
    x = read.table(depths_file, header=T)
    y = data.frame(depth = numeric(102), count = numeric(102))

    for (i in seq(101))
    {
        y$depth[i] = i - 1
        y$count[i] = x[,i] - x[,i+1]
    }
    y$depth[102] = 101
    y$count[102] = x[,101]

    y$genotype = "Heterozygous"
    z = y
    y$genotype = "Homozygous"
    depths = rbind(y, z)
    depths$genotype = as.factor(depths$genotype)

    return(depths)
}

#
# Summarizes expected total recall over a set of target regions
#
# depths_file = output from depth_distribution.pl or sample_cumulative_coverage_counts from GATK DepthOfCoverage
# recall_file = recall.tsv included
# target_size = total size of targets in nt
#
summarize.recall <- function(depths_file, recall_file, target_size)
{
    depths = read.depth.dist.gatk(depths_file)
    recall = read.table(recall_file, header=T)
    
    for (z in levels(recall$genotype))
    {
        for (i in seq(0, 100))
        {
            depths$recall_low[depths$genotype == z & depths$depth == i]  = recall$low[recall$genotype == z & recall$depth == i]
            depths$recall_mean[depths$genotype == z & depths$depth == i] = recall$mean[recall$genotype == z & recall$depth == i]
            depths$recall_high[depths$genotype == z & depths$depth == i] = recall$high[recall$genotype == z & recall$depth == i]
        }
        
        depths$recall_low[depths$genotype == z & depths$depth >= 101]  = recall$low[recall$genotype == z & recall$depth == 101]
        depths$recall_mean[depths$genotype == z & depths$depth >= 101] = recall$mean[recall$genotype == z & recall$depth == 101]
        depths$recall_high[depths$genotype == z & depths$depth >= 101] = recall$high[recall$genotype == z & recall$depth == 101]
    }

    depths$prop = depths$count / target_size

    low  = aggregate(depths$prop * depths$recall_low, by=list(genotype=depths$genotype), sum)
    mid  = aggregate(depths$prop * depths$recall_mean, by=list(genotype=depths$genotype), sum)
    high = aggregate(depths$prop * depths$recall_high, by=list(genotype=depths$genotype), sum)

    res = low
    res$recall_low = res$x
    res$x = NULL
    res$recall_mean = mid$x
    res$recall_high = high$x
    
    return(res)
}