##################################
#                                #
# Last modified 11/19/2012       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math

def run():

    if len(sys.argv) < 5:
        print 'usage: python %s chrom.sizes window_size bed chrFieldID outfilename [-rescaleOutput factor]' % sys.argv[0]
        sys.exit(1)
    
    chromsizes = sys.argv[1]
    window = int(sys.argv[2])
    bed = sys.argv[3]
    chrFieldID = int(sys.argv[4])
    outfilename = sys.argv[5]

    doRescale = False
    if '-rescaleOutput' in sys.argv:
        doRescale = True
        recaleFactor = float(sys.argv[sys.argv.index('-rescaleOutput') + 1])

    outfile = open(outfilename, 'w')

    ScoreDict = {}
    chrDict = {}
    
    linelist  = open(chromsizes)
    for line in linelist:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        chr = fields[0]
        length = int(fields[1])
        ScoreDict[chr]={}
        chrDict[chr]=length
        for i in range(0,length,window):
            left = i
            right = min(length,i+window)
            ScoreDict[chr][(left,right)]=0

    linelist  = open(bed)
    k=0
    for line in linelist:
        k+=1
        if k % 10000 == 0:
            print k
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        chr = fields[chrFieldID]
        left = int(fields[chrFieldID+1])
        right = int(fields[chrFieldID+2])
        for i in range(left,right):
            windowleft = (i/window)*window
            ScoreDict[chr][(windowleft,min(windowleft+window,chrDict[chr]))]+=1

    chrKeys = ScoreDict.keys()
    chrKeys.sort()

    for chr in chrKeys:
        keys = ScoreDict[chr].keys()
        keys.sort()
        for (left,right) in keys:
            if doRescale:
                outline = chr + '\t' + str(left) + '\t' + str(right) + '\t' + str(recaleFactor*(ScoreDict[chr][(left,right)]/(window+0.0)))
            else:
                outline = chr + '\t' + str(left) + '\t' + str(right) + '\t' + str(ScoreDict[chr][(left,right)]/(window+0.0))
            outfile.write(outline + '\n')

    outfile.close()

run()
