################################################################################
#
#   MRC FGU Computational Genomics Group
#
#   $Id: cgat_script_template.py 2871 2010-03-03 10:20:44Z andreas $
#
#   Copyright (C) 2009 Andreas Heger
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; either version 2
#   of the License, or (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################
'''
rnaseq_bam_vs_exons.py - compare bam file against exons
=======================================================

:Author: Andreas Heger
:Release: $Id$
:Date: |today|
:Tags: Python

Purpose
-------

Compare RNASeq reads in a BAM file against reference exons.

The script outputs counts for exon over-/underrun. 

For unspliced reads, any bases extending beyond exon boundaries are counted. 

For spliced reads, both parts of the reads are examined for their overlap.
As a consequence, counts are doubled for spliced reads.

This script is for validation purposes:

   * Exon overrun should be minimal - reads should not 
      extend beyond known exons.

   * Spliced reads should link known exons.

The script requires a list of non-overlapping exons.



Usage
-----

Example::

   python cgat_script_template.py --help

Type::

   python cgat_script_template.py --help

for command line help.

Documentation
-------------

For read counts to be correct the NH flag to be set correctly.

Code
----

'''

import os
import sys
import re
import optparse
import collections
import CGAT.Experiment as E
import CGAT.IOTools as IOTools
import pysam
import CGAT.GTF as GTF

def main( argv = None ):
    """script main.

    parses command line options in sys.argv, unless *argv* is given.
    """

    if not argv: argv = sys.argv

    # setup command line parser
    parser = E.OptionParser( version = "%prog version: $Id: cgat_script_template.py 2871 2010-03-03 10:20:44Z andreas $", 
                                    usage = globals()["__doc__"] )

    parser.add_option( "-e", "--filename-exons", dest="filename_exons", type="string",
                       help = "gtf formatted file with exons locations. [%default]" )

    parser.set_defaults(
        filename_exons = None,
        read_length = 200,
        )

    ## add common options (-h/--help, ...) and parse command line 
    (options, args) = E.Start( parser, argv = argv, add_output_options = True )

    exons = GTF.readAndIndex( GTF.iterator( IOTools.openFile( options.filename_exons ) ) )

    pysam_in = pysam.Samfile( "-", "rb" )
    
    nspliced = 0
    nspliced_ignored = 0
    nspliced_nooverlap = 0
    nspliced_halfoverlap = 0
    nspliced_bothoverlap = 0
    nspliced_overrun = [0] * 2 * (options.read_length + 10)
    nspliced_exact = 0
    nspliced_inexact = 0
    nunspliced = 0
    nunspliced_overlap = 0
    nunspliced_ignored = 0
    nunspliced_nooverlap = 0
    nunspliced_overrun = [0] * (options.read_length + 10)

    overrun_offset = options.read_length + 10

    ninput = 0
    nunmapped = 0

    c = E.Counter()

    def _splice_overrun( start, end, overlap ):
        '''return splicesite over/underrun.

        positive values: overrun
        negative values: underrun
        0: no over/underrun
        '''

        exon_start = min( [x[0] for x in overlap] )
        exon_end = max( [x[1] for x in overlap] )

        if start <= exon_start and end > exon_start:
            # overrun at start or match
            r = exon_start - start
        elif start < exon_end and end >= exon_end:
            # overrun at end or match
            r = end - exon_end
        else:
            # underrun - distance to closest exon boundary
            r = -min( start - exon_start, exon_end - end )

        return r

    for read in pysam_in:
        ninput += 1
        if read.is_unmapped: 
            nunmapped += 1
            continue
        
        # check for BAM_CREF_SKIP code in cigar string
        cigar = read.cigar
        is_spliced = 3 in [ x[0] for x in cigar ]

        contig = pysam_in.getrname( read.tid ) 
        start = read.pos
        end = read.aend
        if is_spliced:
            # count both ends
            nspliced += 1

            if len(cigar) != 3:
                nspliced_ignored += 1
                continue

            start5, end5 = start, start + cigar[0][1]
            start3, end3 = end - cigar[2][1], end
            try:
                overlap3 = list( exons.get( contig, start3, end3 ) )
                overlap5 = list( exons.get( contig, start5, end5 ) )
            except KeyError:
                overlap3 = overlap5 = []

            ovl3 = len(overlap3) 
            ovl5 = len(overlap5) 
            o3 = o5 = None
            if not ovl3 and not ovl5:
                nspliced_nooverlap += 1
            elif ovl3 and not ovl5:
                nspliced_halfoverlap += 1
                o3 = _splice_overrun( start3, end3, overlap3 )
            elif ovl5 and not ovl3:
                nspliced_halfoverlap += 1
                o5 = _splice_overrun( start5, end5, overlap5 )
            else:
                # both overlap
                nspliced_bothoverlap += 1
                o3 = _splice_overrun( start3, end3, overlap3 )
                o5 = _splice_overrun( start5, end5, overlap5 )

            if o3 != None:
                if o3 == 0: nspliced_exact += 1
                else: nspliced_inexact += 1
                nspliced_overrun[ max(0, overrun_offset + o3)] += 1
            if o5 != None:
                if o5 == 0: nspliced_exact += 1
                else: nspliced_inexact += 1
                nspliced_overrun[ max(0, overrun_offset + o5)] += 1
        else:
            nunspliced += 1
            try:
                overlap = list( exons.get( contig, start, end ) )
            except KeyError:
                overlap = []

            if len(overlap) == 0:
                nunspliced_nooverlap += 1
            elif len(overlap) >= 1:
                nunspliced_overlap += 1
                # multiple overlap - merge exons (usually: small introns)
                exon_start = min( [x[0] for x in overlap] )
                exon_end = max( [x[1] for x in overlap] )
                ostart = max( 0, exon_start - start )
                oend = max( 0, end - exon_end )
                o = min(end, exon_end) - max( start, exon_start)
                overrun = ostart + oend
                nunspliced_overrun[overrun] += 1

    # output histograms
    outfile = E.openOutputFile( "overrun" )
    outfile.write( "bases\tunspliced_overrun_counts\tspliced_overrun_counts\tspliced_underrun_counts\n" )
    _nspliced_overrun = nspliced_overrun[overrun_offset:]
    _nspliced_underrun = nspliced_overrun[:overrun_offset+1]
    _nspliced_underrun.reverse()
    for x, v in enumerate( zip( nunspliced_overrun, _nspliced_overrun, _nspliced_underrun) ):
        outfile.write( "%i\t%s\n" % (x, "\t".join(map(str,v))))
    outfile.close()

    # output summary
    # convert to counter
    c.input = ninput
    c.unmapped = nunmapped
    c.mapped = ninput - nunmapped

    c.unspliced = nunspliced
    c.unspliced_nooverlap = nunspliced_nooverlap
    c.unspliced_nooverrun = nunspliced_overrun[0]
    c.unspliced_overlap = nunspliced_overlap
    c.unspliced_overrun = sum( nunspliced_overrun[1:])

    c.spliced = nspliced
    c.spliced_nooverlap = nspliced_nooverlap
    c.spliced_halfoverlap = nspliced_halfoverlap
    c.spliced_bothoverlap = nspliced_bothoverlap
    c.spliced_exact = nspliced_exact
    c.spliced_inexact = nspliced_inexact
    c.spliced_ignored = nspliced_ignored
    c.spliced_underrun = sum( _nspliced_underrun[1:])
    c.spliced_overrun = sum( _nspliced_overrun[1:])
    
    outfile = options.stdout
    outfile.write( "category\tcounts\n" )
    for k,v in c.iteritems():
        outfile.write( "%s\t%i\n" % (k,v))
    
    ## write footer and output benchmark information.
    E.Stop()

if __name__ == "__main__":
    sys.exit( main( sys.argv) )

    
