text: #!/usr/bin/env python
  #
  #     fastq_utils.py: utility functions for operating on fastq files
  #     Copyright (C) University of Manchester 2016-2020 Peter Briggs
  #
  ########################################################################
  #
  # fastq_utils.py
  #
  #########################################################################
  
  """
  Utility classes and functions for operating on Fastq files:
  
  - BaseFastqAttrs: base class for extracting info from Fastq file
  - IlluminaFastqAttrs: class for extracting info from Illumina Fastqs
  - FastqReadCounter: implements various methods for counting reads
    in FASTQ files
  - assign_barcodes_single_end: extract and assign inline barcodes
  - get_read_number: get the read number (1 or 2) from a Fastq file
  - get_read_count: count total reads across one or more Fastqs
  - pair_fastqs: automagically pair up FASTQ files
  - pair_fastqs_by_name: pair up FASTQ files based on their names
  - group_fastqs_by_name: group FASTQ files based on their names
    (more general version of 'pair_fastqs_by_name' which can handle
    arbitrary collections of read IDs)
  - remove_index_fastqs: remove index (I1/I2) Fastqs from a list
  """
  
  #######################################################################
  # Imports
  #######################################################################
  
  import os
  import gzip
  import subprocess
  import logging
  from bcftbx.FASTQFile import FastqIterator
  from bcftbx.FASTQFile import get_fastq_file_handle
  from bcftbx.FASTQFile import nreads
  from bcftbx.qc.report import strip_ngs_extensions
  
  #######################################################################
  # Classes
  #######################################################################
  
  class BaseFastqAttrs(object):
      """
      Base class for extracting information about a Fastq file
  
      Instances of this class provide the follow attributes:
  
      fastq:            the original fastq file name
      basename:         basename with NGS extensions stripped
      extension:        full extension e.g. '.fastq.gz'
      sample_name:      name of the sample
      sample_number:    integer (or None if no sample number)
      barcode_sequence: barcode sequence (string or None)
      lane_number:      integer (or None if no lane number)
      read_number:      integer (or None if no read number)
      set_number:       integer (or None if no set number)
      is_index_read:    boolean (True if index read, False if not)
  
      Subclasses should process the supplied Fastq name and set
      these attributes appropriately.
      """
      def __init__(self,fastq):
          # Store name
          self.fastq = fastq
          # Basename and extension
          self.basename = strip_ngs_extensions(self.fastq)
          self.extension = self.fastq[len(self.basename):]
          self.basename = os.path.basename(self.basename)
          # Values that should be derived from the name
          # (should be set by subclass)
          self.sample_name = None
          self.sample_number = None
          self.barcode_sequence = None
          self.lane_number = None
          self.read_number = None
          self.set_number = None
          self.is_index_read = False
      def __repr__(self):
          return self.basename
  
  class IlluminaFastqAttrs(BaseFastqAttrs):
      """Class for extracting information about Fastq files
  
      Given the name of a Fastq file, extract data about the sample name,
      barcode sequence, lane number, read number and set number.
  
      The name format can be a 'full' Fastq name as generated by CASAVA or
      bcl2fastq 1.8, which follows the general form:
  
      <sample_name>_<barcode_sequence>_L<lane_number>_R<read_number>_<set_number>.astq.gz
  
      e.g. for
  
      NA10831_ATCACG_L002_R1_001.fastq.gz
  
      sample_name = 'NA10831_ATCACG_L002_R1_001'
      barcode_sequence = 'ATCACG'
      lane_number = 2
      read_number = 1
      set_number = 1
  
      Alternatively it can be a full Fastq name as generated by bcl2fastq2,
      of the general form:
  
      <sample_name>_S<sample_number>_L<lane_number>_R<read_number>_001.fastq.gz
  
      e.g. for
  
      ES_exp1_S4_L003_R2_001.fastq.gz
  
      sample_name = 'ES_exp1'
      sample_number = 4
      lane_number = 3
      read_number = 2
      set_number = 1
  
      bcl2fastq can also produce 'index read' Fastq files where the
      R1/R2 is replaced by I1, e.g.:
  
      ES_exp1_S4_L003_I1_001.fastq.gz
  
      Alternatively it can be a 'reduced' version where one or more
      of the components has been omitted (typically because they are
      redundant in uniquely distinguishing a Fastq file within a
      set of Fastqs).
  
      The reduced formats are:
  
      <sample_name>
      <sample_name>_L<lane_number>
      <sample_name>_<barcode_sequence>
      <sample_name>_<barcode_sequence>_L<lane_number>
  
      with an optional suffix '_R<read_number>' for paired end sets.
  
      e.g.
  
      NA10831
      NA10831_L002
      NA10831_ATCACG
      NA10831_ATCACG_L002
  
      Finally, the name can be a non-standard name of the form:
  
      <sample_name>.r<read_number>
  
      or
  
      <sample_name>.<barcode_sequence>.r<read_number>
  
      In this case the sample_names are permitted to include dots.
  
      Provides the follow attributes:
  
      fastq:            the original fastq file name
      sample_name:      name of the sample (leading part of the name)
      sample_number:    integer (or None if no sample number)
      barcode_sequence: barcode sequence (string or None)
      lane_number:      integer (or None if no lane number)
      read_number:      integer (or None if no read number)
      set_number:       integer (or None if no set number)
      is_index_read:    boolean (True if index read, False if not)
      """
  
      def __init__(self,fastq):
          """Create and populate a new IlluminaFastqAttrs object
  
          Arguments:
            fastq: name of the fastq.gz (optionally can include leading path)
          """
          BaseFastqAttrs.__init__(self,fastq)
          # Base name for sample (no leading path or extension)
          fastq_base = self.basename
          # Determine if it's a non-standard (dot-separated) name
          #
          # These are names of the form e.g.
          # NH1.2.r2
          # or
          # NH1_ChIP-seq.ACAGTG.r2
          if fastq_base.count('.') > 0:
              self.delimiter = '.'
              fields = fastq_base.split('.')
              field = fields[-1]
              if len(field) == 2 and field.startswith('r'):
                  # Read number
                  self.read_number = int(field[1])
                  fields = fields[:-1]
                  field = fields[-1]
              if len(fields) > 1:
                  # Barcode sequence
                  # Determine separator for dual index barcodes ('-' or '+')
                  for sep in ('-','+'):
                      if sep in field:
                          break
                  is_tag = True
                  for f in field.split(sep):
                      for c in f:
                          is_tag = is_tag and c in 'ACGTN'
                  if is_tag:
                      self.barcode_sequence = field
                      fields = fields[:-1]
                      field = fields[-1]
              # Remaining fields are the sample name
              self.sample_name = '.'.join(fields)
              assert(self.sample_name != '')
              return
          # Some form of Illumina-derived name
          #
          # Full Illumina-style names are e.g.
          # NH1_ChIP-seq_Gli1_ACAGTG_L001_R1_001
          # or
          # NH1_ChIP-seq_Gli1_S4_L003_R2_001
          #
          # We have shorter name formats where redundant parts are
          # omitted, the patterns are:
          # NAME          e.g. NH1_ChIP-seq_Gli1
          # NAME+LANE     e.g. NH1_ChIP-seq_Gli1_L001
          # NAME+TAG      e.g. NH1_ChIP-seq_Gli1_ACAGTG
          # NAME+TAG+LANE e.g. NH1_ChIP-seq_Gli1_ACAGTG_L001
          #
          # Also read number (i.e. R1 or R2) is appended but only for
          # paired end samples
          #
          # The set number is never included, except for full names
          fields = fastq_base.split('_')
          self.delimiter = '_'
          # Deal with set number first e.g. 001
          field = fields[-1]
          ##logger.debug("Test for set number %s" % field)
          if len(field) == 3 and field.isdigit():
              self.set_number = int(field)
              fields = fields[:-1]
          # Deal with trailing read number e.g. R1
          field = fields[-1]
          ##logger.debug("Test for read number %s" % field)
          if len(field) == 2:
              if field.startswith('R'):
                  self.read_number = int(field[1])
                  fields = fields[:-1]
              elif field.startswith('I'):
                  self.read_number = int(field[1])
                  self.is_index_read = True
                  fields = fields[:-1]
          # Deal with trailing lane number e.g. L001
          field = fields[-1]
          ##logger.debug("Test for lane number %s" % field)
          if len(field) == 4 and field.startswith('L') and field[1:].isdigit():
              self.lane_number = int(field[1:])
              fields = fields[:-1]
          # Deal with trailing index tag e.g. ATTGCT or ATTGCT-CCTAAG
          field = fields[-1]
          ##logger.debug("Test for barcode sequence %s" % field)
          if len(fields) > 1:
              # This mustn't be the last field: if it is then it's
              # not the tag - it's the name
              # Determine separator for dual index barcodes ('-' or '+')
              for sep in ('-','+'):
                  if sep in field:
                      break
              is_tag = True
              for f in field.split(sep):
                  for c in f:
                      is_tag = is_tag and c in 'ACGTN'
              if is_tag:
                  self.barcode_sequence = field
                  fields = fields[:-1]
                  ##logger.debug("Identified barcode sequence as %s" % self.barcode_sequence)
              else:
                  # Alternatively might be the sample number
                  if field.startswith('S'):
                      try:
                          if field[1:].isdigit():
                              self.sample_number = int(field[1:])
                              fields = fields[:-1]
                      except IndexError:
                          pass
          # What's left is the name
          ##logger.debug("Remaining fields: %s" % fields)
          self.sample_name = '_'.join(fields)
          assert(self.sample_name != '')
  
      def __repr__(self):
          """Implement __repr__ built-in
          """
          # Reconstruct name
          fq = ["%s" % self.sample_name]
          if self.sample_number is not None:
              fq.append("S%d" % self.sample_number)
          if self.barcode_sequence is not None:
              fq.append("%s" % self.barcode_sequence)
          if self.lane_number is not None:
              fq.append("L%03d" % self.lane_number)
          if self.read_number is not None:
              if self.delimiter == '.':
                  fq.append("r%d" % self.read_number)
              elif self.is_index_read:
                  fq.append("I%d" % self.read_number)
              else:
                  fq.append("R%d" % self.read_number)
          if self.set_number is not None:
              fq.append("%03d" % self.set_number)
          return self.delimiter.join(fq)
  
  class FastqReadCounter(object):
      """
      Implements various methods for counting reads in FASTQ file
  
      The methods are:
  
      - simple: a wrapper for the FASTQFile.nreads() function
      - fastqiterator: counts reads using FASTQFile.FastqIterator
      - zcat_wc: runs 'zcat | wc -l' in the shell
      - reads_per_lane: counts reads by lane using FastqIterator
  
      """
      @staticmethod
      def simple(fastq=None,fp=None):
          """
          Return number of reads in a FASTQ file
  
          Uses the FASTQFile.nreads function to do the counting.
  
          Arguments:
            fastq: fastq(.gz) file
            fp: open file descriptor for fastq file
  
          Returns:
            Number of reads
  
          """
          return nreads(fastq=fastq,fp=fp)
      @staticmethod
      def fastqiterator(fastq=None,fp=None):
          """
          Return number of reads in a FASTQ file
  
          Uses the FASTQFile.FastqIterator class to do the
          counting.
  
          Arguments:
            fastq: fastq(.gz) file
            fp: open file descriptor for fastq file
  
          Returns:
            Number of reads
  
          """
          nreads = 0
          for r in FastqIterator(fastq_file=fastq,fp=fp):
              nreads += 1
          return nreads
      @staticmethod
      def zcat_wc(fastq=None,fp=None):
          """
          Return number of reads in a FASTQ file
  
          Uses a system call to run 'zcat FASTQ | wc -l' to do
          the counting (or just 'wc -l' if not a gzipped FASTQ).
  
          Note that this can only operate on fastq files (not
          on streams provided via the 'fp' argument; this will
          raise an exception).
  
          Arguments:
            fastq: fastq(.gz) file
            fp: open file descriptor for fastq file
  
          Returns:
            Number of reads
  
          """
          if fastq is None:
              raise Exception("zcat_wc: can only operate on a file")
          if fastq.endswith(".gz"):
              cmd = "zcat %s | wc -l" % fastq
          else:
              cmd = "wc -l %s | cut -d' ' -f1" % fastq
          output = subprocess.check_output(cmd,shell=True)
          try:
              return int(output)//4
          except Exception as ex:
              raise Exception("zcat_wc returned: %s" % output)
      @staticmethod
      def reads_per_lane(fastq=None,fp=None):
          """
          Return counts of reads in each lane of FASTQ file
  
          Uses the FASTQFile.FastqIterator class to do the
          counting, with counts split by lane.
  
          Arguments:
            fastq: fastq(.gz) file
            fp: open file descriptor for fastq file
  
          Returns:
            Dictionary where keys are lane numbers (as integers)
              and values are number of reads in that lane.
  
          """
          nreads = {}
          for r in FastqIterator(fastq_file=fastq,fp=fp):
              lane = int(r.seqid.flowcell_lane)
              try:
                  nreads[lane] += 1
              except KeyError:
                  nreads[lane] = 1
          return nreads
  
  #######################################################################
  # Functions
  #######################################################################
  
  def assign_barcodes_single_end(fastq_in,fastq_out,n=5):
      """
      Extract inline barcodes and assign to Fastq read headers
  
      Strips the first n bases from each read of the input
      FASTQ file and assigns it to the index sequence for that
      read in the output file.
  
      If the supplied output file name ends with '.gz' then it
      will be gzipped.
  
      Arguments:
        fastq_in (str): input FASTQ file (can be gzipped)
        fastq_out (str): output FASTQ file (will be gzipped if
          ending with '.gz')
        n (integer): number of bases to extract and assign as
          index sequence (default: 5)
  
      Returns:
        Integer: number of reads processed.
  
      """
      if fastq_out.endswith('.gz'):
          fp = gzip.GzipFile(filename=fastq_out,mode='wb')
      else:
          fp = open(fastq_out,'w')
      print("Processing reads from %s" % fastq_in)
      nread = 0
      for read in FastqIterator(fastq_in):
          # Extract new barcode sequence
          barcode = read.sequence[:n]
          # Truncate sequence and quality accordingly
          sequence = read.sequence[n:]
          quality = read.quality[n:]
          # Assign new values and write to output
          read.seqid.index_sequence = barcode
          read.sequence = sequence
          read.quality = quality
          fp.write("%s\n" % read)
          nread += 1
      fp.close()
      print("Finished (%d reads processed)" % nread)
      return nread
  
  def get_read_number(fastq):
      """
      Get the read number (1 or 2) from a Fastq file
  
      Arguments:
        fastq (str): path to a Fastq file
  
      Returns:
        Integer: read number (1 or 2) extracted from the first read.
      """
      with get_fastq_file_handle(fastq) as fp:
          for r in FastqIterator(fp=fp):
              seq_id = r.seqid
              break
      return int(seq_id.pair_id)
  
  def get_read_count(fastqs):
      """
      Get the total count of reads across multiple Fastqs
  
      Arguments:
        fastqs (list): lpaths to one or more Fastq files
  
      Returns:
        Integer: total number of reads across all files.
      """
      nreads = 0
      for fq in fastqs:
          n = FastqReadCounter.zcat_wc(fq)
          print("%s:\t%d" % (os.path.basename(fq),n))
          nreads += n
      return nreads
  
  def pair_fastqs(fastqs):
      """
      Automagically pair up FASTQ files
  
      Given a list of FASTQ files, generate a list of R1/R2
      pairs by examining the header for the first read in
      each file.
  
      Arguments:
        fastqs (list): list of paths to FASTQ files which
          will be paired.
  
      Returns:
        Tuple: pair of lists of the form (paired,unpaired),
          where `paired` is a list of tuples consisting of
          FASTQ R1/R2 pairs and `unpaired` is a list of
          FASTQs which couldn't be paired.
      """
      fq_pairs = []
      seq_ids = {}
      bad_files = []
      for fq in [os.path.abspath(fq) for fq in fastqs]:
          # Get header from first read
          seq_id = None
          with get_fastq_file_handle(fq) as fp:
              for r in FastqIterator(fp=fp):
                  seq_id = r.seqid
                  break
          if seq_id is None:
              logging.debug("'Bad' file: %s" % fq)
              bad_files.append(fq)
              continue
          fq_pair = None
          for fq1 in seq_ids:
              if seq_id.is_pair_of(seq_ids[fq1]):
                  # Found a pair
                  if seq_id.pair_id == '1':
                      fq_pair = (fq,fq1)
                  else:
                      fq_pair = (fq1,fq)
                  fq_pairs.append(fq_pair)
                  logging.debug("*** Paired: %s\n"
                                "          : %s" % fq_pair)
                  # Remove paired fastq
                  del(seq_ids[fq1])
                  break
          if fq_pair is None:
              # Unable to pair, store for now
              logging.debug("Unpaired: %s" % fq)
              seq_ids[fq] = seq_id
      # Sort pairs into order
      fq_pairs = sorted(fq_pairs,key=lambda x: x[0])
      unpaired = sorted(list(seq_ids.keys()) + bad_files)
      # Return paired and upaired fastqs
      return (fq_pairs,unpaired)
  
  def pair_fastqs_by_name(fastqs,fastq_attrs=IlluminaFastqAttrs):
      """
      Pair Fastq files based on their name
  
      Pairing is based on the read number for the supplied
      Fastq files being present in the file names; the file
      contents are not examined.
  
      Unpaired Fastqs (i.e. those for which a mate cannot be
      found) are returned as a "pair" where the equivalent R1
      or R2 mate is missing.
  
      Arguments:
        fastqs (list): list of Fastqs to pair
        fastq_attrs (BaseFastqAttrs): optional, class to use
          for extracting data from the filename (default:
          IlluminaFastqAttrs)
  
      Returns:
        List: list of tuples (R1,R2) with the R1/R2 pairs,
         or (R1,) or (R2,) for unpaired files.
      """
      pairs = []
      fastqs_r1 = sorted(list(filter(lambda f:
                                     fastq_attrs(f).read_number != 2,fastqs)))
      fastqs_r2 = sorted(list(filter(lambda f:
                                     fastq_attrs(f).read_number == 2,fastqs)))
      for fqr1 in fastqs_r1:
          # Split up R1 name
          logging.debug("fqr1 %s" % os.path.basename(fqr1))
          dir_path = os.path.dirname(fqr1)
          # Generate equivalent R2 file
          fqr2 = fastq_attrs(fqr1)
          fqr2.read_number = 2
          fqr2 = os.path.join(dir_path,"%s%s" % (fqr2,fqr2.extension))
          logging.debug("fqr2 %s" % os.path.basename(fqr2))
          if fqr2 in fastqs_r2:
              pairs.append((fqr1,fqr2))
          else:
              pairs.append((fqr1,))
      # Looking for unpaired R2 files
      for fqr2 in fastqs_r2:
          for pair in pairs:
              try:
                  if fqr2 == pair[1]:
                      fqr2 = None
                      break
              except IndexError:
                  pass
          if fqr2 is not None:
              pairs.append((fqr2,))
      pairs = sorted(pairs,key=lambda x: x[0])
      return pairs
  
  def group_fastqs_by_name(fastqs,fastq_attrs=IlluminaFastqAttrs):
      """
      Group Fastq files based on their name
  
      Grouping is based on the read number and type for the
      supplied Fastq files being present in the file names; the
      file contents are not examined.
  
      Unpaired Fastqs (i.e. those for which a mate cannot be
      found) are returned as a "pair" where the equivalent R1
      or R2 mate is missing.
  
      Arguments:
        fastqs (list): list of Fastqs to pair
        fastq_attrs (BaseFastqAttrs): optional, class to use
          for extracting data from the filename (default:
          IlluminaFastqAttrs)
  
      Returns:
        List: list of tuples (R1,R2) with the R1/R2 pairs,
          or (R1,) or (R2,) for unpaired files.
      """
      # Get reads and put into groups by read ID
      reads = set()
      index_reads = set()
      fastq_sets = dict()
      for fastq in fastqs:
          fq = fastq_attrs(fastq)
          read_number = fq.read_number if fq.read_number else 1
          if not fq.is_index_read:
              read = "r%d" % read_number
              reads.add(read)
          else:
              read = "i%d" % read_number
              index_reads.add(read)
          try:
              fastq_sets[read].append(fastq)
          except KeyError:
              fastq_sets[read] = [fastq]
      reads = sorted(list(reads)) + sorted(list(index_reads))
      # Rearrange into groups
      groups = []
      for ii,read in enumerate(reads):
          for fastq in fastq_sets[read]:
              # Create reference Fastq name
              fq_ref = fastq_attrs(fastq)
              fq_ref.is_index_read = False
              fq_ref.read_number = 1
              # Initialise a new group
              group = [fastq]
              # Look for matches
              for r in reads[ii+1:]:
                  unmatched_fastqs = list()
                  for fastq1 in fastq_sets[r]:
                      fq1_ref = fastq_attrs(fastq1)
                      fq1_ref.is_index_read = False
                      fq1_ref.read_number = 1
                      if str(fq1_ref) == str(fq_ref):
                          group.append(fastq1)
                      else:
                          unmatched_fastqs.append(fastq1)
                  fastq_sets[r] = unmatched_fastqs
              groups.append(group)
      groups = sorted(groups,key=lambda x: x[0])
      return groups
  
  def remove_index_fastqs(fastqs,fastq_attrs=IlluminaFastqAttrs):
      """
      Remove index (I1/I2) Fastqs from list
  
      Arguments:
        fastqs (list): list of paths to Fastq files
        fastq_attrs (BaseFastqAttrs): class to use for
          extracting attributes from Fastq names
          (defaults to IlluminaFastqAttrs)
  
      Returns:
        List: input Fastq list with any index read
          Fastqs removed.
      """
      return list(filter(lambda fq:
                         not fastq_attrs(fq).is_index_read,
                         fastqs))
  