"""Functions for getting path distances between sets of source vocab codes
potentially from different source vocabs"""

from umls_tools import __version__, __name__ as pkg_name, log, common, orm as o
from sqlalchemy import and_
from sqlalchemy.orm import aliased
from sqlalchemy_config import config as cfg
from tqdm import tqdm
import argparse
import os
import sys
import csv
import re
# import pprint as pp

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

_PROG_NAME = "umls-neo4j-extract"
"""The name of the program that doubles as the logger name (`str`)
"""
_DESC = __doc__
"""The program description given to argparse (`str`)
"""
_DELIMITER = "\t"
"""The output file delimiter value (`str`)
"""
_INTERNAL_DELIMITER = "|"
"""The internal delimiter values for neo4j arrays (`str`)
"""
SOURCE_CUI_REL_NAME = "has_concept"
"""The relationship name for the source->cui relationship (`str`)
"""
ENGLISH_LANGUAGE = "ENG"
"""Constant for the English language in the UMLS (`str`)
"""
NOCODE = "NOCODE"
"""Constant for the NOCODE text in the UMLS (`str`)
"""

# File names and headers
AUI_NODE_BASENAME = "aui_nodes.txt"
"""File basename of the AUI nodes (`str`)
"""
AUI_NODE_HEADER = ["aui:ID", "term:string", "sab:string", ":LABEL"]
"""The Neo4j compatible file header for the AUI node file (`list` of `str`)
"""

CUI_NODE_BASENAME = "cui_nodes.txt"
"""File basename of the CUI nodes (`str`)
"""
CUI_NODE_HEADER = [
    "cui:ID",
    "term:string",
    "naui:long",
    "pcui:double",
    ":LABEL",
]
"""The Neo4j compatible file header for the CUI node file (`list` of `str`)
"""

SOURCE_NODE_BASENAME = "source_nodes.txt"
"""File basename of the source vocab nodes (`str`)
"""
SOURCE_NODE_HEADER = ["code:ID", "term:string[]", ":LABEL"]
"""The Neo4j compatible file header for the source vocab node file
(`list` of `str`)
"""

AUI_REL_BASENAME = "aui_rel.txt"
"""File basename of the AUI->AUI relationship edges (`str`)
"""
AUI_REL_HEADER = [":START_ID", ":END_ID", ":TYPE"]
"""The Neo4j compatible file header for the AUI->AUI relationships file
(`list` of `str`)
"""

SOURCE_REL_BASENAME = "source_rel.txt"
"""File basename of the source->source relationship edges (`str`)
"""
SOURCE_REL_HEADER = [":START_ID", ":END_ID", ":TYPE"]
"""The Neo4j compatible file header for the source->source relationships file
(`list` of `str`)
"""

SOURCE_CUI_REL_BASENAME = "source_cui_rel.txt"
"""File basename of the source->cui relationship edges (`str`)
"""
SOURCE_CUI_REL_HEADER = [":START_ID", ":END_ID", ":TYPE"]
"""The Neo4j compatible file header for the source->source relationships file
(`list` of `str`)
"""

AUI_CUI_REL_BASENAME = "aui_cui_rel.txt"
"""File basename of the aui->cui relationship edges (`str`)
"""
AUI_CUI_REL_HEADER = [":START_ID", ":END_ID", ":TYPE"]
"""The Neo4j compatible file header for the aui->cui relationships file
(`list` of `str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """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,
    )

    try:
        build_files(sm(), args.outdir, verbose=verbose)
        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():
    """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(
        "-o",
        "--outdir",
        type=str,
        help="The output directory for the flat files. The"
        " default it the current dir",
    )
    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(
        "-v",
        "--verbose",
        default=0,
        action="count",
        help="give more output, ``--vv`` turns on progress monitoring.",
    )
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):
    """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)
    args.outdir = args.outdir or os.getcwd()
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_files(session, outdir, verbose=False):
    """Build flat files from UMLS, source terms and concepts (nodes) and their
    relationships (edges).

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    tqdm_kwargs = dict(
        unit=" files", desc="[info] processing files", disable=not verbose
    )

    build_funcs = [
        build_aui_nodes,
        build_cui_nodes,
        build_aui_rel,
        build_cui_aui_rel,
    ]
    for f in tqdm(build_funcs, **tqdm_kwargs):
        f(session, outdir, verbose=verbose)

    # build_aui_nodes(session, outdir, verbose=verbose)
    # build_cui_nodes(session, outdir, verbose=verbose)
    # build_aui_rel(session, outdir, verbose=verbose)
    # build_cui_aui_rel(session, outdir, verbose=verbose)
    # build_source_nodes(session, outdir, verbose=verbose)
    # build_source_rel(session, outdir, verbose=verbose)
    # build_cui_source_rel(session, outdir, verbose=verbose)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_aui_nodes(session, outdir, verbose=False):
    """Build a flat file of the UMLS aui nodes.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    sql = session.query(
        o.MrConso.AUI, o.MrConso.STR, o.MrConso.CODE, o.MrConso.SAB
    ).filter(and_(o.MrConso.LAT == ENGLISH_LANGUAGE, o.MrConso.CODE != NOCODE))

    count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract AUI nodes",
        leave=False,
    )

    with open(os.path.join(outdir, AUI_NODE_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(AUI_NODE_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            outcsv.write(
                "{0}\n".format(
                    _DELIMITER.join(
                        [
                            row.AUI,
                            quote(row.STR),
                            quote(row.CODE),
                            f"{row.SAB};Atom",
                        ]
                    )
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_cui_nodes(session, outdir, verbose=False):
    """Build a flat file from the UMLS concepts.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    sql = session.query(
        o.UniqueUmlsCui.CUI,
        o.UniqueUmlsCui.STR,
        o.UniqueUmlsCui.NAUI,
        o.UniqueUmlsCui.PROB_CUI,
    ).filter(o.UniqueUmlsCui.LAT == ENGLISH_LANGUAGE)
    count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract CUI nodes",
        leave=False,
    )

    with open(os.path.join(outdir, CUI_NODE_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(CUI_NODE_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            outcsv.write(
                "{0}\n".format(
                    _DELIMITER.join(
                        [
                            row.CUI,
                            quote(row.STR),
                            str(row.NAUI),
                            str(row.PROB_CUI),
                            "Concept",
                        ]
                    )
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_source_nodes(session, outdir, verbose=False):
    """Build a flat file for UMLS source vocabulary codes.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    # SELECT CODE, SAB, STR
    # FROM MRCONSO m
    # WHERE CODE != "NOCODE"
    # ORDER BY SAB, CODE
    sql = session.query(o.MrConso.CODE, o.MrConso.SAB, o.MrConso.STR).filter(
        and_(o.MrConso.LAT == ENGLISH_LANGUAGE, o.MrConso.CODE != NOCODE)
    )

    count = sql.count()
    sql = sql.order_by(o.MrConso.SAB, o.MrConso.CODE)

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract source nodes",
        leave=False,
    )

    with open(os.path.join(outdir, SOURCE_NODE_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(SOURCE_NODE_HEADER)))
        row_buffer = []
        cur_sab = None
        cur_code = None
        for row in tqdm(sql, **tqdm_kwargs):
            if cur_sab != row.SAB or cur_code != row.CODE:
                if len(row_buffer) > 0:
                    str_array = _INTERNAL_DELIMITER.join(
                        sorted(set([i.STR for i in row_buffer]))
                    )
                    outcsv.write(
                        "{0}\n".format(
                            _DELIMITER.join(
                                [cur_code, quote(str_array), cur_sab]
                            )
                        )
                    )
                row_buffer = []
                cur_sab = row.SAB
                cur_code = row.CODE
            row_buffer.append(row)
        if len(row_buffer) > 0:
            str_array = _INTERNAL_DELIMITER.join(
                sorted(set([i.STR for i in row_buffer]))
            )
            outcsv.write(
                "{0}\n".format(
                    _DELIMITER.join(
                        [
                            i
                            for i in [cur_code, quote(str_array), cur_sab]
                            if i != ""
                        ]
                    )
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_cui_aui_rel(session, outdir, verbose=False):
    """Build a flat file for the UMLS aui->cui relationships.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    # SELECT CUI, CODE
    # FROM MRCONSO m
    # GROUP BY CUI, CODE
    sql = (
        session.query(o.MrConso.AUI, o.MrConso.CUI)
        .filter(o.MrConso.LAT == ENGLISH_LANGUAGE, o.MrConso.CODE != NOCODE)
        .group_by(o.MrConso.CODE, o.MrConso.CUI)
    )
    count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract source->cui relationships",
        leave=False,
    )

    with open(os.path.join(outdir, AUI_CUI_REL_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(AUI_CUI_REL_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            outcsv.write(
                "{0}\n".format(
                    _DELIMITER.join([row.AUI, row.CUI, SOURCE_CUI_REL_NAME])
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_aui_rel(session, outdir, verbose=False):
    """Build a flat file from the UMLS AUI->AUI relationships.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    mrcons1 = aliased(o.MrConso)
    mrcons2 = aliased(o.MrConso)

    sql = (
        session.query(o.MrRel.AUI1, o.MrRel.AUI2, o.MrRel.REL)
        .join(mrcons1, mrcons1.AUI == o.MrRel.AUI1)
        .join(mrcons2, mrcons2.AUI == o.MrRel.AUI2)
        .filter(
            and_(
                o.MrRel.AUI1.isnot(None),
                o.MrRel.AUI2.isnot(None),
                o.MrRel.AUI1 != o.MrRel.AUI2,
                mrcons1.LAT == ENGLISH_LANGUAGE,
                mrcons2.LAT == ENGLISH_LANGUAGE,
                mrcons1.CODE != NOCODE,
                mrcons2.CODE != NOCODE,
            )
        )
    )

    count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract AUI->AUI relationships",
        leave=False,
    )

    with open(os.path.join(outdir, AUI_REL_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(AUI_REL_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            outcsv.write(
                "{0}\n".format(_DELIMITER.join([row.AUI1, row.AUI2, row.REL]))
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_source_rel(session, outdir, verbose=False):
    """Build a flat file from the UMLS source->source relationships. This may
    take some time to run.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    # SELECT m2.CODE, m3.CODE, REL
    # FROM MRREL m
    # JOIN MRCONSO m2
    # ON m.AUI1 = m2.AUI
    # JOIN MRCONSO m3
    # ON AUI2  = m3.AUI
    # GROUP BY m2.CODE, m3.CODE, REL
    mrcons1 = aliased(o.MrConso)
    mrcons2 = aliased(o.MrConso)

    sql = (
        session.query(
            mrcons1.CODE.label("m1code"),
            mrcons2.CODE.label("m2code"),
            o.MrRel.REL,
        )
        .join(mrcons1, mrcons1.AUI == o.MrRel.AUI1)
        .join(mrcons2, mrcons2.AUI == o.MrRel.AUI2)
        .filter(
            and_(
                mrcons1.LAT == ENGLISH_LANGUAGE,
                mrcons2.LAT == ENGLISH_LANGUAGE,
            )
        )
        .group_by(mrcons1.CODE, mrcons2.CODE, o.MrRel.REL)
    )
    count = sql.count()
    # sql = sql.group_by(mrcons1.CODE, mrcons2.CODE, o.MrRel.REL)

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract source->source relationships",
        leave=False,
    )

    with open(os.path.join(outdir, SOURCE_REL_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(SOURCE_REL_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            if (
                row.m1code != row.m2code
                and row.m1code is not None
                and row.m2code is not None
                and row.m1code != NOCODE
                and row.m2code != NOCODE
            ):
                outcsv.write(
                    "{0}\n".format(
                        _DELIMITER.join([row.m1code, row.m2code, row.REL])
                    )
                )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_cui_source_rel(session, outdir, verbose=False):
    """Build a flat file from the UMLS source->cui relationships. This may
    take some time to run.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outdir : `str`
        The path to the output directory to write the flat files.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    # SELECT CUI, CODE
    # FROM MRCONSO m
    # GROUP BY CUI, CODE
    sql = (
        session.query(o.MrConso.CODE, o.MrConso.CUI)
        .filter(o.MrConso.LAT == ENGLISH_LANGUAGE, o.MrConso.CODE != NOCODE)
        .group_by(o.MrConso.CODE, o.MrConso.CUI)
    )
    count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        total=count,
        disable=not verbose,
        desc="[info] extract source->cui relationships",
        leave=False,
    )

    with open(os.path.join(outdir, SOURCE_CUI_REL_BASENAME), "wt") as outcsv:
        outcsv.write("{0}\n".format(_DELIMITER.join(SOURCE_CUI_REL_HEADER)))
        for row in tqdm(sql, **tqdm_kwargs):
            outcsv.write(
                "{0}\n".format(
                    _DELIMITER.join([row.CODE, row.CUI, SOURCE_CUI_REL_NAME])
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def quote(text):
    """Wrap the text in quotes and translate internal quotes to single quotes.
    Also remove an internal newline characters.

    Parameters
    ----------
    text : `str`
        The text to quote and escape.

    Returns
    -------
    quoted_text : `str`
        The quoted and escaped text.
    """
    text = re.sub(r"[\r\n]", "", text)
    text = re.sub(r'"', "'", text)
    return f'"{text}"'


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