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

import sys
import string
import random
from sets import Set

def run():

    if len(sys.argv) < 3:
        print 'usage: python %s wig chrom.sizes output [-shift bp]' % sys.argv[0]
        sys.exit(1)

    wiggle = sys.argv[1]
    chromSizes = sys.argv[2]
    output = sys.argv[3]

    doShift = False
    if '-shift' in sys.argv:
        doShift = True
        shift = int(sys.argv[sys.argv.index('-shift') + 1])

    chromSizesDict = {}
    linelist = open(chromSizes)
    for line in linelist:
        fields = line.strip().split('\t')
        chr = fields[0]
        chromSizesDict[chr] = int(fields[1])

    outfile = open(output, 'w')

    linelist = open(wiggle)
    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])
        if doShift:
            pos += shift
        score = float(fields[1])
        outline = chr + '\t' + str(pos) + '\t' + str(min(pos + span,chromSizesDict[chr])) + '\t'+ str(score)
        outfile.write(outline + '\n')

    outfile.close()

run()
