<commit_before>"""This module provides essential tools to run reading using indra's own
database. This may also be run as a script; for details run:
`python read_pmids_db --help`
"""
from __future__ import absolute_import, print_function, unicode_literals
from builtins import dict, str
import logging
import random
import zlib
import pickle
import json
from math import log10, floor, ceil
from datetime import datetime
from indra.tools.reading.util.script_tools import get_parser, make_statements, \
                                             StatementData

logger = logging.getLogger('make_db_readings')
if __name__ == '__main__':
    parser = get_parser(
        'A tool to read and process content from the database.',
        ('A file containing a list of ids of the form <id_type>:<id>. '
         'Note that besided the obvious id types (pmid, pmcid, doi, etc.), '
         'you may use trid and tcid to indicate text ref and text content '
         'ids, respectively. Note that these are specific to the database, '
         'and should thus be used with care.')
        )
    parser.add_argument(
        '-m', '--mode',
        choices=['all', 'unread', 'none'],
        default='unread',
        help=('Set the reading mode. If \'all\', read everything, if '
              '\'unread\', only read content that does not have pre-existing '
              'readings of the same reader and version, if \'none\', only '
              'use pre-existing readings. Default is \'unread\'.')
        )
    parser.add_argument(
        '-t', '--temp',
        default='.',
        help='Select the location of the temp file.'
        )
    parser.add_argument(
        '-o', '--output',
        dest='name',
        help=('Pickle all results and save in files labelled as '
              '<NAME>_<output_type>.pkl.'),
        default=None
        )
    parser.add_argument(
        '-b', '--inner_batch',
        dest='b_in',
        help=('Choose the size of the inner batches, which is the number of '
              'text content entires loaded at a given time, and the number of '
              'entries that are read at a time by a reader. The default is '
              '1,000.'),
        default=1000,
        type=int
        )
    parser.add_argument(
        '-B', '--outer_batch',
        dest='b_out',
        default=10000,
        type=int,
        help=('Select the number of ids to read per outer level batch. This '
              'determines the number of readings/statements uploaded/pickled '
              'at a time, and thus also limits the amount of RAM that will be '
              'used. A larger outer batch means more RAM. The default is '
              '10,000.')
        )
    parser.add_argument(
        '--no_reading_upload',
        help='Choose not to upload the reading output to the database.',
        action='store_true'
        )
    parser.add_argument(
        '--no_statement_upload',
        help='Choose not to upload the statements to the databse.',
        action='store_true'
        )
    parser.add_argument(
        '--force_fulltext',
        help='Make the reader only read full text from the database.',
        action='store_true'
        )
    args = parser.parse_args()
    if args.debug and not args.quiet:
        logger.setLevel(logging.DEBUG)

from indra.db import get_primary_db, formats, texttypes
from indra.db import sql_expressions as sql
from indra.db.util import insert_agents
from indra.tools.reading.readers import get_readers, ReadingData, _get_dir


class ReadDBError(Exception):
    pass


# =============================================================================
# Useful functions
# =============================================================================


def _convert_id_entry(id_entry, allowed_types=None):
    ret = [s.strip() for s in id_entry.split(':')]
    if len(ret) != 2:
        raise ReadDBError('Improperly formatted id entry: \"%s\"' % id_entry)
    ret[0] = ret[0].lower()
    if allowed_types is not None and ret[0] not in allowed_types:
        raise ReadDBError('Invalid id type: \"%s\"' % ret[0])
    return ret


def _enrich_reading_data(reading_data_iter, db=None):
    """Get db ids for all ReadingData objects that correspond to a db ref.

    Note that the objects are modified IN PLACE, so nothing is returned, and if
    a copy of the objects is passed as an argument, this function will have no
    effect. This does nothing if the readings are not in the database.
    """
    logger.debug("Enriching the reading data with database refs.")
    if db is None:
        db = get_primary_db()
    possible_matches = db.select_all(
        'readings',
        db.Readings.text_content_id.in_([rd.tcid for rd in reading_data_iter
                                         if rd.reading_id is None])
        )
    for rdata in reading_data_iter:
        for reading in possible_matches:
            if rdata.matches(reading):
                rdata.reading_id = reading.id
                break
    return


# =============================================================================
# Content Retrieval
# =============================================================================


def get_id_dict(id_str_list):
    """Parse the list of id string into a dict."""
    id_types = get_primary_db().TextRef.__table__.columns.keys()
    id_types.remove('id')
    id_types += ['trid', 'tcid']
    id_dict = {id_type: [] for id_type in id_types}
    for id_entry in id_str_list:
        id_type, id_val = _convert_id_entry(id_entry, id_types)
        if id_type in ['trid', 'tcid']:
            id_dict[id_type].append(int(id_val))
        else:
            id_dict[id_type].append(id_val)
    return id_dict


def get_priority_tcids(id_dict, priorities, db=None):
    """For all ids, besides tcids, choose best content available.

    This function will convert all ids to tcids.
    """
    if db is None:
        db = get_primary_db()

    def is_better(new, old):
        return priorities.index(new) < priorities.index(old)

    logger.debug("Getting content prioritized by %s." % str(priorities))
    tcids = set(id_dict.pop('tcid', []))
    clauses = get_clauses(id_dict, db)
    tcid_source = set()
    for clause in clauses:
        q = (db.session.query(db.TextRef.id, db.TextContent.id,
                              db.TextContent.source)
             .filter(db.TextContent.text_ref_id == db.TextRef.id, clause))
        tcid_source |= set(q.all())
    tr_best = {}
    for trid, tcid, source in tcid_source:
        if trid not in tr_best.keys() or is_better(source, tr_best[trid][0]):
            tr_best[trid] = (source, tcid)
    tcids |= {tcid for _, tcid in tr_best.values()}
    return tcids


def get_clauses(id_dict, db):
    """Get a list of clauses to be passed to a db query.

    Note that an empty condition will be returned if id_dict has no ids in it
    (either the dict is empty or all the lists within the dict are empty),
    which will in general have the unexpected effect of selecting everything,
    rather than nothing.

    Parameters
    ----------<commit_msg>Fix prioritization and add some debug messages.<commit_after>"""This module provides essential tools to run reading using indra's own
database. This may also be run as a script; for details run:
`python read_pmids_db --help`
"""
from __future__ import absolute_import, print_function, unicode_literals
from builtins import dict, str
import logging
import random
import zlib
import pickle
import json
from math import log10, floor, ceil
from datetime import datetime
from indra.tools.reading.util.script_tools import get_parser, make_statements, \
                                             StatementData

logger = logging.getLogger('make_db_readings')
if __name__ == '__main__':
    parser = get_parser(
        'A tool to read and process content from the database.',
        ('A file containing a list of ids of the form <id_type>:<id>. '
         'Note that besided the obvious id types (pmid, pmcid, doi, etc.), '
         'you may use trid and tcid to indicate text ref and text content '
         'ids, respectively. Note that these are specific to the database, '
         'and should thus be used with care.')
        )
    parser.add_argument(
        '-m', '--mode',
        choices=['all', 'unread', 'none'],
        default='unread',
        help=('Set the reading mode. If \'all\', read everything, if '
              '\'unread\', only read content that does not have pre-existing '
              'readings of the same reader and version, if \'none\', only '
              'use pre-existing readings. Default is \'unread\'.')
        )
    parser.add_argument(
        '-t', '--temp',
        default='.',
        help='Select the location of the temp file.'
        )
    parser.add_argument(
        '-o', '--output',
        dest='name',
        help=('Pickle all results and save in files labelled as '
              '<NAME>_<output_type>.pkl.'),
        default=None
        )
    parser.add_argument(
        '-b', '--inner_batch',
        dest='b_in',
        help=('Choose the size of the inner batches, which is the number of '
              'text content entires loaded at a given time, and the number of '
              'entries that are read at a time by a reader. The default is '
              '1,000.'),
        default=1000,
        type=int
        )
    parser.add_argument(
        '-B', '--outer_batch',
        dest='b_out',
        default=10000,
        type=int,
        help=('Select the number of ids to read per outer level batch. This '
              'determines the number of readings/statements uploaded/pickled '
              'at a time, and thus also limits the amount of RAM that will be '
              'used. A larger outer batch means more RAM. The default is '
              '10,000.')
        )
    parser.add_argument(
        '--no_reading_upload',
        help='Choose not to upload the reading output to the database.',
        action='store_true'
        )
    parser.add_argument(
        '--no_statement_upload',
        help='Choose not to upload the statements to the databse.',
        action='store_true'
        )
    parser.add_argument(
        '--force_fulltext',
        help='Make the reader only read full text from the database.',
        action='store_true'
        )
    args = parser.parse_args()
    if args.debug and not args.quiet:
        logger.setLevel(logging.DEBUG)

from indra.db import get_primary_db, formats, texttypes
from indra.db import sql_expressions as sql
from indra.db.util import insert_agents
from indra.tools.reading.readers import get_readers, ReadingData, _get_dir


class ReadDBError(Exception):
    pass


# =============================================================================
# Useful functions
# =============================================================================


def _convert_id_entry(id_entry, allowed_types=None):
    ret = [s.strip() for s in id_entry.split(':')]
    if len(ret) != 2:
        raise ReadDBError('Improperly formatted id entry: \"%s\"' % id_entry)
    ret[0] = ret[0].lower()
    if allowed_types is not None and ret[0] not in allowed_types:
        raise ReadDBError('Invalid id type: \"%s\"' % ret[0])
    return ret


def _enrich_reading_data(reading_data_iter, db=None):
    """Get db ids for all ReadingData objects that correspond to a db ref.

    Note that the objects are modified IN PLACE, so nothing is returned, and if
    a copy of the objects is passed as an argument, this function will have no
    effect. This does nothing if the readings are not in the database.
    """
    logger.debug("Enriching the reading data with database refs.")
    if db is None:
        db = get_primary_db()
    possible_matches = db.select_all(
        'readings',
        db.Readings.text_content_id.in_([rd.tcid for rd in reading_data_iter
                                         if rd.reading_id is None])
        )
    for rdata in reading_data_iter:
        for reading in possible_matches:
            if rdata.matches(reading):
                rdata.reading_id = reading.id
                break
    return


# =============================================================================
# Content Retrieval
# =============================================================================


def get_id_dict(id_str_list):
    """Parse the list of id string into a dict."""
    id_types = get_primary_db().TextRef.__table__.columns.keys()
    id_types.remove('id')
    id_types += ['trid', 'tcid']
    id_dict = {id_type: [] for id_type in id_types}
    for id_entry in id_str_list:
        id_type, id_val = _convert_id_entry(id_entry, id_types)
        if id_type in ['trid', 'tcid']:
            id_dict[id_type].append(int(id_val))
        else:
            id_dict[id_type].append(id_val)
    return id_dict


def get_priority_tcids(id_dict, priorities, db=None):
    """For all ids, besides tcids, choose best content available.

    This function will convert all ids to tcids.
    """
    if db is None:
        db = get_primary_db()

    def is_better(new, old):
        if new in priorities and old in priorities:
            return priorities.index(new) < priorities.index(old)
        return False

    logger.debug("Getting content prioritized by %s." % str(priorities))
    tcids = set(id_dict.pop('tcid', []))
    clauses = get_clauses(id_dict, db)
    tcid_source = set()
    for clause in clauses:
        q = (db.session.query(db.TextRef.id, db.TextContent.id,
                              db.TextContent.source)
             .filter(db.TextContent.text_ref_id == db.TextRef.id, clause))
        id_set = set(q.all())
        logger.debug("Got %d more ids." % len(id_set))
        tcid_source |= id_set
    logger.debug("Got %d id's total." % len(tcid_source))
    tr_best = {}
    for trid, tcid, source in tcid_source:
        if trid not in tr_best.keys() or is_better(source, tr_best[trid][0]):
            tr_best[trid] = (source, tcid)
    tcids |= {tcid for _, tcid in tr_best.values()}
    return tcids


def get_clauses(id_dict, db):
    """Get a list of clauses to be passed to a db query.

    Note that an empty condition will be returned if id_dict has no ids in it
    (either the dict is empty or all the lists within the dict are empty),
    which will in general have the unexpected effect of selecting everything,
    rather than nothing.

    Parameters
    ----------
{'fraction_non_alphanumeric': 0.136891592920354, 'fraction_numerical': 0.003595132743362832, 'mean_word_length': 2.5392708588206507, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7046730', 'n_tokens_mistral': 4214, 'n_tokens_neox': 4055, 'n_words': 1404}