"""Add a additional UNIQUE_CUI table to an existing UMLS database build.
This will contain the top ranked AUI for each CUI record, it will also have
the number of AUIs in a CUI and the probability of querying that CUI (ID) at
random from the MRCONSO table.

Although this has been build using SQLAlchemy, so far it has only been tested
on SQLite and MariaDB. The loading process might need to be adapted for other
databases.
"""

from umls_tools import __version__, __name__ as pkg_name, log, orm as o, common
from sqlalchemy_config import config as cfg
from sqlalchemy_utils.functions import database_exists
from tqdm import tqdm
from sqlalchemy import and_
from sqlalchemy.exc import OperationalError, ObjectNotExecutableError
import re
import argparse
import os
import sys
import tempfile
import csv
import gzip
import shutil
import warnings
import pprint as pp

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


_PROG_NAME = "umls-unique-cui"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""
_SUB_DELIM = "|"
"""A sub-delimiter for the semantic types in the UNIQUE_CUI table (`str`)
"""
_TEMP_DELIM = "\t"
"""A delimiter to use in the temp file.
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():  # pragma: no cover
    """The main entry point. For API access see
    ``umls_tools.admin.unique_cui.create_unique_cui``
    """
    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)

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

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

    # Make sure the database exists (or does not exist)
    try:
        if database_exists(sm().bind.url) is False:
            raise FileNotFoundError("database does not exist")
    except ObjectNotExecutableError:
        warnings.warn(
            "unable to check if database exists: sqlalchemy-utils does"
            " not work well with SQLAlchemy >= 1.4"
        )

    tmpobj, tmpfile = tempfile.mkstemp(dir=args.tmpdir)
    os.close(tmpobj)

    try:
        logger.info(
            "building the UMLS unique CUI table, this may take some time..."
        )
        ncuis = create_unique_cui(
            sm(),
            tmpfile,
            verbose=load_verbose,
            commit_every=args.commit_every,
            exists=args.exists,
        )
        logger.info(f"Added {ncuis} CUIs to the UNIQUE_CUI table")
        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)
    finally:
        if args.outfile is not None:
            logger.info(f"Saving to {args.outfile}")
            shutil.move(tmpfile, args.outfile)
        else:
            os.unlink(tmpfile)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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(
        "-s",
        "--config-section",
        type=str,
        default=common.DEFAULT_SECTION,
        help="The section name in the config file",
    )
    parser.add_argument(
        "-T",
        "--tmpdir",
        type=str,
        help="The table is written to tmp first before being loaded",
    )
    parser.add_argument(
        "-O",
        "--outfile",
        type=str,
        help="A file name to use if you want a text copy of the UNIQUE_CUI "
        "table",
    )
    parser.add_argument(
        "--commit-every",
        type=int,
        default=100000,
        help="The commit to the database after every --commit-every rows",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        default=0,
        action="count",
        help="give more output, ``--vv`` turns on progress monitoring",
    )
    parser.add_argument(
        "-e",
        "--exists",
        action="store_true",
        help="The ``UNIQUE_CUI`` table already exists, this will assume that"
        " the user has pre-created a database and will not throw errors if the"
        "the table already exists. ",
    )
    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 create_unique_cui(
    session, outfile, verbose=False, commit_every=50000, exists=False
):
    """The main API entry point to create a ``unique_cui`` table.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAchemy session used to interact with the database. Must return
        the engine from a ``Session.get_bind()`` call.
    outfile : `str`
        A file name to write an intermediate file. This is then read in and
        loaded. This is because of SQLite issues with querying and writing.
    verbose : `bool`, optional, default: `False`
        Show progress of the table build.
    commit_every : `int`, optional, default: `10000`
        The number of rows that are added to a session before
        ``session.commit()`` is called. If you set this to 0 or negative only a
        single commit will be called. Low values will be slow.
    exists : `bool`, optional, default: `False`
        Should the table be expected to already exist? If `False`, then
        existing tables will cause and error, if True then no error will be
        raised. This is to handle situations where the user is unable to create
        databases and tables.

    Returns
    -------
    ncui : `int`
        The number of CUIs added to the ``UNIQUE_CUI`` table.
    """
    if exists is True:
        try:
            o.UniqueUmlsCui.__table__.create(
                bind=session.get_bind(), checkfirst=True
            )
        except OperationalError as e:
            if re.search(r"already exists", str(e), re.IGNORECASE) is None:
                raise
    else:
        o.UniqueUmlsCui.__table__.create(
            bind=session.get_bind(), checkfirst=False
        )

    # How many lines are there
    nlines = session.query(o.MrConso.CUI).distinct().count()

    tqdm_kwargs = dict(
        disable=not verbose,
        unit=" rows",
        desc="[info] collecting unique CUIs",
        total=nlines,
        leave=False,
    )

    added = 0
    ncui = 0

    header = [
        "AUI",
        "CUI",
        "STR",
        "LAT",
        "SAB",
        "TTY",
        "CODE",
        "RANK",
        "NAUI",
        "PROB_CUI",
        "SEMANTIC_TYPES",
    ]

    # try:
    # Query into a temp file
    with gzip.open(outfile, "wt") as outcsv:
        writer = csv.writer(outcsv, delimiter=_TEMP_DELIM)
        writer.writerow(header)
        for row in tqdm(get_rank_unique_cuis(session), **tqdm_kwargs):
            row = list(row)
            row[-1] = _SUB_DELIM.join(row[-1])
            writer.writerow(row)
            ncui += 1
    # finally:
    #     progress.close()

    tqdm_kwargs = dict(
        disable=not verbose,
        unit=" cuis",
        desc="[info] building UNIQUE_CUI",
        total=ncui,
        leave=False,
    )

    added = 0

    # read from temp and load into the DB
    with gzip.open(outfile, "rt") as incsv:
        reader = csv.DictReader(incsv, delimiter=_TEMP_DELIM)
        for row in tqdm(reader, **tqdm_kwargs):
            row["RANK"] = int(row["RANK"])
            row["NAUI"] = int(row["NAUI"])
            row["PROB_CUI"] = float(row["PROB_CUI"])
            try:
                session.add(o.UniqueUmlsCui(**row))
            except Exception:
                pp.pprint(row)
                raise
            added += 1
            if added == commit_every:
                session.commit()
                added = 0
        if added > 0:
            session.commit()
    return ncui


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_rank_unique_cuis(session):
    """Get the unique CUIs by MRRANK. This will return the top ranked AUI for
    each CUI and uses the MRRANK table (highest rank is best).

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAchemy session used to interact with the database.

    Yields
    ------
    aui : `str`
        The AUI code for the top ranked AUI.
    cui : `str`
        The CUI code.
    term : `str`
        The term string.
    lat : `str`
        The language of the term.
    sab : `str`
        The source abbreviation. This is the name of the original source for
        the term in abbreviated form.
    tty : `str`
        The term type.
    code : `str`
        Code for the term in the SAB (original source for the term).
    nauis : `int`
        The number of AUIs for the CUI.
    prob_cui : `float`
        The probability of getting the CUI from a random query of MRCONSO.
    stys : `list` of `str`
        The semantic types of the CUI.
    """
    nlines = session.query(o.MrConso.AUI).count()

    q = (
        session.query(
            o.MrConso.CUI,
            o.MrConso.AUI,
            o.MrConso.STR,
            o.MrConso.LAT,
            o.MrConso.SAB,
            o.MrConso.TTY,
            o.MrConso.CODE,
            o.MrSty.STY,
            o.MrRank.MRRANK_RANK,
        )
        .join(
            o.MrRank,
            and_(o.MrConso.SAB == o.MrRank.SAB, o.MrConso.TTY == o.MrRank.TTY),
        )
        .join(o.MrSty, o.MrConso.CUI == o.MrSty.CUI)
        .order_by(o.MrConso.CUI)
        .yield_per(1000)
    )

    cuis = []
    for i in q:
        try:
            if i.CUI != cuis[0].CUI:
                # Now process
                yield _process_rank_cui(cuis, nlines)
                cuis = []
        except IndexError:
            # First loop
            pass

        cuis.append(i)

    if len(cuis) > 0:
        yield _process_rank_cui(cuis, nlines)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _process_rank_cui(cuis, data_len):
    """Process a batch of cuis to a single high ranked CUI, this is the UMLS
    method.

    Parameters
    ----------
    cuis : `list` of `namedTuple`
        The CUIs to process, queried using SQLAlchemy (so results are named
        tuples)
    data_len : `int`
        The number of AUIs in the MRCONSO table, used to compute the
        probability of a CUI.

    Returns
    -------
    aui : `str`
        The AUI code for the top ranked AUI.
    cui : `str`
        The CUI code.
    term : `str`
        The term string.
    lat : `str`
        The language of the term.
    sab : `str`
        The source abbreviation. This is the name of the original source for
        the term in abbreviated form.
    tty : `str`
        The term type.
    code : `str`
        Code for the term in the SAB (original source for the term).
    rank : `int`
        The rank of the top AUI.
    nauis : `int`
        The number of AUIs for the CUI.
    prob_cui : `float`
        The probability of getting the CUI from a random query of MRCONSO.
    stys : `list` of `str`
        The semantic types of the CUI.
    """
    cuis.sort(key=lambda x: x.MRRANK_RANK, reverse=True)
    lead_aui = cuis[0]
    nauis = len(set([i.AUI for i in cuis]))
    stys = sorted(set([i.STY for i in cuis]))
    return (
        lead_aui.AUI,
        lead_aui.CUI,
        lead_aui.STR,
        lead_aui.LAT,
        lead_aui.SAB,
        lead_aui.TTY,
        lead_aui.CODE,
        lead_aui.MRRANK_RANK,
        nauis,
        nauis / data_len,
        stys,
    )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == "__main__":
    main()
