##################################
#                                #
# Last modified 07/21/2014       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import random
from sets import Set

def run():

    if len(sys.argv) < 3:
        print 'usage: python %s bed chrFieldID wigFile [-varStepWigg]' % sys.argv[0]
        print '\tThis script will calculate the fraction of signal that is in the regions of the supplied BED file'
        print '\tIt will print the answer to stdout'
        sys.exit(1)

    input = sys.argv[1]
    chfFieldID = int(sys.argv[2])
    wiggle = sys.argv[3]

    WantedDict={}

    doVarStepWigg = False
    if '-varStepWigg' in sys.argv:
       doVarStepWigg = True

    linelist = open(input)
    i=0
    for line in linelist:
        if line.startswith('#'):
            continue
        i+=1
        if i % 1000000 == 0:
            print i, 'lines processed'
        fields = line.strip().split('\t')
        chr=fields[chfFieldID]
        left = int(fields[chfFieldID+1])
        right = int(fields[chfFieldID+2])
        if WantedDict.has_key(chr):
            pass
        else:
            WantedDict[chr]={}
        for j in range(left,right):
            WantedDict[chr][j]=0

    TotalScore = 0
    WantedScore = 0

    linelist = open(wiggle)
    if doVarStepWigg:
        for line in linelist:
            if line.startswith('#'):
                continue
            if line.startswith('track'):
                continue
            if line.startswith('variableStep'):
                chr = line.strip().split(' ')[1].split('=')[1]
                span = int(line.strip().split(' ')[2].split('=')[1])
                continue
            fields = line.strip().split('\t')
            pos = int(fields[0])
            score = float(fields[1])
            for j in range(pos,pos + span):
                TotalScore += score
                if WantedDict[chr].has_key(j):
                    WantedScore += score
    else:
        for line in linelist:
            if line.startswith('#'):
                continue
            fields = line.strip().split('\t')
            if len(fields) == 1:
                fields = line.strip().split(' ')
            chr = fields[0]
            left = int(fields[1])
            right = int(fields[2])
            score = float(fields[3])
            if WantedDict.has_key(chr):
                pass
            else:
                continue
            for j in range(left,right):
                TotalScore += score
                if WantedDict[chr].has_key(j):
                    WantedScore += score

    outline = wiggle + '\t' + str(WantedScore) + '\t' + str(TotalScore) + '\t' + str(WantedScore/TotalScore)
    print outline

run()
