# When you have chr1_bed,... in data/:
# python main.py mm10.chrom.sizes data/
#
# When you only have a bed.txt, no chr1_bed in data/
# python main.py mm10.chrom.sizes data/ bed.txt
import os
import sys

def bed(directory, input):
    for filename in os.listdir(directory):
        if filename.endswith('bed'):
            os.remove(directory+filename)

    with open(input, 'r') as f:
        filename = ''
        filename_content = ''
        for line_number, line in enumerate(f):
            if line.strip() == '':
                continue
            l = line.split()[:3]
            if '_' in l[0]:
                continue
            if line_number != 0 and l[0] != filename:
                with open(directory+filename+'_bed', 'w') as x:
                    x.write(filename_content)
                print filename+' writing done...'
                filename_content = ''
            filename = l[0].strip()
            filename_content += ' '.join(l)+'\n'

        with open(directory+filename+'_bed', 'w') as x:
            x.write(filename_content)


def print_count(directory, file, n):
    dict = {}
    T = 50
    if not os.path.isfile(directory+file+'_bed'):
        return
    with open(directory+file+'_bed', 'r') as f:
        for line in f:
            if line.strip() != '':
                l = line.strip().split()
                l1 = int(int(l[1])/T)
                dict[l1] = dict.get(l1, 0) + 1
    with open(directory+file+'_count', 'w') as f:
        for i in range(int(n/T)):
            f.write("%s\t%d\t%d\t%d\n" %(file,i*T,(i+1)*T-1,dict.get(i,0)))
        i = int(n/T)
        f.write("%s\t%d\t%d\t%d\n" %(file,i*T,n,dict.get(i,0)))

def main():
    if len(sys.argv) == 4:
        bed(sys.argv[2], sys.argv[3])

    with open(sys.argv[1], 'r') as f:
        for line in f:
            if line.strip() != '':
                l = line.strip().split()
                if '_' in l[0]:
                    continue
                print_count(sys.argv[2], l[0],int(l[1]))


main()