text: # ----------------------------------------------------------------------------
  # Copyright (c) 2021, QIIME 2 development team.
  #
  # Distributed under the terms of the Modified BSD License.
  #
  # The full license is in the file LICENSE, distributed with this software.
  # ----------------------------------------------------------------------------
  from multiprocessing import cpu_count
  
  import pandas as pd
  import numpy as np
  import scipy
  import warnings
  import qiime2 as q2
  import q2templates
  from os.path import join
  
  from joblib import Parallel, delayed
  from q2_types.feature_data import DNAIterator
  from sklearn.feature_extraction.text import HashingVectorizer
  import seaborn as sns
  import matplotlib.pyplot as plt
  from collections import Counter
  from itertools import zip_longest
  import pkg_resources
  
  from ._utilities import _taxon_to_list
  
  TEMPLATES = pkg_resources.resource_filename('rescript', 'assets')
  
  
  def evaluate_taxonomy(ctx,
                        taxonomies,
                        labels=None,
                        rank_handle_regex=None):
      labels = _process_labels(labels, taxonomies)
      summaries = []
      for name, taxonomy in zip(labels, taxonomies):
          summary = _evaluate_taxonomy(
              taxonomy.view(pd.Series), rank_handle_regex)
          # note: what if n equals an existing name? do we want to error out or
          # just assume a user knows what they are doing? For now I am just
          # adding this info to the warning to make this behavior transparent.
          summary['Dataset'] = str(name)
          summaries.append(summary)
      results = pd.concat(summaries).reset_index()
      # convert index to strings
      results.index = pd.Index(
          [str(i) for i in range(1, len(results.index) + 1)], name='id')
      results = q2.Metadata(results)
      volatility = ctx.get_action('longitudinal', 'volatility')
      plots, = volatility(metadata=results,
                          state_column='Level',
                          default_group_column='Dataset',
                          default_metric='Taxonomic Entropy')
      return plots
  
  
  def _process_labels(labels, taxonomies, description='taxonomies'):
      if labels is None:
          labels = []
      n_labels, n_taxonomies = len(labels), len(taxonomies)
      # warn if fewer labels than taxonomies
      if n_labels < n_taxonomies:
          _warn_uneven_length(description)
          labels = [name if name is not None else n for n, (name, t)
                    in enumerate(zip_longest(labels, taxonomies), 1)]
      # trim labels to number of Taxonomies
      elif n_labels > n_taxonomies:
          labels = labels[:n_taxonomies]
      return labels
  
  
  def _warn_uneven_length(description):
      msg = ("The lists of input {0} and labels are different lengths. "
             "Additional {0} will be labeled numerically by their order in the "
             "inputs. Note that if these numbers match existing labels, those "
             "data will be grouped in the visualization.".format(description))
      warnings.warn(msg, UserWarning)
  
  
  def _evaluate_taxonomy(taxonomy, rank_handle_regex):
      if rank_handle_regex is None:
          rank_handle_regex = ""
      # Count number of unique taxa and unclassifieds at each level
      summary = summarize_taxonomic_depth(
          taxonomy, rank_handle_regex=rank_handle_regex)
      # Measure taxonomic entropy at each level
      max_depth = summary.index.max()
      entropy = _taxonomic_entropy(
          taxonomy, rank_handle_regex=rank_handle_regex, max_depth=max_depth)
      entropy = entropy.merge(summary, left_index=True, right_index=True)
      entropy.index.name = 'Level'
      return entropy
  
  
  def summarize_taxonomic_depth(taxonomy, rank_handle_regex):
      # measure number of levels in each taxonomy
      depths = _taxonomic_depth(taxonomy, rank_handle_regex=rank_handle_regex)
      # count number + proportion of taxonomies per depth
      depths = depths.value_counts()
      total = remaining = len(taxonomy)
      proportions = depths / total
      depths = pd.concat([depths, proportions], axis=1)
      depths.columns = ['Number of Features Terminating at Depth',
                        'Proportion of Features Terminating at Depth']
      # iterate over depths up to max depth
      # fill in empty depths (i.e., indicate no termination)
      # calculate cumulative sum of (un)classifieds per depth
      depths_idx = depths.index
      classified = []
      unclassified = []
      for d in range(0, depths_idx.max() + 1):
          if d not in depths_idx:
              depths.loc[d] = [0, 0]
          classified.append(remaining)
          unclassified.append(total - remaining)
          remaining -= depths.loc[d, 'Number of Features Terminating at Depth']
      depths = depths.sort_index()
      classified_prop = [c / total for c in classified]
      unclassified_prop = [u / total for u in unclassified]
      depths['Number of Features Classified at Depth'] = classified
      depths['Proportion of Features Classified at Depth'] = classified_prop
      depths['Number of Features Unclassified at Depth'] = unclassified
      depths['Proportion of Features Unclassified at Depth'] = unclassified_prop
      depths.index.name = 'Level'
      # drop depth 0 on return (indicating presence of empty taxonomies)
      # these unclassifications are counted, but we don't want to plot at
      # level 0 since it is confusing — level 1 is the real root.
      return depths.drop(0)
  
  
  def _taxonomic_entropy(taxonomy, rank_handle_regex, max_depth):
      # convert each taxonomy to a list
      taxa_lists = [[t for t in _taxon_to_list(v, rank_handle=rank_handle_regex)
                     if t not in [None, '']]
                    for v in taxonomy.values]
      # taxonomic labels should accumulate at each rank (e.g., to avoid
      # counting identical species names as duplicates when genus is unique)
      taxa_lists = [[';'.join(t[:i]) for i in range(1, max_depth + 1)]
                    for t in taxa_lists]
      # coalesce taxonomies at each rank
      ranks = [i for i in zip_longest(*taxa_lists)]
      # Count unique values at each rank, excluding 'None'
      # use Counter instead of set to pass counts to entropy
      unique_counts = [[v for k, v in Counter(r).items()]
                       for r in ranks]
      # measure entropy at each rank, excluding 'None'
      entropy = {n: [len(r), scipy.stats.entropy(r)]
                 for n, r in enumerate(unique_counts, 1)}
      return pd.DataFrame(entropy,
                          index=['Unique Labels', 'Taxonomic Entropy']).T
  
  
  def _taxonomic_depth(taxonomy, rank_handle_regex):
      return taxonomy.apply(lambda x: len([i for i in _taxon_to_list(
          x, rank_handle=rank_handle_regex) if i not in [None, '']]))
  
  
  def evaluate_seqs(output_dir: str, sequences: DNAIterator, labels: list = None,
                    kmer_lengths: list = None, subsample_kmers: float = 1.0,
                    palette: str = 'viridis') -> None:
      labels = _process_labels(labels, sequences, description='sequences')
      results, lengths = _evaluate_seqs(
          sequences, labels, kmer_lengths, subsample_kmers)
      fig = _plot_eval_seqs(results, lengths, palette)
      _visualize(output_dir, results, fig)
      plt.close('all')
  
  
  def _process_kmers(seqs, kmer_len, subsample_kmers):
      if subsample_kmers < 1:
          subsample_size = int(subsample_kmers * len(seqs))
          seqs = np.random.choice(seqs, subsample_size, replace=False)
      vectorizer = HashingVectorizer(
          alternate_sign=False,
          analyzer='char',
          ngram_range=[kmer_len, kmer_len])
      X = vectorizer.fit_transform(seqs)
      kmer_freq = X.sum(axis=0)
      return scipy.stats.entropy(kmer_freq, axis=1)[0]
  
  
  def _evaluate_seqs(sequences, labels, kmer_lengths=None, subsample_kmers=1.0):
      if not kmer_lengths:
          kmer_lengths = []
  
      rownames = ['Length ' + n for n in ['min', '1%', '25%', 'median', '75%',
                                          '99%', 'max']]
      rownames += ['N uniques', 'Sequence Entropy']
      rownames += ['%smer Entropy' % k for k in kmer_lengths]
  
      lengths = []
      results = pd.DataFrame(index=rownames)
      for n, seqs in zip(labels, sequences):
          seqs = [str(s) for s in seqs]
          length_array = np.array([len(s) for s in seqs])
          lengths.append(length_array)
          len_quantiles = np.nanquantile(
              length_array, [0, 0.01, 0.25, 0.5, 0.75, 0.99, 1])
          uniqs = Counter(seqs).values()
          seq_entropy = scipy.stats.entropy(list(uniqs))
          res = list(len_quantiles) + [len(uniqs), seq_entropy]
  
          n_jobs = min(len(kmer_lengths), cpu_count())
          if n_jobs > 0:
              parallel = Parallel(n_jobs=n_jobs, backend='loky')
              kmer_freqs = parallel(
                  delayed(_process_kmers)(
                      seqs, k, subsample_kmers) for k in kmer_lengths)
              res.extend(kmer_freqs)
          results[n] = res
      return results.round(2), lengths
  
  
  def _plot_eval_seqs(results, lengths, palette):
      n_groups = len(lengths)
      cmap = sns.color_palette(palette, n_groups)
  
      # adjust figure width based on number of inputs — minimum 7 inches
      fig_width = 3 * n_groups + 3
      fig = plt.figure(constrained_layout=True, figsize=(fig_width, 6))
      gs1 = fig.add_gridspec(
          nrows=2, ncols=4, left=0.05, right=0.48, wspace=0.05)
  
      # length distribution histogram
      ax = fig.add_subplot(gs1[:-1, :])
      for dat, color in zip(lengths, cmap):
          sns.kdeplot(dat, shade=True, color=color, ax=ax, alpha=0.5)
      ax.set_title('Sequence Length Distribution')
      ax.set_ylabel('Proportion')
      ax.set_xlabel('Length (nt)')
      ax.set_xlim(results.loc['Length 1%'].min(),
                  results.loc['Length 99%'].max())
  
      # barplot of seq counts
      ax = fig.add_subplot(gs1[-1, :-2])
      sns.barplot(data=results.loc[['N uniques']], palette=palette, ax=ax)
      ax.set_title('N Unique Sequences')
      ax.set_ylabel('Count')
      ax.set_xticklabels(ax.get_xticklabels(), rotation=90, ha='right')
  
      # lineplot/barplot of seq/kmer entropy
      ax = fig.add_subplot(gs1[-1, -2:])
      if len(results) > 9:
          data = results[8:]
          data.index = [i.split(' ')[0] for i in data.index]
          sns.lineplot(data=data, sort=False, palette=palette, ax=ax)
          plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
      else:
          # plot barplot if kmer entropy is not calculated
          sns.barplot(data=results.loc[['Sequence Entropy']],
                      palette=palette, ax=ax)
          ax.set_xticklabels(ax.get_xticklabels(), rotation=90, ha='right')
      ax.set_title('Entropy')
      ax.set_ylabel('H')
  
      return fig
  
  
  def _visualize(output_dir, results, plot):
  
      pd.set_option('display.max_colwidth', -1)
  
      # save results
      results.to_csv(join(output_dir, 'evaluate_seqs_results.tsv'), sep='\t')
      results = q2templates.df_to_html(results, index=True)
  
      plot.savefig(join(output_dir, 'evaluate_seqs.png'), bbox_inches='tight')
      plot.savefig(join(output_dir, 'evaluate_seqs.pdf'), bbox_inches='tight')
  
      index = join(TEMPLATES, 'index.html')
      q2templates.render(index, output_dir, context={
          'title': 'Sequence Evaluation Results',
          'running_title': 'evaluate_seqs',
          'results': results,
      })
  