text: import collections
  import functools
  import glob
  import logging
  import os
  import pickle
  import shutil
  import sys
  import tempfile
  from typing import Dict, List, Tuple, Union
  
  import joblib
  import natsort
  import numba as nb
  import numpy as np
  import pandas as pd
  import scipy.sparse as ss
  from sklearn.utils import murmurhash3_32
  from spectrum_utils.spectrum import MsmsSpectrum
  
  from . import __version__
  from .cluster import cluster, spectrum
  from .config import config
  from .ms_io import ms_io
  
  
  logger = logging.getLogger('falcon')
  
  
  def main(args: Union[str, List[str]] = None) -> int:
      # Configure logging.
      logging.captureWarnings(True)
      root = logging.getLogger()
      root.setLevel(logging.DEBUG)
      handler = logging.StreamHandler(sys.stderr)
      handler.setLevel(logging.DEBUG)
      handler.setFormatter(logging.Formatter(
          '{asctime} {levelname} [{name}/{processName}] {module}.{funcName} : '
          '{message}', style='{'))
      root.addHandler(handler)
      # Disable dependency non-critical log messages.
      logging.getLogger('faiss').setLevel(logging.WARNING)
      logging.getLogger('numba').setLevel(logging.WARNING)
      logging.getLogger('numexpr').setLevel(logging.WARNING)
  
      # Load the configuration.
      config.parse(args)
      logger.info('falcon version %s', str(__version__))
      logger.debug('work_dir = %s', config.work_dir)
      logger.debug('overwrite = %s', config.overwrite)
      logger.debug('export_representatives = %s', config.export_representatives)
      logger.debug('usi_pxd = %s', config.usi_pxd)
      logger.debug('precursor_tol = %.2f %s', *config.precursor_tol)
      logger.debug('rt_tol = %s', config.rt_tol)
      logger.debug('fragment_tol = %.2f', config.fragment_tol)
      logger.debug('eps = %.3f', config.eps)
      logger.debug('min_samples = %d', config.min_samples)
      logger.debug('mz_interval = %d', config.mz_interval)
      logger.debug('hash_len = %d', config.hash_len)
      logger.debug('n_neighbors = %d', config.n_neighbors)
      logger.debug('n_neighbors_ann = %d', config.n_neighbors_ann)
      logger.debug('batch_size = %d', config.batch_size)
      logger.debug('n_probe = %d', config.n_probe)
      logger.debug('min_peaks = %d', config.min_peaks)
      logger.debug('min_mz_range = %.2f', config.min_mz_range)
      logger.debug('min_mz = %.2f', config.min_mz)
      logger.debug('max_mz = %.2f', config.max_mz)
      logger.debug('remove_precursor_tol = %.2f', config.remove_precursor_tol)
      logger.debug('min_intensity = %.2f', config.min_intensity)
      logger.debug('max_peaks_used = %d', config.max_peaks_used)
      logger.debug('scaling = %s', config.scaling)
  
      rm_work_dir = False
      if config.work_dir is None:
          config.work_dir = tempfile.mkdtemp()
          rm_work_dir = True
      elif os.path.isdir(config.work_dir):
          logging.warning('Working directory %s already exists, previous '
                          'results might get overwritten', config.work_dir)
      os.makedirs(config.work_dir, exist_ok=True)
      os.makedirs(os.path.join(config.work_dir, 'spectra'), exist_ok=True)
      os.makedirs(os.path.join(config.work_dir, 'nn'), exist_ok=True)
  
      # Clean all intermediate and final results if "overwrite" is specified,
      # otherwise abort if the output files already exist.
      exit_exists = False
      if os.path.isfile(f'{config.output_filename}.csv'):
          if config.overwrite:
              logger.warning('Output file %s (cluster assignments) already '
                             'exists and will be overwritten',
                             f'{config.output_filename}.csv')
              os.remove(f'{config.output_filename}.csv')
          else:
              logger.error('Output file %s (cluster assignments) already '
                           'exists, aborting...',
                           f'{config.output_filename}.csv')
              exit_exists = True
      if os.path.isfile(f'{config.output_filename}.mgf'):
          if config.overwrite:
              logger.warning('Output file %s (cluster representatives) already '
                             'exists and will be overwritten',
                             f'{config.output_filename}.mgf')
              os.remove(f'{config.output_filename}.mgf')
          else:
              logger.error('Output file %s (cluster representatives) already '
                           'exists, aborting...',
                           f'{config.output_filename}.mgf')
              exit_exists = True
      if exit_exists:
          logging.shutdown()
          return 1
      if config.overwrite:
          for filename in os.listdir(os.path.join(config.work_dir, 'spectra')):
              os.remove(os.path.join(config.work_dir, 'spectra', filename))
          for filename in os.listdir(os.path.join(config.work_dir, 'nn')):
              os.remove(os.path.join(config.work_dir, 'nn', filename))
  
      # Read the spectra from the input files and partition them based on their
      # precursor m/z.
      if not any([filename.endswith('.pkl') for filename in os.listdir(
              os.path.join(config.work_dir, 'spectra'))]):
          buckets = _prepare_spectra()
          joblib.dump(buckets, os.path.join(config.work_dir, 'spectra',
                                            'buckets.joblib'))
      else:
          buckets = joblib.load(os.path.join(config.work_dir, 'spectra',
                                             'buckets.joblib'))
  
      vec_len, min_mz, max_mz = spectrum.get_dim(config.min_mz, config.max_mz,
                                                 config.fragment_tol)
      process_spectrum = functools.partial(
          spectrum.process_spectrum,
          min_peaks=config.min_peaks,
          min_mz_range=config.min_mz_range,
          mz_min=min_mz,
          mz_max=max_mz,
          remove_precursor_tolerance=config.remove_precursor_tol,
          min_intensity=config.min_intensity,
          max_peaks_used=config.max_peaks_used,
          scaling=None if config.scaling == 'off' else config.scaling)
  
      # Pre-compute the index hash mappings.
      hash_lookup = np.asarray([murmurhash3_32(i, 0, True) % config.hash_len
                                for i in range(vec_len)], np.uint32)
      vectorize = functools.partial(
          spectrum.to_vector_parallel, dim=config.hash_len, min_mz=min_mz,
          max_mz=max_mz, bin_size=config.fragment_tol, hash_lookup=hash_lookup,
          norm=True)
  
      # Cluster the spectra per charge.
      clusters_all, current_label, representative_info = [], 0, []
      for charge, (n_spectra, bucket_filenames) in buckets.items():
          logger.info('Cluster %d spectra with precursor charge %d',
                      n_spectra, charge)
          dist_filename = os.path.join(
              config.work_dir, 'nn', f'dist_{charge}.npz')
          metadata_filename = os.path.join(
              config.work_dir, 'nn', f'metadata_{charge}.parquet')
          if (not os.path.isfile(dist_filename)
                  or not os.path.isfile(metadata_filename)):
              pairwise_dist_matrix, metadata = \
                  cluster.compute_pairwise_distances(
                      n_spectra, bucket_filenames, process_spectrum, vectorize,
                      config.precursor_tol[0], config.precursor_tol[1],
                      config.rt_tol, config.n_neighbors, config.n_neighbors_ann,
                      config.batch_size, config.n_probe)
              metadata.insert(1, 'precursor_charge', charge)
              logger.debug('Export pairwise distance matrix to file %s',
                           dist_filename)
              ss.save_npz(dist_filename, pairwise_dist_matrix, False)
              metadata.to_parquet(metadata_filename, index=False)
          else:
              logger.debug('Load previously computed pairwise distance matrix '
                           'from file %s', dist_filename)
              pairwise_dist_matrix = ss.load_npz(dist_filename)
              metadata = pd.read_parquet(metadata_filename)
          # No valid spectra found with the current charge.
          if len(metadata) == 0:
              continue
          # Cluster using the pairwise distance matrix.
          clusters = cluster.generate_clusters(
              pairwise_dist_matrix, config.eps, config.min_samples,
              metadata['precursor_mz'].values, metadata['retention_time'].values,
              config.precursor_tol[0], config.precursor_tol[1], config.rt_tol)
          # Make sure that different charges have non-overlapping cluster labels.
          mask_no_noise = clusters != -1
          clusters[mask_no_noise] += current_label
          current_label = (np.amax(clusters[mask_no_noise]) + 1
                           if any(mask_no_noise) else current_label)
          # Save cluster assignments.
          metadata['cluster'] = clusters
          clusters_all.append(metadata)
          # Extract identifiers for cluster representatives (medoids).
          if config.export_representatives:
              charge_repr = cluster.get_cluster_representatives(
                  clusters, pairwise_dist_matrix.indptr,
                  pairwise_dist_matrix.indices, pairwise_dist_matrix.data)
              if charge_repr is not None:
                  representative_info.append(metadata.iloc[charge_repr])
              if config.export_include_singletons:
                  representative_info.append(metadata.loc[~mask_no_noise])
              logger.debug('Extract %d cluster representative %sidentifiers',
                           len(charge_repr) if charge_repr is not None else 0,
                           f'and {(~mask_no_noise).sum()} singleton spectra '
                           if config.export_include_singletons else '')
  
      # Export cluster memberships and representative spectra.
      n_clusters, n_spectra_clustered = 0, 0
      for clust in clusters_all:
          clust_no_noise = clust[clust['cluster'] != -1]
          n_clusters += clust_no_noise['cluster'].nunique()
          n_spectra_clustered += len(clust_no_noise)
      if n_spectra_clustered == 0:
          logger.error('No valid spectra found for clustering')
          logging.shutdown()
          return 1
      logger.info('Export cluster assignments of %d spectra to %d unique '
                  'clusters to output file %s', n_spectra_clustered, n_clusters,
                  f'{config.output_filename}.csv')
      clusters_all = (pd.concat(clusters_all, ignore_index=True)
                      .sort_values('identifier', key=natsort.natsort_keygen()))
      with open(f'{config.output_filename}.csv', 'a') as f_out:
          # Metadata.
          f_out.write(f'# falcon version {__version__}\n')
          f_out.write(f'# work_dir = {config.work_dir}\n')
          f_out.write(f'# overwrite = {config.overwrite}\n')
          f_out.write(f'# export_representatives = '
                      f'{config.export_representatives}\n')
          f_out.write(f'# usi_pxd = {config.usi_pxd}\n')
          f_out.write(f'# precursor_tol = {config.precursor_tol[0]:.2f} '
                      f'{config.precursor_tol[1]}\n')
          f_out.write(f'# rt_tol = {config.rt_tol}\n')
          f_out.write(f'# fragment_tol = {config.fragment_tol:.2f}\n')
          f_out.write(f'# eps = {config.eps:.3f}\n')
          f_out.write(f'# min_samples = {config.min_samples}\n')
          f_out.write(f'# mz_interval = {config.mz_interval}\n')
          f_out.write(f'# hash_len = {config.hash_len}\n')
          f_out.write(f'# n_neighbors = {config.n_neighbors}\n')
          f_out.write(f'# n_neighbors_ann = {config.n_neighbors_ann}\n')
          f_out.write(f'# batch_size = {config.batch_size}\n')
          f_out.write(f'# n_probe = {config.n_probe}\n')
          f_out.write(f'# min_peaks = {config.min_peaks}\n')
          f_out.write(f'# min_mz_range = {config.min_mz_range:.2f}\n')
          f_out.write(f'# min_mz = {config.min_mz:.2f}\n')
          f_out.write(f'# max_mz = {config.max_mz:.2f}\n')
          f_out.write(f'# remove_precursor_tol = '
                      f'{config.remove_precursor_tol:.2f}\n')
          f_out.write(f'# min_intensity = {config.min_intensity:.2f}\n')
          f_out.write(f'# max_peaks_used = {config.max_peaks_used}\n')
          f_out.write(f'# scaling = {config.scaling}\n')
          f_out.write('#\n')
          # Cluster assignments.
          clusters_all.to_csv(f_out, index=False, chunksize=1000000)
      if config.export_representatives:
          representative_info = pd.concat(representative_info, ignore_index=True)
          logger.info('Export %d cluster representative spectra %sto output '
                      'file %s', len(representative_info),
                      ('(including singletons) '
                       if config.export_include_singletons else ''),
                      f'{config.output_filename}.mgf')
          # Get the spectra corresponding to the cluster representatives.
          representative_info['filename'] = representative_info.apply(
              lambda row: os.path.join(
                  config.work_dir, 'spectra',
                  f"""{row.precursor_charge}_{_precursor_to_interval(
                      row.precursor_mz, row.precursor_charge,
                      config.mz_interval)}.pkl"""),
              axis='columns')
          representatives = []
          for spectra in joblib.Parallel(n_jobs=-1)(
                  joblib.delayed(_find_spectra_pkl)(
                      fn, spectra.set_index('identifier')['cluster'].to_dict())
                  for fn, spectra in representative_info.groupby('filename')):
              representatives.extend(spectra)
          representatives.sort(key=lambda spec: spec.cluster)
          ms_io.write_spectra(f'{config.output_filename}.mgf', representatives)
  
      if rm_work_dir:
          shutil.rmtree(config.work_dir)
  
      logging.shutdown()
      return 0
  
  
  def _prepare_spectra() -> Dict[int, Tuple[int, List[str]]]:
      """
      Read the spectra from the input peak files and partition to intermediate
      files split and sorted by precursor m/z.
  
      Returns
      -------
      Dict[int, Tuple[int, List[str]]]
          A dictionary with per precursor charge the total number of spectra and
          a list of file names of the spectrum buckets.
      """
      input_filenames = [fn for pattern in config.input_filenames
                         for fn in glob.glob(pattern)]
      logger.info('Read spectra from %d peak file(s)', len(input_filenames))
      spectra = collections.defaultdict(lambda: collections.defaultdict(list))
      for file_spectra in joblib.Parallel(n_jobs=-1)(
              joblib.delayed(_read_spectra)(filename, config.usi_pxd,
                                            config.mz_interval, config.work_dir)
              for filename in input_filenames):
          for spec, pkl_filename in file_spectra:
              spectra[spec.precursor_charge][pkl_filename].append(spec)
      # Make sure the spectra in the individual files are sorted by their
      # precursor m/z and count the number of spectra per precursor charge.
      buckets, n_spectra, n_buckets = {}, 0, 0
      for charge in sorted(spectra.keys()):
          spectra_charge = spectra[charge]
          n_spectra_charge = sum(
              joblib.Parallel(n_jobs=-1)
              (joblib.delayed(_write_spectra_pkl)(filename, file_spectra)
               for filename, file_spectra in spectra_charge.items()))
          buckets[charge] = (n_spectra_charge,
                             natsort.natsorted(spectra_charge.keys()))
          n_spectra += n_spectra_charge
          n_buckets += len(spectra_charge.keys())
      logger.debug('%d spectra written to %d buckets by precursor charge and '
                   'precursor m/z', n_spectra, n_buckets)
      return buckets
  
  
  def _read_spectra(filename: str, usi_pxd: str, mz_interval: int,
                    work_dir: str) -> List[Tuple[MsmsSpectrum, str]]:
      """
      Get MS/MS spectra from the given file.
  
      Parameters
      ----------
      filename : str
          The path of the peak file to be read.
      usi_pxd : str
          ProteomeXchange dataset identifier to create USI references.
      mz_interval : int
          The width of each m/z interval.
      work_dir : str
          The directory in which the spectrum buckets will be stored.
  
  
      Returns
      -------
      List[Tuple[MsmsSpectrum, str]]
          The spectra read from the given file and their bucket filenames
          (based on precursor charge and m/z).
      """
      spectra = []
      for spec in ms_io.get_spectra(filename):
          spec.identifier = f'mzspec:{usi_pxd}:{spec.identifier}'
          interval = _precursor_to_interval(
              spec.precursor_mz, spec.precursor_charge, mz_interval)
          pkl_filename = os.path.join(work_dir, 'spectra',
                                      f'{spec.precursor_charge}_{interval}.pkl')
          spectra.append((spec, pkl_filename))
      return spectra
  
  
  @nb.njit
  def _precursor_to_interval(mz: float, charge: int, interval_width: int) -> int:
      """
      Convert the precursor m/z to the neutral mass and get the interval index.
  
      Parameters
      ----------
      mz : float
          The precursor m/z.
      charge : int
          The precursor charge.
      interval_width : int
          The width of each m/z interval.
  
      Returns
      -------
      int
          The index of the interval to which a spectrum with the given m/z and
          charge belongs.
      """
      hydrogen_mass, cluster_width = 1.00794, 1.0005079
      neutral_mass = (mz - hydrogen_mass) * charge
      return round(neutral_mass / cluster_width) // interval_width
  
  
  def _write_spectra_pkl(filename: str, spectra: List[MsmsSpectrum]) -> int:
      """
      Write the spectra to a pickled file sorted by their precursor m/z.
  
      Parameters
      ----------
      filename : str
          The pickled spectrum file name.
      spectra : List[MsmsSpectrum]
          The spectra to write to the file.
  
      Returns
      -------
      int
          The number of spectra in the file.
      """
      # Radix sort on integers with O(n) time complexity.
      order = np.argsort(np.asarray([int(spec.precursor_mz * 10000)
                                     for spec in spectra], dtype=np.int32),
                         kind='stable')
      with open(filename, 'wb') as f_in:
          pickle.dump([spectra[i] for i in order], f_in, protocol=5)
      return len(spectra)
  
  
  def _find_spectra_pkl(filename: str, usis_to_read: Dict[str, int]) \
          -> List[MsmsSpectrum]:
      """
      Read spectra with the specified USIs from a pkl file.
  
      Parameters
      ----------
      filename : str
          Name of the pkl file from which the spectra are read.
      usis_to_read : Dict[str, int]
          Dict with as keys the USIs of the spectra to read from the pkl file,
          and as values the cluster labels that will be assigned to the spectra.
  
      Returns
      -------
      Iterable[MsmsSpectrum]
          The spectra with the specified USIs.
      """
      spectra_found = []
      with open(filename, 'rb') as f_in:
          for spec in pickle.load(f_in):
              if spec.identifier in usis_to_read:
                  spec.cluster = usis_to_read[spec.identifier]
                  spectra_found.append(spec)
                  if len(spectra_found) == len(usis_to_read):
                      return spectra_found
      raise ValueError(f'{len(usis_to_read) - len(spectra_found)} spectra not '
                       f'found in file {filename}')
  
  
  if __name__ == '__main__':
      sys.exit(main())
  