text: #!/usr/bin/env python
  
  # HiC-Pro
  # Copyleft 2015 Institut Curie
  # Author(s): Nelle Varoquaux, Nicolas Servant
  # Contact: nicolas.servant@curie.fr
  # This software is distributed without any guarantee under the terms of the
  # GNU General
  # Public License, either Version 2, June 1991 or Version 3, June 2007.
  
  """
  Script to extract restriction fragment from a fasta file and output a BED file
  """
  
  import argparse
  import re
  import os
  import sys
  import numpy as np
  
  RE_cutsite = {
      "mboi": ["^GATC"],
      "dpnii": ["^GATC"],
      "bglii": ["A^GATCT"],
      "hindiii": ["A^AGCTT"]}
  
  
  def find_re_sites(filename, sequences, offset):
      with open(filename, 'r') as infile:
          chr_id = None
          big_str = ""
          indices = []
          all_indices = []
          contig_names = []
          c = 0
          for line in infile:
              c += 1
              if line.startswith(">"):
                  print("{}...".format(line.split()[0][1:]))
                  # If this is not the first chromosome, find the indices and append
                  # them to the list
                  if chr_id is not None:
                       for rs in range(len(sequences)):
                           pattern = "(?={})".format(sequences[rs].lower())
                           indices += [m.start() + offset[rs]\
                           for m in re.finditer(pattern, big_str)]
                       indices.sort()
                       all_indices.append(indices)
                       indices = []
  
                  # This is a new chromosome. Empty the sequence string, and add the
                  # correct chrom id
                  big_str = ""
                  chr_id = line.split()[0][1:]
                  if chr_id in contig_names:
                      print("The fasta file contains several instance of {}. Exit.".format(chr_id))
                      sys.exit(-1)
                  contig_names.append(chr_id)
              else:
                  # As long as we don't change chromosomes, continue reading the
                  # file, and appending the sequences
                  big_str += line.lower().strip()
          # Add the indices for the last chromosome
          for rs in range(len(sequences)):
              pattern = "(?={})".format(sequences[rs].lower())
              indices += [m.start() + offset[rs]
                          for m in re.finditer(pattern, big_str)]
          indices.sort()
          all_indices.append(indices)
      
      return contig_names, all_indices
  
  
  def find_chromsomose_lengths(reference_filename):
      chromosome_lengths = []
      chromosome_names = []
      length = None
      with open(reference_filename, 'r') as infile:
          for line in infile:
              if line.startswith(">"):
                  chromosome_names.append(line[1:].strip())
                  if length is not None:
                      chromosome_lengths.append(length)
                  length = 0
              else:
                  length += len(line.strip())
          chromosome_lengths.append(length)
      return chromosome_names, np.array(chromosome_lengths)
  
  
  def replaceN(cs):
      npos = int(cs.find('N'))
      cseql = []
      if npos != -1:
          for nuc in ["A","C","G","T"]:
              tmp = cs.replace('N', nuc, 1)
              tmpl = replaceN(tmp)
              if type(tmpl) == list:
                  cseql = cseql + tmpl
              else:
                  cseql.append(tmpl)
      else:
          cseql.append(cs)
      return cseql
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument('fastafile')
      parser.add_argument('-r', '--restriction_sites',
                          dest='res_sites',
                          nargs='+',
                          help=("The cutting position has to be specified using "
                                "'^'. For instance, -r A^AGCTT for HindIII "
                                "digestion. Several restriction enzyme can be "
                                "specified."))
      parser.add_argument('-o', '--out', default=None)
      args = parser.parse_args()
  
      filename = args.fastafile
      out = args.out
      
      # Split restriction sites if comma-separated
      cutsites=[]
      for s in args.res_sites:
          for m in s.split(','):
              cutsites.append(m)
                  
      # process args and get restriction enzyme sequences
      sequences = []
      offset = []
      for cs in cutsites:
          if cs.lower() in RE_cutsite:
              cseq = ''.join(RE_cutsite[cs.lower()])
          else:
              cseq = cs
  
          offpos = int(cseq.find('^'))
          if offpos == -1:
              print("Unable to detect offset for {}. Please, use '^' to specify the cutting position,\
                     i.e A^GATCT for HindIII digestion.".format(cseq))
              sys.exit(-1)
  
          for nuc in list(set(cs)):
              if nuc not in ['A','T','G','C','N','^']:
                  print("Find unexpected character ['{}']in restriction motif".format(nuc))
                  print("Note that multiple motifs should be separated by a space (not a comma !)")
  
                  sys.exit(-1)
  
          offset.append(offpos)
          sequences.append(re.sub('\^', '', cseq))
  
      # replace all N in restriction motif
      sequences_without_N = []
      offset_without_N = []
      for rs in range(len(sequences)):
          nrs = replaceN(sequences[rs])
          sequences_without_N = sequences_without_N + nrs
          offset_without_N = offset_without_N + [offset[rs]] * len(nrs)
            
      sequences = sequences_without_N
      offset = offset_without_N
      
      if out is None:
          out = os.path.splitext(filename)[0] + "_fragments.bed"
  
      print("Analyzing", filename)
      print("Restriction site(s)", ",".join(sequences))
      print("Offset(s)",  ','.join(str(x) for x in offset))
  
      # Read fasta file and look for rs per chromosome
      contig_names, all_indices = find_re_sites(filename, sequences,  offset=offset)
      _, lengths = find_chromsomose_lengths(filename)
  
      valid_fragments = []
      for i, indices in enumerate(all_indices):
          valid_fragments_chr = np.concatenate(
              [np.concatenate([[0], indices])[:, np.newaxis],
               np.concatenate([indices, [lengths[i]]])[:, np.newaxis]],
              axis=1)
          valid_fragments.append(valid_fragments_chr)
  
      # Write results
      print("Writing to {} ...".format(out))
      with open(out, 'w') as outfile:
          for chrom_name, indices in zip(contig_names, valid_fragments):
              frag_id = 0
              for begin, end in indices:
                  # allow to remove cases where the enzyme cut at
                  # the first position of the chromosome
                  if end > begin:
                      frag_id += 1
                      frag_name = "HIC_{}_{}".format(str(chrom_name), int(frag_id))
                      outfile.write("{}\t{}\t{}\t{}\t0\t+\n".format(str(chrom_name), int(begin), int(end), str(frag_name)))
  