##################################
#                                #
# Last modified 2017/05/09       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math
import os	

def run():

    if len(sys.argv) < 5:
        print 'usage: python %s FIMO.txt_file1 FIMO.txt_file2 minDistance maxDistance outfilename' % sys.argv[0]
        print '\tNote: the FIMO files can be in .gz or .bzip2 format'
        sys.exit(1)

    FIMO1 = sys.argv[1]
    FIMO2 = sys.argv[2]
    minDist = int(sys.argv[3])
    maxDist = int(sys.argv[4])
    outfilename = sys.argv[5]

    maxMotLen = 0

    FIMO1Dict = {}
    if FIMO1.endswith('.bz2'):
        cmd = 'bzip2 -cd ' + FIMO1
    elif FIMO1.endswith('.gz'):
        cmd = 'gunzip -c ' + FIMO1
    else:
        cmd = 'cat ' + FIMO1
    p = os.popen(cmd, "r")
    line = 'line'
    while line != '':
        line = p.readline()
        if line == '':
            break
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        mot = fields[0]
        chr = fields[1] 
        left = int(fields[2])
        right = int(fields[3])
        if (right - left) > maxMotLen:
            maxMotLen = (right - left)
        strand = fields[4]
        score = fields[5]
        pval = fields[6]
        if FIMO1Dict.has_key(chr):
            pass
        else:
            FIMO1Dict[chr] = {}
        if FIMO1Dict[chr].has_key(left):
            pass
        else:
            FIMO1Dict[chr][left]=[]
        motif = (mot,left,right,strand,score,pval)
        FIMO1Dict[chr][left].append(motif)

    FIMO2Dict = {}
    if FIMO2.endswith('.bz2'):
        cmd = 'bzip2 -cd ' + FIMO2
    elif FIMO2.endswith('.gz'):
        cmd = 'gunzip -c ' + FIMO2
    else:
        cmd = 'cat ' + FIMO2
    p = os.popen(cmd, "r")
    line = 'line'
    while line != '':
        line = p.readline()
        if line == '':
            break
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        mot = fields[0]
        chr = fields[1] 
        left = int(fields[2])
        right = int(fields[3])
        if (right - left) > maxMotLen:
            maxMotLen = (right - left)
        strand = fields[4]
        score = fields[5]
#        pval = fields[6]
        if FIMO2Dict.has_key(chr):
            pass
        else:
            FIMO2Dict[chr] = {}
        if FIMO2Dict[chr].has_key(left):
            pass
        else:
            FIMO2Dict[chr][left]=[]
#        motif = (mot,left,right,strand,score,pval)
        motif = (mot,left,right,strand,score)
        FIMO2Dict[chr][left].append(motif)

    outfile = open(outfilename, 'w')

    outline = '#chr\tmotif1\tleft\tright\tstrand\tscore\t\tmotif2\tleft\tright\tstrand\tscore'
    outfile.write(outline + '\n')
    for chr in FIMO1Dict.keys():
        if FIMO2Dict.has_key(chr):
            pass
        else:
            continue
        for pos in FIMO1Dict[chr].keys():
            for motif in FIMO1Dict[chr][pos]:
                mot1left = motif[1]
                mot1right = motif[2]
                motlen = (mot1right - mot1left)
                for i in range(mot1left - maxDist - maxMotLen, mot1left + motlen + maxDist):
                    if FIMO2Dict[chr].has_key(i):
                        for motif2 in FIMO2Dict[chr][i]:
                            if motif2 == motif:
                                continue
                            mot2left = motif2[1]
                            mot2right = motif2[2]
                            if mot2left < mot1left:
                                if (mot1left - mot2right) <= maxDist and (mot1left - mot2right) >= minDist:
                                    outline = chr + '\t' + motif[0] + '\t' + str(motif[1]) + '\t' + str(motif[2]) + '\t' + str(motif[3]) + '\t' + str(motif[4])
                                    outline = outline + '\t' + motif2[0] + '\t' + str(motif2[1]) + '\t' + str(motif2[2]) + '\t' + str(motif2[3]) + '\t' + str(motif2[4])
                                    outfile.write(outline + '\n')
                            if mot2left >= mot1left:
                                if (mot2left - mot1right) <= maxDist and (mot2left - mot1right) >= minDist:
                                    outline = chr + '\t' + motif[0] + '\t' + str(motif[1]) + '\t' + str(motif[2]) + '\t' + str(motif[3]) + '\t' + str(motif[4])
                                    outline = outline + '\t' + motif2[0] + '\t' + str(motif2[1]) + '\t' + str(motif2[2]) + '\t' + str(motif2[3]) + '\t' + str(motif2[4])
                                    outfile.write(outline + '\n')

    outfile.close()
   
run()