#!/usr/bin/env python
"""
intron_histogram.py

Outputs list of number of simulated samples from the 112 in which each intron
occurs.

We executed python intron_histogram.py -t
    /scratch0/langmead-fs1/geuvadis_sims_for_paper_v2
    >intron_histogram; the argument of -t
    should be where the Flux BEDs for the 112 simulated RNA-seq samples
    generated by generate_bioreps.py are.
"""

import os
from collections import defaultdict
import sys
import glob
import time
import math
import multiprocessing
import signal
import re

def init_worker():
    """ Prevents KeyboardInterrupt from reaching a pool's workers.

        Exiting gracefully after KeyboardInterrupt or SystemExit is a
        challenge. The solution implemented here is by John Reese and is from
        http://noswap.com/blog/python-multiprocessing-keyboardinterrupt .

        No return value.
    """
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def introns_from_bed(index_and_bed):
    """ Converts BED to dictionary that maps RNAMES to sets of introns.
        
        index_and_bed is composed of:
        index: index of bed file; used for tracking which bed is which
            when using multiple cores
        bed: input BED file characterizing splice junctions

        Return value: (index, a dictionary). Each key is an RNAME, typically a
            chromosome, and its corresponding value is a set of tuples, each
            denoting an intron on RNAME. Each tuple is of the form
            (start position, end position).
    """
    index, bed = index_and_bed
    introns = set()
    with open(bed) as bed_stream:
        for line in bed_stream:
            tokens = line.rstrip().split('\t')
            if len(tokens) < 12:
                continue
            chrom = tokens[0]
            chrom_start = int(tokens[1])
            chrom_end = int(tokens[2])
            block_sizes = tokens[10].split(',')
            block_starts = tokens[11].split(',')
            # Handle trailing commas
            try:
                int(block_sizes[-1])
            except ValueError:
                block_sizes = block_sizes[:-1]
            try:
                int(block_starts[-1])
            except ValueError:
                block_starts = block_starts[:-1]
            block_count = len(block_sizes)
            if block_count < 2:
                # No introns
                continue
            assert block_count == len(block_starts)
            junctions = []
            # First block characterizes junction on left side of intron
            junctions.append(chrom_start + int(block_starts[0]) 
                                    + int(block_sizes[0]))
            for i in xrange(1, block_count - 1):
                # Any intervening blocks characterize two junctions
                intron_start = chrom_start + int(block_starts[i])
                junctions.append(intron_start)
                junctions.append(intron_start + int(block_sizes[i]))
            # Final block characterizes junction on right side of intron
            junctions.append(chrom_start + int(block_starts[-1]))
            for i in xrange(len(junctions)/2):
                introns.add((chrom, junctions[2*i]+1, junctions[2*i+1]+1))
    return (index, introns)

if __name__ == '__main__':
    import argparse
    import subprocess
    # Print file's docstring if -h is invoked
    parser = argparse.ArgumentParser(description=__doc__, 
                formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-t', '--true-introns-bed-dir', type=str,
            required=True,
            help=('Path to directory containing Flux Simulator BEDs; the only '
                  'BED files in this directory should be the 112 from Flux')
        )
    args = parser.parse_args()
    # Ingest true introns
    true_introns = defaultdict(list)
    bed_paths = glob.glob(os.path.join(args.true_introns_bed_dir, '*_sim.bed'))
    print >>sys.stderr, 'Loading true introns...'
    # Use multiple cores
    pool = multiprocessing.Pool(multiprocessing.cpu_count() - 1,
                                    init_worker, maxtasksperchild=5)
    returned_introns = []
    pool.map_async(
                    introns_from_bed,
                    list(enumerate(bed_paths)),
                    callback=returned_introns.extend
                )
    pool.close()
    pool.join()
    for index, true_intron_set in returned_introns:
        for intron in true_intron_set:
            true_introns[intron].append(index)
    for intron in true_introns:
        print ';'.join(str(el) for el in intron) + (
                                    '\t%d' % len(true_introns[intron])
                                )
