# Source code for pybedtools.contrib.long_range_interaction

    from __future__ import print_function
    import os
    import sys
    import itertools
    import six
    import time
    import pysam
    import pybedtools

    def tag_bedpe(bedpe, queries, verbose=False):
        """
        Tag each end of a BEDPE with a set of (possibly many) query BED files.

        For example, given a BEDPE of interacting fragments from a Hi-C experiment,
        identify the contacts between promoters and ChIP-seq peaks. In this case,
        promoters and ChIP-seq peaks of interest would be provided as BED files.

        The strategy is to split the BEDPE into two separate files.  Each file is
        intersected independently with the set of queries.  The results are then
        iterated through in parallel to tie the ends back together. It is this
        iterator that is returned (see example below).

        Parameters
        ----------

        bedpe : str
            BEDPE-format file. Must be name-sorted.

        queries : dict
            Dictionary of BED/GFF/GTF/VCF files to use. After splitting the BEDPE,
            these query files (values in the dictionary) will be passed as the `-b`
            arg to `bedtools intersect`. The keys are passed as the `names`
            argument for `bedtools intersect`

            *Features in each file must have unique names*. Use
            :func:`pybedtools.featurefuncs.UniqueID` to help fix this.

            *Each file must be BED3 to BED6*.

        Returns
        -------
        Tuple of (iterator, n, extra).

        `iterator` is described below. `n` is the total number of lines in the
        BEDPE file, which is useful for calculating percentage complete for
        downstream work. `extra` is the number of extra fields found in the BEDPE
        (also useful for downstream processing).

        `iterator` yields tuples of (label, end1_hits, end2_hits) where `label` is
        the name field of one line of the original BEDPE file. `end1_hits` and
        `end2_hits` are each iterators of BED-like lines representing all
        identified intersections across all query BED files for end1 and end2 for
        this pair.

        Recall that BEDPE format defines a single name and a single score for each
        pair. For each item in `end1_hits`, the fields are::

            chrom1
            start1
            end1
            name
            score
            strand1
            [extra fields]
            query_label
            fields_from_query_intersecting_end1

        where `[extra fields]` are any additional fields from the original BEDPE,
        `query_label` is one of the keys in the `beds` input dictionary, and the
        remaining fields in the line are the intersecting line from the
        corresponding BED file in the `beds` input dictionary.

        Similarly, each item in `end2_hits` consists of:

            chrom2
            start2
            end2
            name
            score
            strand2
            [extra fields]
            query_label
            fields_from_query_intersecting_end2

        At least one line is reported for every line in the BEDPE file. If there
        was no intersection, the standard BEDTools null fields will be shown. In
        `end1_hits` and `end2_hits`, a line will be reported for each hit in each
        query.

        Example
        -------

        Consider the following BEDPE (where "x1" is an aribtrary extra field).

        >>> bedpe = pybedtools.example_bedtool('test_bedpe.bed')
        >>> print(bedpe) # doctest: +NORMALIZE_WHITESPACE
        chr1  1  10  chr1  50   90   pair1  5  +  -  x1
        chr1  2  15  chr1  200  210  pair2  1  +  +  y1
        <BLANKLINE>

        And the following transcription start sites (TSSes) in BED4 format:

        >>> tsses = pybedtools.example_bedtool('test_tsses.bed')
        >>> print(tsses) # doctest: +NORMALIZE_WHITESPACE
        chr1  5   6   gene1
        chr1  60  61  gene2
        chr1  88  89  gene3
        <BLANKLINE>

        And the following called peaks as BED6:

        >>> peaks = pybedtools.example_bedtool('test_peaks.bed')
        >>> print(peaks) # doctest: +NORMALIZE_WHITESPACE
        chr1  3  4  peak1  50  .
        <BLANKLINE>

        Then we can get the following iterator, n, and extra. Note that the
        OrderedDict is only for testing to ensure output is always consistend; in
        practice a regular dictionary is fine:

        >>> from pybedtools.contrib.long_range_interaction import tag_bedpe
        >>> from collections import OrderedDict
        >>> queries = OrderedDict()
        >>> queries['tss'] = tsses
        >>> queries['pk'] = peaks
        >>> iterator, n, extra = tag_bedpe(bedpe, queries)
        >>> print(n)
        2
        >>> print(extra)
        1

        The following illustrates that each item in the iterator represents one
        pair, and each item in each group represents an intersection with one end.
        Note that the sorting is necessary only for the doctests to be output in
        consistent format; this not typically needed:

        >>> for (label, end1_hits, end2_hits) in iterator:
        ...    end1_hits = sorted(end1_hits, key=lambda x: str(x))
        ...    end2_hits = sorted(end2_hits, key=lambda x: str(x))
        ...    print('PAIR = {}'.format(label))
        ...    print('end1_hits:')
        ...    for i in end1_hits:
        ...        print(i, end='')
        ...    print('end2_hits:')
        ...    for i in end2_hits:
        ...        print(i, end='')  # doctest: +NORMALIZE_WHITESPACE
        PAIR = pair1
        end1_hits:
        chr1       1       10      pair1   5       +       x1      pk      chr1    3       4       peak1   50      .       1
        chr1       1       10      pair1   5       +       x1      tss     chr1    5       6       gene1   1
        end2_hits:
        chr1       50      90      pair1   5       -       x1      tss     chr1    60      61      gene2   1
        chr1       50      90      pair1   5       -       x1      tss     chr1    88      89      gene3   1
        PAIR = pair2
        end1_hits:
        chr1       2       15      pair2   1       +       y1      pk      chr1    3       4       peak1   50      .       1
        chr1       2       15      pair2   1       +       y1      tss     chr1    5       6       gene1   1
        end2_hits:
        chr1       200     210     pair2   1       +       y1      .       .       -1      -1      .       0

        See the `cis_trans_interactions()` function for one way of summarizing
        these data.
        """
        b = pybedtools.BedTool(bedpe)

        # Figure out if the supplied bedpe had any extra fields. If so, the fields
        # are repeated in each of the split output files.
        observed = b.field_count()
        extra = observed - 10
        extra_inds = [10 + i for i in range(extra)]

        end1_fn = pybedtools.BedTool._tmp()
        end2_fn = pybedtools.BedTool._tmp()

        # Performance notes:
        # We don't need the overhead of converting every line into
        # a pybedtools.Interval object just so we can grab the fields. Doing so
        # takes 3.5x more time than simply splitting each line on a tab.
        if verbose:
            print('splitting BEDPE into separate files.')
            print('end1 is going to %s' % end1_fn)
            print('end2 is going to %s' % end2_fn)

        n = 0
        with open(end1_fn, 'w') as end1_out, open(end2_fn, 'w') as end2_out:
            for line in open(b.fn):
                n += 1
                f = line.strip().split('\t')
                end1_out.write(
                    '\t'.join(
                        (f[i] for i in [0, 1, 2, 6, 7, 8] + extra_inds)) + '\n')
                end2_out.write(
                    '\t'.join(
                        (f[i] for i in [3, 4, 5, 6, 7, 9] + extra_inds)) + '\n')

        # Performance notes:
        #
        # For small BEDPE and large set of query files, it would be faster to sort
        # these independently, intersect with sorted=True, and then re-sort by name
        # for the grouping. For large BEDPE, I don't think the sorted=True
        # performance gain outweighs the hit from sorting twice.
        #
        # On the other hand, if BEDPE was coord-sorted in the first place, only
        # end2 would need to be sorted and re-sorted. On the other (third!?) hand,
        # BEDPE creation from BAM implies name-sorting, so it's probably not
        # reasonable to assume coord-sorted.
        #
        # In the end: don't do any sorting.

        end1_bt = pybedtools.BedTool(end1_fn)
        end2_bt = pybedtools.BedTool(end2_fn)
        names, fns = [], []
        for name, fn in queries.items():
            names.append(name)
            if isinstance(fn, pybedtools.BedTool):
                fns.append(fn.fn)
            else:
                fns.append(fn)

        if verbose:
            print('intersecting end 1')
        end1_hits = end1_bt.intersect(list(fns), names=names, wao=True)
        if verbose:
            print('intersecting end 2')
        end2_hits = end2_bt.intersect(list(fns), names=names, wao=True)
        if verbose:
            print('intersection with end1 is in %s' % (end1_hits.fn))
            print('intersection with end2 is in %s' % (end2_hits.fn))

        grouped_end1 = itertools.groupby(end1_hits, lambda f: f[3])
        grouped_end2 = itertools.groupby(end2_hits, lambda f: f[3])

        def gen():
            for (label1, group1), (label2, group2) \
                in six.moves.zip(grouped_end1, grouped_end2):
                assert label1 == label2
                yield label1, group1, group2

        return gen(), n, extra

    def cis_trans_interactions(iterator, n, extra, verbose=True):
        """
        Converts the output from `tag_bedpe` into a pandas DataFrame containing
        information about regions that contact each other in cis (same fragment) or
        trans (different fragments).

        For example, given a BEDPE file representing 3D interactions in the genome,
        we want to identify which transcription start sites are connected to distal
        regions containing a peak.

        >>> bedpe = pybedtools.example_bedtool('test_bedpe.bed')
        >>> print(bedpe) # doctest: +NORMALIZE_WHITESPACE
        chr1  1  10  chr1  50   90   pair1  5  +  -  x1
        chr1  2  15  chr1  200  210  pair2  1  +  +  y1
        <BLANKLINE>

        >>> tsses = pybedtools.example_bedtool('test_tsses.bed')
        >>> print(tsses) # doctest: +NORMALIZE_WHITESPACE
        chr1  5   6   gene1
        chr1  60  61  gene2
        chr1  88  89  gene3
        <BLANKLINE>

        >>> peaks = pybedtools.example_bedtool('test_peaks.bed')
        >>> print(peaks) # doctest: +NORMALIZE_WHITESPACE
        chr1  3  4  peak1  50  .
        <BLANKLINE>

        Here's what the tracks look like. Note that pair1 is evidence of
        a gene1-gene2 interaction and a gene1-gene3 interaction::

            TRACKS:

                          1         2  /         5         6    / 8         9     / 20
                0123456789012345678901 / 2345678901234567890123 / 012345678901234 / 0123456789
           pair1 |||||||||------------ / --------|||||||||||||| / ||||||||||||||| /
           pair2  |||||||||||||------- / ---------------------- / --------------- / ||||||||||
           tsses     1                 /                   2    /         3
           peaks   1

        >>> from collections import OrderedDict
        >>> queries = OrderedDict()
        >>> queries['tss'] = tsses
        >>> queries['pk'] = peaks
        >>> iterator, n, extra = tag_bedpe(bedpe, queries)
        >>> for (label, group1, group2) in iterator:
        ...    group1 = sorted(group1, key=lambda x: str(x))
        ...    group2 = sorted(group2, key=lambda x: str(x))
        ...    for i in group1:
        ...        print(i, end='')  # doctest: +NORMALIZE_WHITESPACE
        ...    for i in group2:
        ...        print(i, end='')  # doctest: +NORMALIZE_WHITESPACE
        chr1       1       10      pair1   5       +       x1      pk      chr1    3       4       peak1   50      .       1
        chr1       1       10      pair1   5       +       x1      tss     chr1    5       6       gene1   1
        chr1       50      90      pair1   5       -       x1      tss     chr1    60      61      gene2   1
        chr1       50      90      pair1   5       -       x1      tss     chr1    88      89      gene3   1
        chr1       2       15      pair2   1       +       y1      pk      chr1    3       4       peak1   50      .       1
        chr1       2       15      pair2   1       +       y1      tss     chr1    5       6       gene1   1
        chr1       200     210     pair2   1       +       y1      .       .       -1      -1      .       0

        Now we run the same thing, but now aggregate it. Note that each piece of
        interaction evidence has its own line. The first line shows that pair1 has
        gene1 and peak1 in the same fragment, and that they are connected to gene2.
        The second line shows again that gene1 and peak1 are in the same fragmet
        and that they are also connected to gene3:

        >>> import pandas; pandas.set_option('display.max_columns', 10)
        >>> iterator, n, extra = tag_bedpe(bedpe, {'tss': tsses, 'pk': peaks})
        >>> df =  cis_trans_interactions(iterator, n, extra)
        >>> print(df.sort_values(list(df.columns)).reset_index(drop=True))
          target_label target_name cis_label cis_name distal_label distal_name  label
        0           pk       peak1       tss    gene1            .           .  pair2
        1           pk       peak1       tss    gene1          tss       gene2  pair1
        2           pk       peak1       tss    gene1          tss       gene3  pair1
        3          tss       gene1        pk    peak1            .           .  pair2
        4          tss       gene1        pk    peak1          tss       gene2  pair1
        5          tss       gene1        pk    peak1          tss       gene3  pair1
        6          tss       gene2       tss    gene3           pk       peak1  pair1
        7          tss       gene2       tss    gene3          tss       gene1  pair1
        8          tss       gene3       tss    gene2           pk       peak1  pair1
        9          tss       gene3       tss    gene2          tss       gene1  pair1

        If we only care about genes:

        >>> print((df[df.target_label == 'tss']).sort_values(list(df.columns)).reset_index(drop=True))
          target_label target_name cis_label cis_name distal_label distal_name  label
        0          tss       gene1        pk    peak1            .           .  pair2
        1          tss       gene1        pk    peak1          tss       gene2  pair1
        2          tss       gene1        pk    peak1          tss       gene3  pair1
        3          tss       gene2       tss    gene3           pk       peak1  pair1
        4          tss       gene2       tss    gene3          tss       gene1  pair1
        5          tss       gene3       tss    gene2           pk       peak1  pair1
        6          tss       gene3       tss    gene2          tss       gene1  pair1

        Note that in pair2, there is no evidence of interaction between gene1 and
        gene2.

        What interacts distally with gene2's TSS?

        >>> assert set(df.loc[df.target_name == 'gene2', 'distal_name']).difference('.') == set([u'gene1', u'peak1'])

        """
        try:
            import pandas
        except ImportError:
            raise ImportError("pandas must be installed to use this function")
        c = 0
        lines = []
        for label, end1_hits, end2_hits in iterator:
            c += 1
            if c % 1000 == 0:
                print('%d (%.1f%%)\r' % (c, c / float(n) * 100), end='')
                sys.stdout.flush()

            # end1_hits has the full lines of all intersections with end1
            end1_hits = list(end1_hits)
            end2_hits = list(end2_hits)

            def get_name_hits(f):
                """
                Returns the key (from which file the interval came) and the name
                (of the individual feature).
                """
                # this is the "name" reported if there was no hit.
                if f[6 + extra] == '.':
                    return ('.', '.')
                interval = pybedtools.create_interval_from_list(f[7 + extra:])
                return ([f[6 + extra], interval.name])

            names1 = set(
                six.moves.map(tuple, six.moves.map(get_name_hits, end1_hits)))
            names2 = set(
                six.moves.map(tuple, six.moves.map(get_name_hits, end2_hits)))

            for cis, others in [(names1, names2), (names2, names1)]:
                for target in cis:
                    if target == ('.', '.'):
                        continue
                    non_targets = set(cis).difference([target])
                    if len(non_targets) == 0:
                        non_targets = [('.', '.')]
                    for non_target in non_targets:
                        for other in others:
                            line = []
                            line.extend(target)
                            line.extend(non_target)
                            line.extend(other)
                            line.append(label)
                            lines.append(line)

        df = pandas.DataFrame(
            lines,
            columns=['target_label', 'target_name', 'cis_label', 'cis_name',
                     'distal_label', 'distal_name', 'label'])
        df = df.drop_duplicates()
        return df