"""Map text strings into the UMLS. The query processes is as follows.

1. Query the term for an exact (case insensitive) match, use optional LAT
   filtering.
2. If step one turns out nothing then query for a substring match.
3. If step two gives nothing and Metamap is not available then search for a
   tokenised word match (all words must match but order is irrelevant).
4. If Metamap is available (if the --metamap flag is in use), the use that to
   perform the search instead of the tokeniser.
5. The UMLS CUI identified from either steps 1-4 is back queried to get all
   other concepts that map to it or the highest ranking concept.

Entries that do not match steps 1-4 and 7 are returned as blank records. So
the output file will have at least the same number of rows as the input file.

The performance of this script is slow approx 2 rows/second using an SQLite
 back end on a mechanical enterprise HDD. I aim to improve this in future.

This script requires the UMLS database to be built with the ``UNIQUE_CUI``
table added, please see `the command line scripts
<https://cfinan.gitlab.io/umls-tools/scripts.html>`_.
"""

from umls_tools import __version__, __name__ as pkg_name, log, common, query
from umls_tools.metamap import metamap
from sqlalchemy_config import config as cfg
from tqdm import tqdm
import stdopen
import argparse
import os
import sys
import csv
import gzip
# import pprint as pp

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

_PROG_NAME = "umls-mapper"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""
_META_ARGS = [
    "-L",
    "{0}",
    "-Z",
    "{0}",
    "-sAIy",
    "--negex",
    "--conj",
    "-V",
    "USAbase",
]
"""The metamap arguments to use (`list` of `str`)
"""
_INTERNAL_DELIM = ["|", "/"]
"""The internal delimiter values (`list` of `str`)
"""


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

    verbose = False
    if int(args.verbose) > 1:
        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,
    )

    metamap_server = None
    if args.metamap is not None:
        metamap_args = [i.format(args.metamap) for i in _META_ARGS]
        metamap_server = metamap.MetamapServer(
            keep_alive=True, binary="metamap", args=metamap_args
        )
        metamap_server.start()

    try:
        total = 0
        match_type_count = {
            query.EXACT_MATCH: 0,
            query.SUBSTRING_MATCH: 0,
            query.TOKEN_STRING_MATCH: 0,
            query.METAMAP_MATCH: 0,
            query.NO_MATCH: 0,
        }
        match_type_desc = {
            query.EXACT_MATCH: "Exact ({0}):             {1} ({2}%)",
            query.SUBSTRING_MATCH: "Substring mapping ({0}): {1} ({2}%)",
            query.TOKEN_STRING_MATCH: "Token mapping ({0}):     {1} ({2}%)",
            query.METAMAP_MATCH: "Metamap mapping ({0}):   {1} ({2}%)",
            query.NO_MATCH: "No mapping ({0}):        {1} ({2}%)",
        }

        try:
            sty_delim = [i for i in _INTERNAL_DELIM if args.delimiter != i][0]
        except IndexError as e:
            raise ValueError("delimiter conflict") from e

        # Check to see if we need to output compressed
        write_method = open
        if args.outfile is not None and args.outfile.endswith(".gz"):
            write_method = gzip.open

        with stdopen.open(
            args.outfile,
            "wt",
            tmpdir=args.tmp,
            use_tmp=True,
            method=write_method,
        ) as outfile:
            # Set up the mapper and the output file
            map_kwargs = dict(
                delimiter=args.delimiter,
                verbose=verbose,
                encoding=args.encoding,
                map_columns=args.map_columns,
                lats=args.lat,
                top_rank=args.top_rank,
                metamap_server=metamap_server,
            )
            mapper = map_file(sm(), args.infile, **map_kwargs)
            header, first_row = _get_header(mapper)
            prev_idx = first_row["ROW_IDX"]
            writer = csv.DictWriter(outfile, header, delimiter=args.delimiter)
            writer.writeheader()

            try:
                first_row["STY"] = sty_delim.join(first_row["STY"])
            except TypeError:
                pass

            writer.writerow(first_row)
            for row in mapper:
                # Multiple mappings can be returned foreach input row
                if prev_idx != row["ROW_IDX"]:
                    match_type_count[row["MATCH_TYPE"]] += 1
                    total += 1
                try:
                    row["STY"] = sty_delim.join(row["STY"])
                except TypeError:
                    pass
                writer.writerow(row)
                prev_idx = row["ROW_IDX"]

        logger.info(f"total:                 {total}")
        for k, count in match_type_count.items():
            desc = match_type_desc[k]
            logger.info(desc.format(k, count, round((count / total) * 100, 2)))

        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 metamap_server is not None:
            metamap_server.stop()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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(
        "-i",
        "--infile",
        type=str,
        help="The input file to map, if not provided then "
        "input is expected from STDIN.",
    )
    parser.add_argument(
        "-C",
        "--map-columns",
        type=str,
        nargs="+",
        default=["term"],
        help="The columns in the input file with terms that we want to map.",
    )
    parser.add_argument(
        "-o",
        "--outfile",
        type=str,
        help="The output file to map, if not provided then "
        "output is to STDOUT. If the output file has a .gz "
        "extension then the output will be compressed.",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default=common.get_default_config_cli_path(),
        help="The location of the config file.",
    )
    parser.add_argument(
        "-L",
        "--lat",
        nargs="+",
        type=str,
        default=["ENG"],
        help="The language of the terms to return.",
    )
    parser.add_argument(
        "-t",
        "--top-rank",
        action="store_true",
        help="Return only the top ranking AUI for each CUI match.",
    )
    parser.add_argument(
        "-m",
        "--metamap",
        type=str,
        help="Use this metamap version (must be installed locally), if not"
        " defined then Metamap is not used.",
    )
    parser.add_argument(
        "-E",
        "--encoding",
        type=str,
        default="utf-8",
        help="The encoding of the input 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",
        "--tmp",
        type=str,
        help="The location of tmp, if not provided will use the system tmp.",
    )
    parser.add_argument(
        "-d",
        "--delimiter",
        type=str,
        default="\t",
        help="The location of tmp, if not provided will use the system tmp.",
    )
    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 _get_header(map_file_gen):
    """Get the header from the first row of a map file.

    Parameters
    ----------
    map_file_gen : `generator`
        A generator from the ``map_file`` function with no rows extracted.
    """
    first_row = next(map_file_gen)
    return list(first_row.keys()), first_row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def map_file(
    session,
    infile,
    map_columns=None,
    delimiter="\t",
    verbose=False,
    encoding="utf-8",
    lats=None,
    top_rank=False,
    metamap_server=None,
):
    """Map a file of text terms against the UMLS. This looks for exact matches
    only.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    infile : `str`
        The path to the input file containing the query terms.
    map_columns : `list` of `str`, optional, default: `NoneType`
        The column names in the input file that contain terms that you want to
        map against the UMLS. If ``None`` this defaults to ``['term']``.
    delimiter : `str`, optional, default: `\t`
        The delimiter of the input file.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    encoding : `str`, optional, default: `utf-8`
        The encoding of the input file.
    lats : `list` of `str`, optional, default: `NoneType`
        The restrict to AUIs with any of the specific languages. If ``None``
        this defaults to ``['ENG']``.
    top_rank : `bool`, optional, default: `False`
        Restrict matches to top ranked AUIs.
    metamap_server : `umls_tools.metamap.MetamapServer`, optional, default: `NoneType`
        A server object to issue queries against. If ``NoneType`` then  Metamap
        will not be used. This assumes that the metamap arguments were given
        when the server object was created.

    Yields
    ------
    mapped_row : `dict`
        The mappings for each input row. If there are no mappings then a blank
        mapping is yielded. i.e. all input rows will be yielded regardless of
        if they have a mapping against a UMLS term.

    See also
    --------
    umls_tools.metamap.MetamapServer
    """
    if map_columns is None:
        map_columns = ["term"]
    if lats is None:
        lats = ["ENG"]

    # Get a rank lookup dictionary
    ranks = query.get_ranks(session)

    mapper = query.CuiMapper(session, lats=lats, metamap_server=metamap_server)
    with stdopen.open(infile, "rt", encoding=encoding) as incsv:
        reader = csv.DictReader(incsv, delimiter=delimiter)
        for idx, row in enumerate(
            tqdm(
                reader,
                disable=not verbose,
                desc="[info] mapping file",
                unit=" rows",
            )
        ):
            match = False
            # Each row of the input file can have multiple columns to look
            # in
            for i in map_columns:
                try:
                    term = row[i]
                except KeyError as e:
                    raise KeyError(f"can't find map-column: {i}") from e
                # If a term has been defined
                if term is not None and term != "":
                    results, match_type = mapper.map_term(term)

                    if len(results) > 0:
                        match = True

                    if top_rank is False:
                        for r, start, end in results:
                            return_row = create_row(
                                idx, row, i, term, r.STR, match_type, start, end
                            )
                            return_row["CUI"] = r.CUI
                            return_row["AUI"] = r.AUI
                            return_row["LAT"] = r.LAT
                            return_row["SAB"] = r.SAB
                            return_row["CODE"] = r.CODE
                            return_row["STR"] = r.STR
                            return_row["TTY"] = r.TTY
                            return_row["RANK"] = ranks[(r.SAB, r.TTY)][0]
                            return_row["STY"] = [i.STY for i in r.STY]
                            yield return_row
                    else:
                        seen_cui = set()
                        for r, start, end in results:
                            if r.CUI not in seen_cui:
                                return_row = create_row(
                                    idx,
                                    row,
                                    i,
                                    term,
                                    r.STR,
                                    match_type,
                                    start,
                                    end,
                                )

                                u = query.get_unique_cui(session, r.CUI)
                                return_row["CUI"] = u.CUI
                                return_row["AUI"] = u.AUI
                                return_row["LAT"] = u.LAT
                                return_row["SAB"] = u.SAB
                                return_row["CODE"] = u.CODE
                                return_row["STR"] = u.STR
                                return_row["TTY"] = u.TTY
                                return_row["RANK"] = ranks[(u.SAB, u.TTY)][0]
                                return_row["STY"] = u.SEMANTIC_TYPES.split("|")
                                seen_cui.add(r.CUI)
                                yield return_row

                if match is False:
                    yield create_null_row(idx, row, i, term)
                else:
                    break


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_null_row(idx, row, query_column, query_term):
    """Create an empty no match row.

    Parameters
    ----------
    idx : `int`
        The input row number.
    row : `dict`
        The input row.
    column : `str`
        The column name in the input row that we are getting the query terms
        from.
    query_term : `str`
        The actual query term.

    Returns
    -------
    output_row : `dict`
        An unmapped row that can be written to file.
    """
    return_row = create_row(
        idx, row, query_column, query_term, None, query.NO_MATCH, None, None
    )
    return_row["CUI"] = None
    return_row["AUI"] = None
    return_row["LAT"] = None
    return_row["SAB"] = None
    return_row["CODE"] = None
    return_row["STR"] = None
    return_row["TTY"] = None
    return_row["RANK"] = None
    return_row["STY"] = None
    return return_row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_row(
    idx, row, query_column, query_term, match_term, match_type, start, end
):
    """Create an output row.

    Parameters
    ----------
    idx : `int`
        The input row number.
    row : `dict`
        The input row.
    query_column : `str`
        The column name in the input row that we are getting the query terms
        from.
    query_term : `str`
        The actual query term.
    match_term : `str`
        The UMLS CUI term that matched with the query string.
    match_type : `str`
        the match type, so exact (``E``), substring (``S``), tokenised (``T``)
    start : `int`
        The start position of the match string within the query_term.
    end : `int`
        The end position of the match string within the query_term.

    Returns
    -------
    output_row : `dict`
        An unmapped row that can be written to file.
    """
    return_row = row.copy()
    return_row["ROW_IDX"] = idx
    return_row["QUERY_COL"] = query_column
    return_row["QUERY_STR"] = query_term
    return_row["MATCH_TERM"] = match_term
    return_row["MATCH_TYPE"] = match_type
    return_row["MATCH_START"] = start
    return_row["MATCH_END"] = end

    return return_row


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