#!/usr/bin/env python
import argparse, sys, itertools
from Bio import Restriction
from Bio import SeqIO
import numpy as numpy
###########################################################
# Get command line arguments
def getParser():
    parser = argparse.ArgumentParser(description='Counts the number of restriction sites for any given interval.')
    parser.add_argument('-f',type=str,dest="fFile",help="FASTA file.")
    parser.add_argument('-r',type=str,dest="rFile",help="BED file. Each line contain the region under which the restriction sites will be counted")
    parser.add_argument('-o',type=str,dest="oFile",help="STR. Name of output file.")

    if len(sys.argv) == 1:
        print >> sys.stderr,parser.print_help()
        exit(0)
    return parser
#########################################################
def main():
    args = getParser().parse_args()
    fFile = args.fFile  
    rFile = args.rFile  
    oFile = args.oFile 
    # Get restrcition sites
    restr_sites={}
    handle = open(fFile, "rU")
    for record in SeqIO.parse(handle, "fasta"):
        print record.id
        sites=Restriction.HaeIII.search(record.seq)
        restr_sites[record.id]=numpy.array(sites)
        #distances=[t - s for s, t in zip(sites, sites[1:])]
    handle.close()
    # Print results
    out=open(oFile,"w")
    for n,line in enumerate(open(rFile,"r")):
        line=line.strip().split("\t")
        chrom=line[0]
        start=int(line[1])
        end=int(line[2])
        count=sum(numpy.logical_and(restr_sites[chrom]>start, restr_sites[chrom]<=end))
        print >>out, "\t".join(map(str,[chrom,start,end,count]))
    out.close()
##########################################################
if __name__ == '__main__':
    main()

