content: """
  Read a rvcf file with stability selection scores for taxa.
  Sort the dataframe by rsq_median.
  Print results.
  
  usage:
      python sort_results.py \
          ../example/stability_selection_example_output.vcf \
          ../example/hominid_example_taxon_table_input.txt \
          arcsinsqrt \
          0.5 \
          10
  """
  import argparse
  import sys
  
  import pandas as pd
  
  from hominid.hominid import read_taxon_file, align_snp_and_taxa
  
  
  def sort_results(rvcf_input_file_path, taxon_table_file_path, transform,
                   r_sqr_median_cutoff, stability_cutoff, snp_count, no_tables,
                   extra_columns):
      print('plotting {} SNPs from {}'.format(snp_count, rvcf_input_file_path))
  
      # read the rvcf file and sort by rsq_median
      df = pd.read_csv(rvcf_input_file_path, sep='\t', dtype={'CHROM': str})
      #print('df.shape: {}'.format(df.shape))
  
      sorted_rsq_best_medians_df = df.sort_values(by='rsq_median', ascending=False)
  
      x_df = sorted_rsq_best_medians_df[sorted_rsq_best_medians_df.rsq_median > r_sqr_median_cutoff]
      print('{} SNPs with r_sqr > {:5.3f}'.format(x_df.shape[0], r_sqr_median_cutoff))
  
      taxon_table_df = read_taxon_file(taxon_table_file_path, transform=transform)
  
      for row_i in range(sorted_rsq_best_medians_df.shape[0]):
          if row_i >= snp_count:
              break
          else:
              # get a 1-row dataframe
              snp_df = sorted_rsq_best_medians_df.iloc[[row_i]]
              aligned_snp_df, aligned_taxa_df = align_snp_and_taxa(
                  snp_df,
                  taxon_table_df
              )
              # get the taxon stability selection scores
              # use the taxon table df index to get column names for snp_df
              taxon_scores_df = snp_df.loc[:, taxon_table_df.index].transpose()
              sorted_taxon_scores_df = taxon_scores_df.sort_values(by=taxon_scores_df.columns[0], ascending=False)
              #sorted_taxon_scores_df = taxon_scores_df.sort(taxon_scores_df.columns[0], ascending=False)
  
              p_df_list = []
              print('{} {} {:5.3f}'.format(snp_df.iloc[0].GENE, snp_df.iloc[0].ID, snp_df.iloc[0].rsq_median))
              summary_line = '{}\t{}\t'.format(snp_df.iloc[0].GENE, snp_df.iloc[0].ID)
              for i, (selected_taxon, selected_taxon_row) in enumerate(sorted_taxon_scores_df.iterrows()):
                  # use selected_taxon_row.index[0] to index the first and only column
                  selected_taxon_score = selected_taxon_row.iloc[0]
                  if selected_taxon_score < stability_cutoff:
                      #print('done with selected taxa')
                      break
                  else:
                      # trim 'Root;' from the front of the taxon name
                      if selected_taxon.startswith('Root;'):
                          taxon_name = selected_taxon[5:]
                      else:
                          taxon_name = selected_taxon
                      print('  {:5.3f} {}'.format(selected_taxon_score, taxon_name))
                      summary_line += '{}, '.format(taxon_name)
                      gts = [
                          snp_df.iloc[0].REF + snp_df.iloc[0].REF,  # 0
                          snp_df.iloc[0].REF + snp_df.iloc[0].ALT,  # 1
                          snp_df.iloc[0].ALT + snp_df.iloc[0].ALT   # 2
                      ]
                      aligned_snp_value_list = aligned_snp_df.values.flatten().tolist()
                      data_dict = {
                          'chromosome': [snp_df.iloc[0].CHROM] * aligned_snp_df.shape[1],
                          'snp_id': [snp_df.iloc[0].ID] * aligned_snp_df.shape[1],
                          'gene': [snp_df.iloc[0].GENE] * aligned_snp_df.shape[1],
                          'taxon': [selected_taxon] * aligned_snp_df.shape[1],
                          'abundance': aligned_taxa_df[selected_taxon].values.tolist(),
                          'variant_allele_count': [str(int(v)) for v in aligned_snp_value_list],
                          'genotype': [gts[int(v)] for v in aligned_snp_value_list],
                          'sample_id' : aligned_snp_df.columns
                      }
                      columns_to_display = ['abundance', 'variant_allele_count', 'genotype', 'sample_id']
                      if extra_columns:
                          for extra_column in extra_columns.split(','):
                              data_dict[extra_column] = snp_df.iloc[0][extra_column]
                              columns_to_display.append(extra_column)
                      p_df = pd.DataFrame(data_dict)
                      p_df_list.append(p_df)
                      if no_tables:
                          pass
                      else:
                          p_df[columns_to_display].to_csv(
                              sys.stdout,
                              sep='\t'
                          )
              # save a stacked bar plot
              if len(p_df_list) > 0:
                  file_name = 'stacked_bar_plot_selected_taxa_{}_{}.pdf'.format(
                      snp_df.iloc[0].GENE,
                      snp_df.iloc[0].ID
                  )
                  p_df = pd.concat(p_df_list, axis=0)
                  # at this point the index for p_df looks like
                  #   0...76.0...76.0...76
                  # replace the index
                  p_df.index = range(p_df.shape[0])
                  #p_df.to_csv(file_path, sep='\t')
                  stacked_bar_title = '{}\n{}'.format(snp_df.iloc[0].GENE, snp_df.iloc[0].ID)
  
  
  def main():
      argparser = argparse.ArgumentParser()
      argparser.add_argument('rvcf_input_file_path')
      argparser.add_argument('taxon_table_file_path')
      argparser.add_argument('transform')
      argparser.add_argument(
          'r_sqr_median_cutoff',
          type=float
      )
      argparser.add_argument(
          'stability_cutoff',
          type=float
      )
      argparser.add_argument(
          'snp_count',
          type=int
      )
      argparser.add_argument(
          '--no-tables',
          action='store_true'
      )
      argparser.add_argument(
          '--extra-columns',
          type=str
      )
      args = argparser.parse_args()
      print(args)
      sort_results(**vars(args))
  
  if __name__ == '__main__':
      main()
  