"""For creating and querying index files that can be used for
searching and mapping to clinical term databases.

The command line program in here that will build a index tables for the
UMLS database, but the actual API code is generalisable.

The index tables are similar to what already exists in the database but
generated using NLTK as opposed to the UMLS norm program.
"""

from umls_tools import __version__, __name__ as pkg_name, orm, common
from sqlalchemy_config import config as cfg
from tqdm import tqdm
from pyaddons import utils, log
from merge_sort import chunks, merge
from operator import itemgetter
from nltk.tokenize.destructive import NLTKWordTokenizer
from nltk import stem
from nltk.corpus import stopwords
import gzip
import shutil
import argparse
import os
import re
import tempfile
import sys
import csv
import math
import unicodedata
import pprint as pp

# deal with long fields in csv
csv.field_size_limit(sys.maxsize)

_PROG_NAME = "umls-index"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class IndexTerms(object):
    """Create an index files for term/ID combinations. This is designed to
    generalise outside of the UMLS.

    Parameters
    ----------
    chunksize : `int` optional, default: `100000`
        The number of rows to hold in memory during the chunking phase of an
        external merge sort. More rows will use more memory but will be
        faster.
    max_files : `int`, optional, default: `16`
        The maximum number of tempfiles to open at a single time during the
        merge phase of an external merge sort.
    tmpdir : `str`, optional, default: `NoneType`
        The temp directory to use for the merge sort. If not provided then the
        default temp location is used.
    """

    END_WORDS = [
        re.compile(r"\s+NOS$", re.IGNORECASE),
        re.compile(r"not otherwise specified", re.IGNORECASE),
    ]
    """Some comment stop words/phrases that occur at the end of terms
    (`list` of `re.Pattern`)
    """
    NON_TEXT_REGEXP = re.compile(r"[^A-Za-z0-9]")
    """A regular expression for identifying non-characters/digits
    (`re.Pattern`)
    """
    _TOKENISER = NLTKWordTokenizer()
    """The tokeniser to use (`nltk.tokenize.NLTKWordTokenizer`)
    """
    _STEMMER = stem.SnowballStemmer("english")
    """The stemmer to use (`nltk.stem.SnowballStemmer`)
    """

    try:
        STOPWORDS = set(stopwords.words("english"))
        """Get the stopwords to exclude from the terms being indexed
        (`set` of `str`)
        """
    except (LookupError, OSError, ValueError):
        STOPWORDS = set()

    SORT_COL = 4
    """The column number in the temp file that holds the token value to sort
    on (`int`)
    """
    INDEX_ID = "index_id"
    """The name of the index identifer column, this is an increment for each
    row (`str`)
    """
    RANK_ORDER = "rank_order"
    """The name of the rank order column. This is the integer rank position
    that the token occupies from the start of the term (0-based) (`str`).
    """
    START_POS = "start_pos"
    """The name of the start position column. This is the integer character
    position that the word that the token represents occupies from the start
    of the term (0-based) (`str`).
    """
    END_POS = "start_end"
    """The name of the end position column. This is the integer character
    position that the word that the token represents occupies from the start
    of the term (0-based) (`str`).
    """
    TERM_ID = "term_id"
    """The name of the term ID column. The term ID is the concept identifier
    that maps to a concept term (`str`).
    """
    TERM_STR = "term_str"
    """The name of the full term string column, this contains the unprocessed
    term (`str`).
    """
    TOKEN_STR = "token"
    """The name of the token string column, this contains the processed token
    term (`str`).
    """
    TOKEN_ID = "token_id"
    """The name of the token ID column. Each unique token is given an integer
    identifier (1-based) (`str`).
    """

    HEADER = [
        # Integer increment
        INDEX_ID,
        # Integer ID for the token
        TOKEN_ID,
        START_POS,
        END_POS,
        RANK_ORDER,
        TERM_ID,
        TOKEN_STR,
        TERM_STR,
    ]
    """The header of the temp output file and the key names of the yielded
    sorted dict rows of the index terms (`list` of `str`).
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, chunksize=100000, max_files=16, tmpdir=None):
        # The root temp dir, a directory will be created in here for sorting
        # files
        self.tmproot = tmpdir
        self.max_files = max_files
        self.chunksize = chunksize
        self.tmpdir = None
        self.chunker = None
        # Will hold the status of the sorted index file
        self._is_closed = True
        self.max_term_len = 0
        self.max_token_len = 0
        self.nterms = 0
        self.ntokens = 0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __enter__(self):
        self.open()
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __exit__(self, *args):
        self.close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Open the index add rows to."""
        if self._is_closed is True:
            self.tmpdir = tempfile.mkdtemp(dir=self.tmproot)
            self.max_term_len = 0
            self.max_token_len = 0
            self.nterms = 0
            self.ntokens = 0
            self.chunker = chunks.CsvSortedChunks(
                self.tmpdir,
                key=itemgetter(self.SORT_COL),
                chunksize=self.chunksize,
                chunk_prefix="index_",
                delimiter="\t",
                lineterminator=os.linesep,
            )
            self.chunker.open()
            self._is_closed = False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close the index and remove all temporary files."""
        if self._is_closed is False:
            self.chunker.close()
        self._is_closed = True

        try:
            self.chunker.delete()
            shutil.rmtree(self.tmpdir)
        except AttributeError:
            # Chunker not created?
            pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_term(self, term_id, term_str):
        """Add a term/ID combination to the index.

        Parameters
        ----------
        term_id : `str`
            The term identifier.
        term_str : `str`
            The term string
        """
        if self._is_closed is True:
            raise IOError("can't add terms to a closed index")
        self.max_term_len = max(self.max_term_len, len(term_str))
        self.nterms += 1

        for span, rank_idx, text in self.process_str(term_str):
            start_idx, end_idx = span
            self.ntokens += 1
            self.max_token_len = max(self.max_token_len, len(text))
            self.chunker.add_row(
                [start_idx, end_idx, rank_idx, term_id, text, term_str]
            )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def fetch(self):
        """Yield all the tokens in token sort order.

        Yields
        ------
        ordered_row : `list`
            A row containing the token information in token sort order. The
            element order is ``[0]`` - row number ``[1]`` - token identifier
            (a counter for each unique token) ``[2]`` - The start position of
            the token (pre-processing).  ``[3]`` - The start position of the
            token (pre-processing).  ``[2]`` - The rank order of the token
            (pre-processing). ``[3]`` The term identifer in the original
            database. This may be a string or and integer, for different
            databases. ``[4]`` The token text. ``[5]`` - The full term string.
        """
        if self._is_closed is False:
            self.chunker.close()

        self._is_closed = True

        with merge.CsvIterativeHeapqMerge(
            self.chunker.chunk_files,
            max_files=self.max_files,
            key=itemgetter(self.SORT_COL),
            tmpdir=self.tmpdir,
            header=False,
            delete=False,
            csv_kwargs=dict(delimiter="\t", lineterminator=os.linesep),
        ) as m:
            # Row ID is an integer representation of the term ID
            # which is treated as a string
            token_id = None
            current_token = None
            for idx, row in enumerate(m, 1):
                if row[self.SORT_COL] != current_token:
                    try:
                        token_id += 1
                        current_token = row[self.SORT_COL]
                    except TypeError:
                        token_id = 1
                yield [idx, token_id] + row

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def process_str(cls, text):
        """Process a term string.

        Parameters
        ----------
        text : `str`
            The term string to process.

        Returns
        -------
        tokens : `list` of `tuple`
            The tokens, each tuple is the span in the term
            (tuple start, end (0-based)). The rank order in the term
            (1-based). The final token string.

        Notes
        -----
        The text processing involves:

        1. Remove end words such as NOS, Not Otherwise Specified
        2. Strip the punctuation (non character/numerics).
        3. Normalize accented unicode characters to ASCII.
        4. Tokenise using the default NLTK tokeniser.
        5. Make lowercase
        6. Remove stop words.
        7. Stemming
        """
        text = cls.remove_end_words(text)
        text = cls.strip_punct(text)
        text = unicodedata.normalize("NFKD", text)
        text = text.encode("ascii", "ignore").decode("ascii")
        tokens = cls.tokenise(text.lower())
        tokens = cls.remove_stop_words(tokens)
        return [(span, idx, cls.stem(text)) for span, idx, text in tokens]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def strip_punct(cls, text):
        """Remove all non-words/numbers and replace with spaces.

        Parameters
        ----------
        text : `str`
            The text to process.

        Returns
        -------
        text : `str`
            The potentially processed text.
        """
        return cls.NON_TEXT_REGEXP.sub(" ", text)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def remove_end_words(cls, text):
        """Remove ending words like NOS or Not Otherwise Specified.

        Parameters
        ----------
        text : `str`
            The text to process.

        Returns
        -------
        text : `str`
            The potentially processed text.
        """
        for i in cls.END_WORDS:
            text = i.sub("", text)
        return text

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def remove_stop_words(cls, tokens):
        """Remove stop words.

        Parameters
        ----------
        tokens : `list` of `str`
            The tokens to process.

        Returns
        -------
        text : `str`
            The potentially processed text.
        """
        return [
            (span, idx, text)
            for span, idx, text in tokens
            if text not in cls.STOPWORDS
        ]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def tokenise(cls, text):
        """Texenise the text.

        Parameters
        ----------
        text : `str`
            The text to process.

        Returns
        -------
        tokens : `list` or (`tuple`, `int`, `str`)
            Each nested tuple is the token span (`start`, `end`), the token
            rank order and the token text.
        """
        tokens = cls._TOKENISER.tokenize(text)
        spans = list(cls._TOKENISER.span_tokenize(text))
        return [(i, j, k) for i, j, k in zip(spans, range(len(spans)), tokens)]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def stem(cls, text):
        """Tokenise the text.

        Parameters
        ----------
        text : `str`
            The text to process.

        Returns
        -------
        tokens : `list` or (`tuple`, `int`, `str`)
            Each nested tuple is the token span (`start`, `end`), the token
            rank order and the token text.
        """
        return cls._STEMMER.stem(text)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class BaseSearchIndex(object):
    """A base class for searching indexed tables. This is not UMLS specific
    and will operate on any database implementing the index tables.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session to issue the queries against the database.
    orm_module : `module`
        The imported module that contains the SQLAlchemy ORM models. It is
        expected that this has a `TermIndexLookup` table, see Notes for
        details.
    nterms : `int`
        The total number of terms that have been indexed.

    Notes
    -----
    There are some ORM mixin classes that can be used to build the correct
    index tables. The usage of the mixins can be seen in the ``umls_tools.orm``
    module.

    See also
    --------
    umls_tools.orm_mixin.TermIndexLookupMixin
    umls_tools.orm_mixin.TermIndexMapMixinInt
    umls_tools.orm_mixin.TermIndexMapMixinStr
    umls_tools.orm.TermIndexLookup
    umls_tools.orm.TermIndexMap
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, session, orm_module, nterms):
        self.session = session
        self.nterms = nterms
        self.orm_module = orm_module
        self.SQL = self.session.query(
            self.orm_module.TermIndexLookup, self.orm_module.TermIndexMap
        ).join(
            self.orm_module.TermIndexMap,
            self.orm_module.TermIndexMap.token_id
            == self.orm_module.TermIndexLookup.token_id,
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def search_term(self, term):
        """Search for a term against an indexed table. The term is processed in
        a similar way to the indexed terms.

        Parameters
        ----------
        term : `str`
            The term to search for.

        Returns
        -------
        match : `list` of `list`
            The elements are ordered from highest score to lowest and each
            sub-list  (matching term) has the following elements:

            0. The match score source_coverage * target_coverage *
               sum(log10(total index terms/token freq)).
            1. The length of the match.
            2. The length of the whole matching term
            3. A set of tuples of the containing the matching positions of
               individual tokens.
            4. The token mapping rows in the database (SQLAlchemy objects)
        """
        tokens = IndexTerms.process_str(term)
        if not tokens:
            return []
        term_len = len(term)

        matches = {}
        for pos, rank, token in tokens:
            for row in self.query_index(token):
                try:
                    data = matches[row[1].term_id]
                    data[0] += math.log10(self.nterms / row[0].token_prob)
                    data[1] += pos[1] - pos[0]
                    data[3].add((pos, rank, token))
                    data[4].append(row)
                except KeyError:
                    x = set()
                    x.add((pos, rank, token))
                    matches[row[1].term_id] = [
                        math.log10(self.nterms / row[0].token_prob),
                        pos[1] - pos[0],
                        row[1].term_len,
                        x,
                        [row],
                    ]

        for i in matches.values():
            # This approximates the proportion of the user specified term
            # that is covered by the matches
            search_found_len = sum(
                [e - s for s, e in set([pos for pos, rank, token in i[3]])]
            )
            search_wt = min(1, search_found_len / term_len)
            match_wt = i[1] / i[2]
            i.append(search_wt)
            i.append(match_wt)
            i[0] = i[0] * search_wt * match_wt

        return sorted(matches.values(), key=itemgetter(0), reverse=True)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def query_index(self, token):
        """Query the index for the token."""
        for row in self.SQL.filter(
            self.orm_module.TermIndexLookup.token_str == token
        ):
            yield row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():  # pragma: no cover
    """The main entry point. For API access see
    ``umls_tools.admin.index.build_umls_index``
    """
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)

    logger = log.init_logger(_PROG_NAME, verbose=args.verbose)
    log.log_prog_name(logger, pkg_name, __version__)
    log.log_args(logger, args)

    # Get a sessionmaker to create sessions to interact with the database
    sm = cfg.get_new_sessionmaker(
        args.dburl,
        conn_prefix=common.DEFAULT_PREFIX,
        config_arg=args.config,
        config_env=None,
        config_default=common.DEFAULT_CONFIG,
        exists=True,
    )

    load_verbose = False
    if args.verbose > 1:
        load_verbose = True

    try:
        logger.info("building a UMLS index, this may take some time...")
        if not IndexTerms.STOPWORDS:
            logger.warning(
                "No NLTK stopwords available; indexing will run without "
                'stopword filtering. Install with: python -c "import nltk; '
                "nltk.download('stopwords')\""
            )
        nterms, ntokens, mterm_len, mtoken_len = build_umls_index(
            sm,
            tmpdir=args.tmp,
            verbose=load_verbose,
            commit_every=args.commit_every,
            chunksize=args.chunksize,
        )
        logger.info(f"# terms: {nterms}")
        logger.info(f"# tokens: {ntokens}")
        logger.info(f"max term length: {mterm_len}")
        logger.info(f"max token_length: {mtoken_len}")
        log.log_end(logger)
    except (OSError, FileNotFoundError):
        raise
    except (BrokenPipeError, IOError, KeyboardInterrupt):
        # Python flushes standard streams on exit; redirect remaining
        # output to devnull to avoid another BrokenPipeError at shutdown
        devnull = os.open(os.devnull, os.O_WRONLY)

        try:
            os.dup2(devnull, sys.stdout.fileno())
        except Exception:
            pass
        log.log_interrupt(logger)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():  # pragma: no cover
    """Use argparse to parse the command line arguments

    Returns
    -------
    parser : `argparse.ArgumentParser`
        The arguments expected on the command line.
    """
    parser = argparse.ArgumentParser(description=_DESC)

    parser.add_argument(
        "dburl",
        nargs="?",
        type=str,
        help="An SQLAlchemy connection URL or filename if "
        "using SQLite. If you do not want to put full "
        "connection parameters on the cmd-line use the "
        "config file (--config) and config section"
        "(--config-section) to supply the parameters",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default=common.get_default_config_cli_path(),
        help="The location of the config file",
    )
    parser.add_argument(
        "-T",
        "--tmp",
        type=str,
        help="The location of tmp, if not provided will use the system tmp",
    )
    parser.add_argument(
        "--commit-every",
        type=int,
        default=10000,
        help="The commit to the database after every --commit-every rows",
    )
    parser.add_argument(
        "--chunksize",
        type=int,
        default=100000,
        help="The max rows to keep in memory when sorting",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        default=0,
        action="count",
        help="give more output, ``--vv`` turns on progress monitoring",
    )
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):  # pragma: no cover
    """Use argparse to parse the command line arguments.

    Parameters
    ----------
    parser : `argparse.ArgumentParser`
        The arguments expected on the command line.

    Returns
    -------
    args : `ArgumentParser.args`
        The arguments from parsing the cmd line args.
    """
    args = parser.parse_args()
    # Because the default is ~/ for Sphinx docs
    args.config = os.path.expanduser(args.config)
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_umls_index(
    sm,
    tmpdir=None,
    verbose=False,
    commit_every=10000,
    lang="ENG",
    chunksize=100000,
    max_files=16,
):
    """Build an index data for all the UMLS terms.

    Parameters
    ----------
    sm : `sqlalchemy.Sessionmaker`
        The SQLAlchemy session maker.
    tmpdir : `str`, optional, default: `NoneType`
        The temp directory to use for the merge sort. If not provided then the
        default temp location is used.
    verbose : `bool`, optional, default: `False`
        Monitor progress.
    commit_every : `int`, optional, default: `10000`
        Commit to the database every 10000 rows.
    lang : `str`, optional, default: `ENG`
        Restrict the building of the index to a specific language.
    chunksize : `int` optional, default: `100000`
        The number of rows to hold in memory during the chunking phase of an
        external merge sort. More rows will use more memory but will be
        faster.
    max_files : `int`, optional, default: `16`
        The maximum number of tempfiles to open at a single time during the
        merge phase of an external merge sort.

    Returns
    -------
    nterms : `int`
        The total number of input terms processed.
    ntokens : `int`
        The total number of tokens generated.
    max_term_len : `int`
        The maximum input term length in characters.
    max_token_len : `int`
        The maximum token length in characters.

    Notes
    -----
    The index tables are similar to the ones that will already exist in the
    UMLS however, will be build from Python libraries and not the norm
    program.
    """
    session = sm()

    try:
        q = session.query(orm.MrConso.AUI, orm.MrConso.STR)

        if lang is not None:
            q = q.filter(orm.MrConso.LAT == lang)
        q = q.distinct()
        return build_index(
            session,
            q,
            orm,
            tmpdir=tmpdir,
            verbose=verbose,
            commit_every=commit_every,
            chunksize=chunksize,
        )
    finally:
        session.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_index(
    session,
    query,
    orm_module,
    tmpdir=None,
    verbose=False,
    commit_every=10000,
    chunksize=100000,
    max_files=16,
):
    """Build an index data from a generic input. This can be used to
    generalise the index table building process.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session.
    query : `sqlalchemy.Query`
        The query object.
    orm_module : `module`
        the ORM module. Must have a ``TermIndexLookup`` class and a
        ``TermIndexMap``. These can be created from mixins in the
        ``umls_tools.orm_mixin`` module.
        See ``umls_tools.orm.TermIndexLookup`` and
        ``umls_tools.orm.TermIndexMap``
    tmpdir : `str`, optional, default: `NoneType`
        The temp directory to use for the merge sort. If not provided then the
        default temp location is used.
    verbose : `bool`, optional, default: `False`
        Monitor progress.
    commit_every : `int`, optional, default: `10000`
        Commit to the database every 10000 rows.
    lang : `str`, optional, default: `ENG`
        Restrict the building of the index to a specific language.
    chunksize : `int` optional, default: `100000`
        The number of rows to hold in memory during the chunking phase of an
        external merge sort. More rows will use more memory but will be
        faster.
    max_files : `int`, optional, default: `16`
        The maximum number of tempfiles to open at a single time during the
        merge phase of an external merge sort.

    Returns
    -------
    nterms : `int`
        The total number of input terms processed.
    ntokens : `int`
        The total number of tokens generated.
    max_term_len : `int`
        The maximum input term length in characters.
    max_token_len : `int`
        The maximum token length in characters.

    Notes
    -----
    The index tables are similar to the one that will already exist in the
    UMLS however, will be build from Python libraries and not the norm
    program. To use this you will need to implement two index tables in your
    database, ``TermIndexMap`` and ``TermIndexLookup``. There are mixins, for
    these in the ``umls_orm.mixin`` module.

    Note that each time this is called, it will drop the old index tables
    and create new ones.

    See also
    --------
    umls_tools.orm_mixin.TermIndexLookupMixin
    umls_tools.orm_mixin.TermIndexMapMixinInt
    umls_tools.orm_mixin.TermIndexMapMixinStr
    umls_tools.orm.TermIndexLookup
    umls_tools.orm.TermIndexMap
    """
    tmpout = utils.get_temp_file(dir=tmpdir, suffix=".txt.gz")
    try:
        tqdm_kwargs = dict(
            disable=not verbose,
            desc="[info] generating index...",
            unit=" terms",
            leave=False,
            total=query.count(),
        )

        break_at = -10000
        nrows = 0
        with gzip.open(tmpout, "wt") as outf:
            writer = csv.writer(outf, delimiter="\t", lineterminator=os.linesep)
            writer.writerow(IndexTerms.HEADER)

            id_cast = str
            try:
                id_cast = orm_module.TermIndexMap.term_id.property.columns[
                    0
                ].type.python_type
            except Exception:
                id_cast = str

            with IndexTerms(tmpdir=tmpdir, chunksize=chunksize) as idx:
                for term_id, term in tqdm(query, **tqdm_kwargs):
                    try:
                        idx.add_term(term_id, term)
                    except (TypeError, ValueError) as e:
                        e.error_record = (term_id, term)
                        raise
                    nrows += 1
                    if nrows == break_at:
                        break

                tqdm_kwargs = dict(
                    disable=not verbose,
                    desc="[info] writing to temp...",
                    unit=" rows",
                    leave=False,
                    total=idx.ntokens,
                )

                for row in tqdm(idx.fetch(), **tqdm_kwargs):
                    writer.writerow(row)
        ntokens = idx.ntokens

        tqdm_kwargs = dict(
            disable=not verbose,
            desc="[info] loading into the database...",
            unit=" rows",
            leave=False,
            total=ntokens,
        )

        # Drop and recreate the token tables
        for table in (
            orm_module.TermIndexMap.__table__,
            orm_module.TermIndexLookup.__table__,
        ):
            table.drop(session.get_bind(), checkfirst=True)

        orm_module.TermIndexLookup.__table__.create(session.get_bind())
        orm_module.TermIndexMap.__table__.create(session.get_bind())

        # Now we write to the database
        with gzip.open(tmpout, "rt") as infile:
            reader = csv.DictReader(
                infile, delimiter="\t", lineterminator=os.linesep
            )
            row = next(reader)
            curr_token_str = row[IndexTerms.TOKEN_STR]
            curr_token_id = int(row[IndexTerms.TOKEN_ID])
            ncurr_token = 1
            mappings = [
                orm_module.TermIndexMap(
                    token_id=curr_token_id,
                    start_pos=int(row[IndexTerms.START_POS]),
                    end_pos=int(row[IndexTerms.END_POS]),
                    rank_order=int(row[IndexTerms.RANK_ORDER]),
                    term_id=id_cast(row[IndexTerms.TERM_ID]),
                    term_len=len(row[IndexTerms.TERM_STR]),
                )
            ]
            inserts = []
            proc_rows = 1
            for row in tqdm(reader, **tqdm_kwargs):
                if row[IndexTerms.TOKEN_STR] != curr_token_str:
                    prob_token = ncurr_token / ntokens
                    curr_token = orm_module.TermIndexLookup(
                        token_id=curr_token_id,
                        token_str=curr_token_str,
                        token_count=ncurr_token,
                        token_prob=prob_token,
                    )
                    inserts.append(curr_token)

                    if proc_rows >= commit_every:
                        try:
                            session.bulk_save_objects(inserts)
                            session.bulk_save_objects(mappings)
                            session.commit()
                        except Exception:
                            pp.pprint(inserts[:10])
                            raise
                        inserts = []
                        mappings = []
                        proc_rows = 0
                    curr_token_str = row[IndexTerms.TOKEN_STR]
                    curr_token_id = int(row[IndexTerms.TOKEN_ID])
                    ncurr_token = 0

                ncurr_token += 1
                mappings.append(
                    orm_module.TermIndexMap(
                        token_id=curr_token_id,
                        start_pos=int(row[IndexTerms.START_POS]),
                        end_pos=int(row[IndexTerms.END_POS]),
                        rank_order=int(row[IndexTerms.RANK_ORDER]),
                        term_id=id_cast(row[IndexTerms.TERM_ID]),
                        term_len=len(row[IndexTerms.TERM_STR]),
                    )
                )
                proc_rows += 1
            # A final write
            if len(mappings) > 0:
                prob_token = ncurr_token / ntokens
                curr_token = orm_module.TermIndexLookup(
                    token_id=curr_token_id,
                    token_str=curr_token_str,
                    token_count=ncurr_token,
                    token_prob=prob_token,
                )
                inserts.append(curr_token)
                try:
                    session.bulk_save_objects(inserts)
                    session.bulk_save_objects(mappings)
                    session.commit()
                except Exception:
                    pp.pprint(inserts[0:10])
                    raise

    except Exception:
        os.unlink(tmpout)
        raise

    return idx.nterms, ntokens, idx.max_term_len, idx.max_token_len
