#!/usr/bin/python
"""
Create a matrix of many pairwise intersections; see \
:mod:`pybedtools.contrib.IntersectionMatrix` for more flexibility
"""

import collections
import time
import sys
import os.path as op
import argparse
import pybedtools
from pybedtools import BedTool, example_filename

usage = """

    Send in a list of `N` bed files, and this script will create an N by
    N matrix of their intersections, or optionally, co-localization scores.

    Run the example with::

        %s --test > matrix.txt

    You can then plot a quick heatmap in R with::

        > m = as.matrix(read.table("matrix.txt"))
        > heatmap(m)

""" % sys.argv[0]


def get_name(fname):
    return op.splitext(op.basename(fname))[0]


def actual_intersection(a, b):
    return len(a.intersect(b, u=True))


def frac_of_a(a, b):
    len_a = float(len(a))
    return len(a.intersect(b, u=True)) / len_a


def enrichment_score(a, b, genome_fn, iterations=None, processes=None):
    results = a\
            .randomstats(b, new=True, genome_fn=genome_fn, iterations=iterations, processes=processes)
    return (results['actual'] + 1) / (results['median randomized'] + 1)


def create_matrix(beds, func, verbose=False, **kwargs):
    nfiles = len(beds)
    total = nfiles ** 2
    i = 0
    matrix = collections.defaultdict(dict)
    for fa in beds:
        a = BedTool(fa)
        for fb in beds:
            i += 1
            b = BedTool(fb)

            if verbose:
                sys.stderr.write(
                        '%(i)s of %(total)s: %(fa)s + %(fb)s\n' % locals())
                sys.stderr.flush()

            matrix[get_name(fa)][get_name(fb)] = func(a, b, **kwargs)

    return matrix


def main():
    """
    Creates a pairwise matrix containing overlapping feature counts for many
    BED files
    """
    ap = argparse.ArgumentParser(usage=usage)
    ap.add_argument('beds', nargs="*", help='BED/GTF/GFF/VCF filenames, e.g., '
                    'in a directory of bed files, you can use *.bed')
    ap.add_argument('--frac', action='store_true',
                    help='Instead of counts, report fraction overlapped')
    ap.add_argument('--enrichment', action='store_true',
                    help='Run randomizations (default 1000, specify otherwise '
                    'with --iterations) on each pairwise comparison and '
                    'compute the enrichment score as '
                    '(actual intersection count + 1) / (median randomized + 1)'
                    )
    ap.add_argument('--genome', help='Required argument if --enrichment is '
                    'used. Needs to be a string assembly name like "dm3" or '
                    '"hg19"')
    ap.add_argument('--iterations', default=1000, type=int,
                    help='Number of randomizations to perform for enrichement '
                    'scores')
    ap.add_argument('--processes', default=None, type=int,
                    help='Number of CPUs to use for randomization')
    ap.add_argument('--test', action='store_true', help='Ignore any input BED '
                    'files and use test BED files')
    ap.add_argument('-v', '--verbose', action='store_true',
                    help='Be verbose: print which files are '
                    'currently being intersected and timing info at the end.')
    args = ap.parse_args()

    if not args.beds and not args.test:
        ap.print_help()
        sys.exit(1)

    if args.test:
        # insulator binding sites from ChIP-chip -- 4 proteins, 2 cell types
        # Genes Dev. 2009 23(11):1338-1350
        args.beds = [example_filename(i) for i in  [
                'Cp190_Kc_Bushey_2009.bed',
                'Cp190_Mbn2_Bushey_2009.bed',
                'CTCF_Kc_Bushey_2009.bed',
                'CTCF_Mbn2_Bushey_2009.bed',
                'SuHw_Kc_Bushey_2009.bed',
                'SuHw_Mbn2_Bushey_2009.bed',
                'BEAF_Mbn2_Bushey_2009.bed',
                'BEAF_Kc_Bushey_2009.bed'
                ]]

    if args.enrichment:
        FUNC = enrichment_score
        genome_fn = pybedtools.chromsizes_to_file(pybedtools.chromsizes(args.genome))
        kwargs = dict(genome_fn=genome_fn, iterations=args.iterations,
                processes=args.processes)

    elif args.frac:
        FUNC = frac_of_a
        kwargs = {}
    else:
        FUNC = actual_intersection
        kwargs = {}

    t0 = time.time()
    matrix = create_matrix(beds=args.beds, func=FUNC, verbose=args.verbose, **kwargs)
    t1 = time.time()

    nfiles = len(args.beds)

    if args.verbose:
        sys.stderr.write('Time to construct %s x %s matrix: %.1fs' \
                % (nfiles, nfiles, (t1 - t0)) + '\n')
    keys = sorted(matrix.keys())

    sys.stdout.write("\t" + "\t".join(keys) + '\n')
    for k in keys:
        sys.stdout.write(k)
        for j in keys:
            sys.stdout.write('\t' + str(matrix[k][j]))
        sys.stdout.write('\n')

if __name__ == "__main__":
    main()