# Source code for pybedtools.contrib.intersection_matrix

    import os
    import sys
    import sqlite3
    import pybedtools
    import time
    import collections

    def now():
        return time.time()

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

    class IntersectionMatrix(object):
        """
        Class to handle many pairwise comparisons of interval files
        """
        def __init__(self, beds, genome, iterations, dbfn=None, force=False):
            """
            Class to handle and keep track of many pairwise comparisons of interval
            files.

            A lightweight database approach is used to minimize computational time.

            The database stores filenames and calculation timestamps;
            re-calculating a matrix using the same interval files will only
            re-calculate values for those files whose modification times are newer
            than the timestamp in the database.

            `beds` is a list of bed files.

            `genome` is the string assembly name, e.g., "hg19" or "dm3".

            `dbfn` is the filename of the database you'd like to use to track
            what's been completed.

            Example usage:

            First, get a list of bed files to use:
            #>>> beds = [
            #... pybedtools.example_filename(i) for i in  [
            #... 'Cp190_Kc_Bushey_2009.bed',
            #... 'CTCF_Kc_Bushey_2009.bed',
            #... 'SuHw_Kc_Bushey_2009.bed',
            #... 'BEAF_Kc_Bushey_2009.bed'
            #... ]]

            Set some parameters.  "dm3" is the genome to use; info will be stored
            in "ex.db".  `force=True` means to overwrite what's in the database
            #>>> # In practice, you'll want many more iterations...
            #>>> im = IntersectionMatrix(beds, 'dm3',
            #...            dbfn='ex.db', iterations=3, force=True)
            #>>> # Use 4 CPUs for randomization
            #>>> matrix = im.create_matrix(verbose=True, processes=4)
            """
            self.beds = beds
            self.genome = genome
            self.dbfn = dbfn
            self.iterations = iterations

            if self.dbfn:
                self._init_db(force)
                self.conn = sqlite3.connect(dbfn)
                self.conn.row_factory = sqlite3.Row
                self.c = self.conn.cursor()

        def _init_db(self, force=False):
            """
            Prepare the database if it doesn't already exist
            """
            if self.dbfn is None:
                return
            if os.path.exists(self.dbfn) and not force:
                return
            conn = sqlite3.connect(self.dbfn)
            c = conn.cursor()
            if force:
                c.execute('DROP TABLE IF EXISTS intersections;')
            c.executescript("""
            CREATE TABLE intersections (
                filea TEXT,
                fileb TEXT,
                timestamp FLOAT,
                actual FLOAT,
                median FLOAT,
                iterations INT,
                self INT,
                other INT,
                fractionabove FLOAT,
                fractionbelow FLOAT,
                percentile FLOAT,
                PRIMARY KEY (filea, fileb, iterations));
            """)
            conn.commit()

        def get_row(self, fa, fb, iterations):
            """
            Return the sqlite3.Row from the database corresponding to files `fa`
            and `fb`; returns None if not found.
            """
            if self.dbfn is None:
                return

            results = list(self.c.execute(
                    """
                    SELECT * FROM intersections
                    WHERE
                    filea=:fa AND fileb=:fb AND iterations=:iterations
                    """, locals()))
            if len(results) == 0:
                return
            assert len(results) == 1
            return results[0]

        def done(self, fa, fb, iterations):
            """
            Retrieves row from db and only returns True if there's something in
            there and the timestamp is newer than the input files.
            """
            row = self.get_row(fa, fb, iterations)
            if row:
                tfa = os.path.getmtime(fa)
                tfb = os.path.getmtime(fb)
                if (row['timestamp'] > tfa) and (row['timestamp'] > tfb):
                    return True
            return False

        def run_and_insert(self, fa, fb, **kwargs):
            a = pybedtools.BedTool(fa).set_chromsizes(self.genome)
            kwargs['iterations'] = self.iterations
            results = a.randomstats(fb, **kwargs)
            self.add_row(results)

        def add_row(self, results):
            """
            Inserts data into db.  `results` is a dictionary as returned by
            BedTool.randomstats with keys like::

                'iterations'
                'actual'
                'file_a'
                'file_b'
                self.fn
                other.fn
                'self'
                'other'
                'frac randomized above actual'
                'frac randomized below actual'
                'median randomized'
                'normalized'
                'percentile'
                'lower_%sth' % lower_thresh
                'upper_%sth' % upper_thresh
            """
            # translate results keys into db-friendly versions
            translations = [
                    ('file_a', 'filea'),
                    ('file_b', 'fileb'),
                    ('median randomized', 'median'),
                    ('frac randomized above actual', 'fractionabove'),
                    ('frac randomized below actual', 'fractionbelow'),
                    ]
            for orig, new in translations:
                results[new] = results[orig]

            results['timestamp'] = now()

            sql = """
            INSERT OR REPLACE INTO intersections (

                filea,
                fileb,
                timestamp,
                actual,
                median,
                iterations,
                self,
                other,
                fractionabove,
                fractionbelow,
                percentile)

                VALUES (

                :filea,
                :fileb,
                :timestamp,
                :actual,
                :median,
                :iterations,
                :self,
                :other,
                :fractionabove,
                :fractionbelow,
                :percentile)

            """
            self.c.execute(sql, results)
            self.conn.commit()

        def create_matrix(self, verbose=False, **kwargs):
            """
            Matrix (implemented as a dictionary), where the final values are
            sqlite3.ROW objects from the database::

                {
                    filea: {
                                filea: ROW,
                                fileb: ROW,
                                ...},
                    fileb: {
                                filea: ROW,
                                fileb: ROW,
                                ...},

                            }
                }
            """
            nfiles = len(self.beds)
            total = nfiles ** 2
            i = 0
            matrix = collections.defaultdict(dict)
            for fa in self.beds:
                for fb in self.beds:
                    i += 1

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

                    if not self.done(fa, fb, self.iterations):
                        self.run_and_insert(fa, fb, **kwargs)

                    matrix[get_name(fa)][get_name(fb)] = \
                            self.get_row(fa, fb, self.iterations)

            return matrix

        def print_matrix(self, matrix, key):
            """
            Prints a pairwise matrix of values. `matrix` is a dict-of-dicts from
            create_matrix(), and `key` is a field name from the database -- one of:

            ['filea', 'fileb', 'timestamp', 'actual', 'median', 'iterations',
            'self', 'other', 'fractionabove', 'fractionbelow', 'percentile']
            """