##################################
#                                #
# Last modified 2024/08/23       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math
from sets import Set

def run():

    if len(sys.argv) < 6:
        print 'usage: python %s TSS-bed chrFieldID1 Enhancers-bed chrFieldID2 genomic_span outfile' % sys.argv[0]
        print '\tlist_of_files format: label <tab> file_name'
        sys.exit(1)

    TSS = sys.argv[1]
    chrFieldID1 = int(sys.argv[2])
    Enh = sys.argv[3]
    chrFieldID2 = int(sys.argv[4])
    span = int(sys.argv[5])
    outfilename = sys.argv[6]

    TSSDict = {}
    linelist = open(TSS)
    for line in linelist:
        fields = line.strip().split('\t')
        if line.startswith('#'):
            continue
        chr = fields[chrFieldID1]
        left = int(fields[chrFieldID1 + 1])
        right = int(fields[chrFieldID1 + 2])
        if TSSDict.has_key(chr):
            pass
        else:
            TSSDict[chr] = {}
        TSSDict[chr][(left,right)] = fields

    EnhDict = {}
    linelist = open(Enh)
    for line in linelist:
        fields = line.strip().split('\t')
        if line.startswith('#'):
            continue
        chr = fields[chrFieldID1]
        left = int(fields[chrFieldID1 + 1])
        right = int(fields[chrFieldID1 + 2])
        if EnhDict.has_key(chr):
            pass
        else:
            EnhDict[chr] = {}
        EnhDict[chr][(left,right)] = fields

    outfile = open(outfilename,'w')

    for chr in TSSDict.keys():
        print chr
        for (left,right) in TSSDict[chr].keys():
            if EnhDict.has_key(chr):
                pass
            else:
                continue
            for (L,R) in EnhDict[chr].keys():
                if (L <= right and L >= left) or (R <= right and R >= left):
                    continue
                if math.fabs(L-left) < span and math.fabs(R-left) < span and math.fabs(L-right) < span and math.fabs(R-right) < span:
                    outline = ''
                    for F in TSSDict[chr][(left,right)]:
                        outline += F
                        outline += '\t'
                    for F in EnhDict[chr][(L,R)]:
                        outline += F
                        outline += '\t'
                    outfile.write(outline.strip() + '\n')

    outfile.close()
        
run()