"""Perform UMLS code lookups to get all the CUIs that mat to a specific code.
This queries the ``CODE`` column in the MRCONSO table and then back queries for
all the matching CUIs for that code.
"""

from umls_tools import (
    __version__,
    __name__ as pkg_name,
    common,
    query,
    orm as o,
)
from pyaddons import utils, log
from sqlalchemy.orm import aliased
from sqlalchemy_config import config as cfg
from tqdm import tqdm
import stdopen
import argparse
import sys
import csv
# import pprint as pp

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

_PROG_NAME = "umls-code-lookup"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""
_INTERNAL_DELIM = "|"
"""The internal delimiter value (`list` of `str`)
"""

ROW_IDX_COL_NAME = "ROW_IDX"
"""The name for the input row index column in the output file (`str`)
"""
QUERY_TERM_COL_NAME = "QUERY_CODE"
"""The name for the query term column in the output file (`str`)
"""
CUI_COL_NAME = "CUI"
"""The name for the concept ID column in the output file (`str`)
"""
AUI_COL_NAME = "AUI"
"""The name for the atom ID column in the output file (`str`)
"""
LAT_COL_NAME = "LAT"
"""The name for the language column in the output file (`str`)
"""
SAB_COL_NAME = "SAB"
"""The name for the source abbreviation column in the output file (`str`)
"""
CODE_COL_NAME = "CODE"
"""The name for the term code column in the output file (`str`)
"""
STR_COL_NAME = "STR"
"""The name for the term name column in the output file (`str`)
"""
TTY_COL_NAME = "TTY"
"""The name for the term type column in the output file (`str`)
"""
RANK_COL_NAME = "RANK"
"""The name for the concept rank column in the output file (`str`)
"""
STY_COL_NAME = "STY"
"""The name for the semantic type column in the output file (`str`)
"""

OUTPUT_HEADER = [
    ROW_IDX_COL_NAME,
    QUERY_TERM_COL_NAME,
    CUI_COL_NAME,
    AUI_COL_NAME,
    LAT_COL_NAME,
    SAB_COL_NAME,
    CODE_COL_NAME,
    STR_COL_NAME,
    TTY_COL_NAME,
    RANK_COL_NAME,
    STY_COL_NAME,
]
"""The output column names added by the code mapper (`list` of `str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():  # pragma: no cover
    """The main entry point."""
    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)

    # See if we have progress verbosity on
    prog_verbose = log.progress_verbose(args.verbose)

    # 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,
    )

    utils.test_same_files(
        args.infile, args.outfile, error_msg="infile and outfile and the same"
    )

    total = 0
    matches = 0

    query_func = match_all_codes
    if args.top_rank is True:
        query_func = match_top_rank_codes

    try:
        inopen_method = utils.get_open_method(args.infile)
    except TypeError:
        inopen_method = open

    try:
        outopen_method = utils.get_open_method(args.outfile)
    except TypeError:
        outopen_method = open

    with stdopen.open(
        args.infile, method=inopen_method, encoding=args.encoding
    ) as infile:
        reader = csv.DictReader(infile, delimiter=args.delimiter)
        with stdopen.open(
            args.outfile,
            mode="wt",
            method=outopen_method,
            use_tmp=True,
            tmpdir=args.tmpdir,
        ) as outfile:
            session = sm()

            # A rank lookup table
            ranks = query.get_ranks(session)

            try:
                first_row = next(reader)
                total += 1

                header = list(first_row.keys()) + OUTPUT_HEADER
                writer = csv.DictWriter(
                    outfile, fieldnames=header, delimiter=args.delimiter
                )
                writer.writeheader()

                tqdm_kwargs = dict(
                    unit=" rows",
                    desc="[info] mapping codes",
                    disable=not prog_verbose,
                )
                for row in tqdm(reader, **tqdm_kwargs):
                    found = False
                    sty_lookup = dict()
                    for code in query_func(
                        session, row[args.code_column], lats=args.lat
                    ):
                        rank = ranks[(code.SAB, code.TTY)]
                        try:
                            sty = sty_lookup[code.CUI]
                        except KeyError:
                            try:
                                sty_lookup[code.CUI] = [i.STY for i in code.STY]
                            except AttributeError:
                                # The Unique CUI table has different column
                                # names
                                sty_lookup[code.CUI] = (
                                    code.SEMANTIC_TYPES.split(_INTERNAL_DELIM)
                                )
                            sty = sty_lookup[code.CUI]

                        writer.writerow(
                            create_match_row(
                                row,
                                total,
                                row[args.code_column],
                                code,
                                rank[0],
                                sty,
                            )
                        )
                        found = True
                    if found is False:
                        writer.writerow(
                            create_null_row(row, total, row[args.code_column])
                        )
                    matches += found
                    total += 1

                logger.info(f"total tested: {total}")
                logger.info(
                    f"matched: {matches} ({round((matches / total) * 100, 2)}%)"
                )
            except KeyError as e:
                raise KeyError(f"column not found: {e.args[0]}") from e
            except StopIteration as e:
                raise ValueError("no data in input file") from e
            finally:
                session.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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. Compressed input files are allowed.",
    )
    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",
        "--code-column",
        type=str,
        default="code",
        help="The column name containing the code definition.",
    )
    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 CODE/CUI match.",
    )
    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 location of tmp, if not provided will use the system tmp.",
    )
    parser.add_argument(
        "-E",
        "--encoding",
        type=str,
        default="utf-8",
        help="The encoding of the input file.",
    )
    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()
    args.infile = utils.get_full_path(args.infile, allow_none=True)
    args.outfile = utils.get_full_path(args.outfile, allow_none=True)
    args.config = utils.get_full_path(args.config, allow_none=True)
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def match_all_codes(session, code, lats=None):
    """Search the UMLS for the code and return all CUIs that match the CUIs
    that matches the code.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    code : `str`
        The code to map.
    lats : `list` of `str` or None, optional, default: None
        The restrict to AUIs with any of the specific languages.

    Yields
    ------
    match : `umls_tools.orm.MrConso`
        One of the AUIs representing the a CUI of the AUI that matched the
        code.
    """
    if lats is None:
        lats = ["ENG"]

    code_alias = aliased(o.MrConso)
    cui_alias = aliased(o.MrConso)
    q = (
        session.query(cui_alias)
        .filter(code_alias.CODE == code)
        .join(code_alias, cui_alias.CUI == code_alias.CUI)
        .group_by(cui_alias.MRCONSO_IDX)
    )
    for row in q:
        if row.LAT in lats:
            yield row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def match_top_rank_codes(session, code, lats=None):
    """Match a code against the MRCONSO table and extract the top ranked CUI
    that matches the code.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    code : `str`
        The code to map.
    lats : `list` of `str` or None, optional, default: None
        The restrict to AUIs with any of the specific languages.

    Yields
    ------
    match : `umls_tools.orm.MrConso`
        One of the AUIs representing the a CUI of the AUI that matched the
        code.
    """
    if lats is None:
        lats = ["ENG"]

    q = (
        session.query(o.UniqueUmlsCui)
        .filter(o.MrConso.CODE == code)
        .join(o.MrConso, o.UniqueUmlsCui.CUI == o.MrConso.CUI)
        .group_by(o.MrConso.MRCONSO_IDX)
    )
    for row in q:
        if row.LAT in lats:
            yield row


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

    Parameters
    ----------
    in_row : `dict`
        The input row that has not been matched.
    idx : `int`
        The input row number.
    query_code : `str`
        The actual query code that could not be matched.

    Returns
    -------
    output_row : `dict`
        An unmapped row that can be written to file.
    """
    out_row = in_row.copy()

    for i in OUTPUT_HEADER:
        out_row[i] = None
    out_row[ROW_IDX_COL_NAME] = idx
    out_row[QUERY_TERM_COL_NAME] = query_code
    return out_row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_match_row(in_row, idx, query_code, match, rank, sty):
    """Create an output row from a matching code.

    Parameters
    ----------
    in_row : `dict`
        The input row that has not been matched.
    idx : `int`
        The input row number.
    query_code : `str`
        The actual query code that could not be matched.
    match : `umls_orm.MrConso`
        The matching code row.
    rank : `int`
        The UMLS sab/tty rank for the match.
    sty : `list` or `str`
        The semantic types of the match.

    Returns
    -------
    output_row : `dict`
        An mapped row that can be written to file.
    """
    out_row = in_row.copy()
    out_row[ROW_IDX_COL_NAME] = idx
    out_row[QUERY_TERM_COL_NAME] = query_code
    out_row[CUI_COL_NAME] = match.CUI
    out_row[AUI_COL_NAME] = match.AUI
    out_row[LAT_COL_NAME] = match.LAT
    out_row[SAB_COL_NAME] = match.SAB
    out_row[CODE_COL_NAME] = match.CODE
    out_row[STR_COL_NAME] = match.STR
    out_row[TTY_COL_NAME] = match.TTY
    out_row[RANK_COL_NAME] = rank
    out_row[STY_COL_NAME] = _INTERNAL_DELIM.join(sty)

    return out_row


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