<commit_before>import os, json, datetime, sys, re
import rethinkdb as r
from Bio import SeqIO
import numpy as np

sys.path.append('')  # need to import from base
from base.rethink_io import rethink_io


def get_parser():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-db', '--database', default='vdb', help="database to download from")
    parser.add_argument('--rethink_host', help="rethink host url")
    parser.add_argument('--auth_key', help="auth_key for rethink database")
    parser.add_argument('--local', action="store_true",  help ="connect to local instance of rethinkdb database")
    parser.add_argument('-v', '--virus', help="virus name")
    parser.add_argument('--ftype', default='fasta', help="output file format, default \"fasta\", other options are \"json\" and \"tsv\"")
    parser.add_argument('--fstem', help="default output file name is \"VirusName_Year_Month_Date\"")
    parser.add_argument('--path', default='data', help="path to dump output files to")
    parser.add_argument('--fasta_fields', nargs='+', default=['strain', 'virus', 'accession', 'collection_date', 'region', 'country', 'division', 'location', 'source', 'locus', 'authors'], help="fasta fields for output fasta")

    parser.add_argument('--public_only', action="store_true", help="include to subset public sequences")
    parser.add_argument('--select', nargs='+', type=str, default=[], help="Select specific fields ie \'--select field1:value1 field2:value1,value2\'")
    parser.add_argument('--present', nargs='+', type=str, default=[], help="Select specific fields to be non-null ie \'--present field1 field2\'")
    parser.add_argument('--interval', nargs='+', type=str, default=[], help="Select interval of values for fields \'--interval field1:value1,value2 field2:value1,value2\'")
    parser.add_argument('--years_back', type=str, help='number of past years to sample sequences from \'--years_back field:value\'')
    parser.add_argument('--relaxed_interval', action="store_true", help="Relaxed comparison to date interval, 2016-XX-XX in 2016-01-01 - 2016-03-01")

    def duplicate_resolver(resolve_method):
        method = str(resolve_method)
        accepted_methods = ('keep_duplicates', 'choose_longest', 'choose_genbank', 'split_passage')
        if method not in accepted_methods:
            msg = '"{}" is not a supported duplicate resolve method'.format(method)
            raise argparse.ArgumentTypeError(msg)
        return method

    parser.add_argument('--resolve_method', type=duplicate_resolver, default="choose_longest", help="Set method of resolving duplicates for the same locus, options are \'keep_duplicates\', \'choose_longest\', \'choose_genbank\' and \'split_passage\'")
    return parser


class download(object):
    def __init__(self, database, virus, **kwargs):
        '''
        parser for virus, fasta fields, output file names, output file format path, interval
        '''
        self.virus = virus.lower()
        self.viruses_table = virus + "_viruses"
        self.sequences_table = virus + "_sequences"
        self.database = database.lower()

    def connect_rethink(self, **kwargs):
        if self.database not in ['vdb', 'test_vdb', 'test']:
            raise Exception("Can't download from this database: " + self.database)
        self.rethink_io = rethink_io()
        self.rethink_host, self.auth_key = self.rethink_io.assign_rethink(**kwargs)
        self.rethink_io.connect_rethink(self.database, self.rethink_host, self.auth_key)
        self.rethink_io.check_table_exists(self.database, self.viruses_table)
        self.rethink_io.check_table_exists(self.database, self.sequences_table)

    def count_documents(self, table):
        '''
        return integer count of number of documents in table
        '''
        return r.db(self.database).table(table).count().run()

    def download(self, output=True, **kwargs):
        '''
        download documents from table
        '''

        import time
        start_time = time.time()
        self.connect_rethink(**kwargs)
        select, present, interval= self.parse_subset_arguments(**kwargs)
        sequence_count = r.table(self.sequences_table).count().run()
        # print(sequence_count, "sequences in table:", self.sequences_table)
        virus_count = r.table(self.viruses_table).count().run()
        # print(virus_count, "viruses in table:", self.viruses_table)
        print("Downloading documents from the sequence table \"{}\" (n={}) & virus table \"{}\" (n={})".format(self.sequences_table, sequence_count, self.viruses_table, virus_count))
        sequences = self.rethinkdb_download(self.sequences_table, self.viruses_table, presents=present, selections=select, intervals=interval, **kwargs)
        print("Downloaded " + str(len(sequences)) + " sequences")
        sequences = self.resolve_duplicates(sequences, **kwargs)
        if output:
            self.output(sequences, **kwargs)
        print("--- %s minutes to download ---" % ((time.time() - start_time)/60))

    def parse_subset_arguments(self, select=[], present=[], interval=[], years_back=None, **kwargs):
        '''
        Parse arguments needed for subsetting on the server side
        Return a tuple of the arguments
        '''
        selections = self.parse_select_argument(select)
        intervals = self.parse_select_argument(interval)
        if years_back is not None and len(intervals)==0:
            intervals = self.parse_years_back_argument(years_back)
        return selections, present, intervals,

    def parse_select_argument(self, groupings=[]):
        '''
        :arg groupings like country:brazil,argentina
        parse the 'select' parameter to determine which field name to filter and for what values
        :return: [(grouping name, [group values])] ex. [(country, [brazil, argentina)]
        '''
        selections = []
        if groupings is not None and len(groupings)>0:
            for group in groupings:
                result = group.split(':')
                selections.append((result[0].lower(), result[1].lower().split(',')))
        return selections

    def parse_years_back_argument(self, argument):
        field = argument.split(':')[0].lower()
        years = int(argument.split(':')[1])
        date_now = datetime.datetime.utcnow()
        older_date = str(datetime.datetime.strftime(date_now.replace(year=date_now.year-years),'%Y-%m-%d'))
        newer_date = str(datetime.datetime.strftime(date_now,'%Y-%m-%d'))
        return [(field, [older_date, newer_date])]

    def rethinkdb_download(self, sequence_table, virus_table, index='strain', **kwargs):
        '''
        Default command merges documents from the sequence table and virus table
        Chain rethinkdb filter and has_fields commands to the default command
        Return documents from the database that are left after filtering
        '''
        # take each sequence and merge with corresponding virus document
        # NO! you can't define the spec as a lookup from viruses[accessions] to the primary key of sequences and then do the merge this way! don't do this
        command = r.table(sequence_table).merge(lambda sequence: r.table(virus_table).get(sequence[index]))
        command = self.add_present_command(command, **kwargs)
        command = self.add_selections_command(command, **kwargs)
        command = self.add_intervals_command(command, **kwargs)
        command = self.add_public_command(command, **kwargs)
        sequences = list(command.run())
        sequences = filter(None, sequences)
        return list(sequences)

    def add_present_command(self, command, presents=[], **kwargs):
        '''
        Add present fields check to command
        '''
        if len(presents)>0:
            print("Only downloading documents with fields: " + str(presents))<commit_msg>Add commented code for eq_join merge strategy<commit_after>import os, json, datetime, sys, re
import rethinkdb as r
from Bio import SeqIO
import numpy as np

sys.path.append('')  # need to import from base
from base.rethink_io import rethink_io


def get_parser():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-db', '--database', default='vdb', help="database to download from")
    parser.add_argument('--rethink_host', help="rethink host url")
    parser.add_argument('--auth_key', help="auth_key for rethink database")
    parser.add_argument('--local', action="store_true",  help ="connect to local instance of rethinkdb database")
    parser.add_argument('-v', '--virus', help="virus name")
    parser.add_argument('--ftype', default='fasta', help="output file format, default \"fasta\", other options are \"json\" and \"tsv\"")
    parser.add_argument('--fstem', help="default output file name is \"VirusName_Year_Month_Date\"")
    parser.add_argument('--path', default='data', help="path to dump output files to")
    parser.add_argument('--fasta_fields', nargs='+', default=['strain', 'virus', 'accession', 'collection_date', 'region', 'country', 'division', 'location', 'source', 'locus', 'authors'], help="fasta fields for output fasta")

    parser.add_argument('--public_only', action="store_true", help="include to subset public sequences")
    parser.add_argument('--select', nargs='+', type=str, default=[], help="Select specific fields ie \'--select field1:value1 field2:value1,value2\'")
    parser.add_argument('--present', nargs='+', type=str, default=[], help="Select specific fields to be non-null ie \'--present field1 field2\'")
    parser.add_argument('--interval', nargs='+', type=str, default=[], help="Select interval of values for fields \'--interval field1:value1,value2 field2:value1,value2\'")
    parser.add_argument('--years_back', type=str, help='number of past years to sample sequences from \'--years_back field:value\'')
    parser.add_argument('--relaxed_interval', action="store_true", help="Relaxed comparison to date interval, 2016-XX-XX in 2016-01-01 - 2016-03-01")

    def duplicate_resolver(resolve_method):
        method = str(resolve_method)
        accepted_methods = ('keep_duplicates', 'choose_longest', 'choose_genbank', 'split_passage')
        if method not in accepted_methods:
            msg = '"{}" is not a supported duplicate resolve method'.format(method)
            raise argparse.ArgumentTypeError(msg)
        return method

    parser.add_argument('--resolve_method', type=duplicate_resolver, default="choose_longest", help="Set method of resolving duplicates for the same locus, options are \'keep_duplicates\', \'choose_longest\', \'choose_genbank\' and \'split_passage\'")
    return parser


class download(object):
    def __init__(self, database, virus, **kwargs):
        '''
        parser for virus, fasta fields, output file names, output file format path, interval
        '''
        self.virus = virus.lower()
        self.viruses_table = virus + "_viruses"
        self.sequences_table = virus + "_sequences"
        self.database = database.lower()

    def connect_rethink(self, **kwargs):
        if self.database not in ['vdb', 'test_vdb', 'test']:
            raise Exception("Can't download from this database: " + self.database)
        self.rethink_io = rethink_io()
        self.rethink_host, self.auth_key = self.rethink_io.assign_rethink(**kwargs)
        self.rethink_io.connect_rethink(self.database, self.rethink_host, self.auth_key)
        self.rethink_io.check_table_exists(self.database, self.viruses_table)
        self.rethink_io.check_table_exists(self.database, self.sequences_table)

    def count_documents(self, table):
        '''
        return integer count of number of documents in table
        '''
        return r.db(self.database).table(table).count().run()

    def download(self, output=True, **kwargs):
        '''
        download documents from table
        '''

        import time
        start_time = time.time()
        self.connect_rethink(**kwargs)
        select, present, interval= self.parse_subset_arguments(**kwargs)
        sequence_count = r.table(self.sequences_table).count().run()
        # print(sequence_count, "sequences in table:", self.sequences_table)
        virus_count = r.table(self.viruses_table).count().run()
        # print(virus_count, "viruses in table:", self.viruses_table)
        print("Downloading documents from the sequence table \"{}\" (n={}) & virus table \"{}\" (n={})".format(self.sequences_table, sequence_count, self.viruses_table, virus_count))
        sequences = self.rethinkdb_download(self.sequences_table, self.viruses_table, presents=present, selections=select, intervals=interval, **kwargs)
        print("Downloaded " + str(len(sequences)) + " sequences")
        sequences = self.resolve_duplicates(sequences, **kwargs)
        if output:
            self.output(sequences, **kwargs)
        print("--- %s minutes to download ---" % ((time.time() - start_time)/60))

    def parse_subset_arguments(self, select=[], present=[], interval=[], years_back=None, **kwargs):
        '''
        Parse arguments needed for subsetting on the server side
        Return a tuple of the arguments
        '''
        selections = self.parse_select_argument(select)
        intervals = self.parse_select_argument(interval)
        if years_back is not None and len(intervals)==0:
            intervals = self.parse_years_back_argument(years_back)
        return selections, present, intervals,

    def parse_select_argument(self, groupings=[]):
        '''
        :arg groupings like country:brazil,argentina
        parse the 'select' parameter to determine which field name to filter and for what values
        :return: [(grouping name, [group values])] ex. [(country, [brazil, argentina)]
        '''
        selections = []
        if groupings is not None and len(groupings)>0:
            for group in groupings:
                result = group.split(':')
                selections.append((result[0].lower(), result[1].lower().split(',')))
        return selections

    def parse_years_back_argument(self, argument):
        field = argument.split(':')[0].lower()
        years = int(argument.split(':')[1])
        date_now = datetime.datetime.utcnow()
        older_date = str(datetime.datetime.strftime(date_now.replace(year=date_now.year-years),'%Y-%m-%d'))
        newer_date = str(datetime.datetime.strftime(date_now,'%Y-%m-%d'))
        return [(field, [older_date, newer_date])]

    def rethinkdb_download(self, sequence_table, virus_table, index='strain', **kwargs):
        '''
        Default command merges documents from the sequence table and virus table
        Chain rethinkdb filter and has_fields commands to the default command
        Return documents from the database that are left after filtering
        '''
        command = r.table(sequence_table).merge(lambda sequence: r.table(virus_table).get(sequence[index])) # via merge
        # command = r.table(sequence_table).eq_join("strain", r.table(virus_table)).zip() # via eq_join, testing suggests that this is slower than the above merge strategy
        command = self.add_present_command(command, **kwargs)
        command = self.add_selections_command(command, **kwargs)
        command = self.add_intervals_command(command, **kwargs)
        command = self.add_public_command(command, **kwargs)
        sequences = list(command.run())
        sequences = filter(None, sequences)
        return list(sequences)

    def add_present_command(self, command, presents=[], **kwargs):
        '''
        Add present fields check to command
        '''
        if len(presents)>0:
            print("Only downloading documents with fields: " + str(presents))
{'fraction_non_alphanumeric': 0.13132085084702586, 'fraction_numerical': 0.005349636988918609, 'mean_word_length': 3.547639733565016, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '756424', 'n_tokens_mistral': 4557, 'n_tokens_neox': 4349, 'n_words': 1159}