"""Extract node/relationship flat files to import into a neo4j graph. These
are for use with the Neo4j admin `import tool <https://neo4j.com/docs/operations-manual/current/tools/neo4j-admin/neo4j-admin-import/>`_

This may take a while to run, it may also produce some SQLAlchemy warnings.
This should but impact anything.

Please note, there is a bug that prevents the root concepts for hierarchical
SABs being linked properly, I hope to fix this in the near future.

Also see these URLs:
https://neo4j.com/developer/guide-import-csv/#batch-importer
https://neo4j.com/docs/operations-manual/current/tutorial/neo4j-admin-import/
"""

from umls_tools import __version__, __name__ as pkg_name, common, orm as o
from pyaddons import utils, log
from sqlalchemy import and_
from sqlalchemy.orm import aliased
from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
from sqlalchemy_config import config as cfg
from tqdm import tqdm
from collections import namedtuple
from itertools import combinations
import argparse
import os
import sys
import csv
import re
import shelve
# import pprint as pp


Sab = namedtuple("Sab", ["name", "id_name", "file_suffix", "code"])
"""named tuple to hold the raw/processed source abbreviation data
(`namedtuple`)
"""

# 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`)
"""


# Relationship names i.e. :TYPE
SOURCE_CUI_REL_NAME = "has_cui"
"""The relationship name for the source->cui relationship (`str`)
"""
CUI_SOURCE_REL_NAME = "has_source"
"""The relationship name for the cui->source relationship (`str`)
"""
SOURCE_MAPPING_REL_NAME = "cui_mapping"
"""The relationship name for the source->source mapping relationship
 that is a relationship that is mediated by a cui (although not on the
 graph)(`str`)
"""
PAR_REL = "PAR"
"""The UMLS parent relationship abbreviation (`str`)
"""
CHD_REL = "CHD"
"""The UMLS child relationship abbreviation (`str`)
"""

# Labels i.e. :LABEL
SOURCE_LABEL = "Source"
"""The label for the source codes (`str`)
"""
CONCEPT_LABEL = "Concept"
"""The label for the CUI codes (`str`)
"""
SOURCE_ROOT_LABEL = "SourceRoot"
"""The label for the source root (`str`)
"""
CONCEPT_ROOT_LABEL = "ConceptRoot"
"""The label for the concept root (`str`)
"""
UMLS_ROOT_LABEL = "UMLS_ROOT"
"""The label for the UMLS root term (`str`)
"""

# Constants for some SABs
CUI_TYPE = "CUI"
"""The SAB of the 'concept' (`str`)
"""
UMLS_SAB = "UMLS"
"""The SAB for the hierarchical root UMLS code (`str`)
"""

# Some other constant values using in UMLS database queries
ENGLISH_LANGUAGE = "ENG"
"""Constant for the English language in the UMLS (`str`)
"""
NOCODE = "NOCODE"
"""Constant for the NOCODE text in the UMLS (`str`)
"""
ROOT_TTY = "RHT"
"""The term type for a root node (`str`)
"""

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

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

# The source mapping relationships file
SOURCE_MAPPING_BASENAME = "source_mapping_rel.txt"
"""File basename of the source->source mapping edges (`str`)
"""
SOURCE_CUI_REL_BASENAME = "source_cui_rel.txt"
"""File basename of the source->cui relationship edges (`str`)
"""

# The source hierarchical relationships file
SOURCE_REL_BASENAME = "source_rel.txt"
"""File basename of the source->source relationship edges (`str`)
"""

# The concept (CUI) hierarchical relationships file
CUI_REL_BASENAME = "cui_rel.txt"
"""File basename of the source->source relationship edges (`str`)
"""

# A general header for a relationships file
GENERAL_REL_HEADER = [":START_ID", ":END_ID", ":TYPE"]
"""The Neo4j compatible file header a general relationships file
(`list` of `str`)
"""


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class IdTracker(object):
    """Track mappings of nodes+properties to numeric IDs.

    Parameters
    ----------
    *args
        Ignored.
    **kargs
        Ignored.
    """

    ID = 1
    """Class wide node ID variable, this shared by all objects (`int`)
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.nodes = {}

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __len__(self):
        """Get the number of mappings held (`int`)"""
        return len(self.nodes)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Does nothing, for interface.

        Returns
        -------
        self : `IdMapper`
            For chaining.
        """
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Does nothing, for interface."""
        pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_node(self, sab, code, term, label):
        """Add a node to the ID tracker.

        Parameters
        ----------
        sab : `str`
            The source abbreviation of the node.
        code : `str`
            The code (source ID) of the node.
        term : `str`
            The description of the node.
        label : `str`
            The label for the node.

        Returns
        -------
        node_id : `int`
            The ID that the node has been stored under.
        """
        try:
            node_id = self.nodes[(sab, code)][1]
            self.nodes[(sab, code)][0].add(term)
        except KeyError:
            self.nodes[(sab, code)] = [set([term]), self.ID, label]
            node_id = self.ID
            self.__class__.ID += 1
        return node_id

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_node_id(self, sab, code):
        """Retrieve a node Id from the tracker.

        Parameters
        ----------
        sab : `str`
            The source abbreviation of the node.
        code : `str`
            The code (source ID) of the node.

        Returns
        -------
        node_id : `int`
            The ID that the node has been stored under.
        """
        return self.nodes[(sab, code)][1]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def fetch_nodes(self):
        """Yield all the nodes as rows that can be written to file.

        Yields
        ------
        row : `list`
            The node list that can be written to file.
            ``[0]`` node_id, ``[1]`` code, ``[2]`` term ``[3]`` label.
        """
        for k, v in self.nodes.items():
            term_str = _INTERNAL_DELIMITER.join(sorted(v[0]))
            yield [
                str(v[1]),
                quote(k[1]),
                quote(term_str),
                f"{k[0]}{_INTERNAL_DELIMITER}{v[2]}",
            ]


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class IdTrackerShelve(IdTracker):
    """Track mappings of nodes+properties to numeric IDs.

    Parameters
    ----------
    db_file : `str`
        The file to store the shelve node store.

    Notes
    -----
    This implementation is similar to IdTracker but using shelve instead of
    dictionaries, so it is lower memory.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, db_file, **kwargs):
        self.dbfile = db_file
        self.store = None
        self.nodes = {}

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Open the shelve.

        Returns
        -------
        self : `IdMapperShelve`
            For chaining.
        """
        self.store = shelve.open(self.dbfile, flag="n")
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close the shelve."""
        self.store.close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_node(self, sab, code, term, label):
        """Add a node to the ID tracker.

        Parameters
        ----------
        sab : `str`
            The source abbreviation of the node.
        code : `str`
            The code (source ID) of the node.
        term : `str`
            The description of the node.
        label : `str`
            The label for the node.

        Returns
        -------
        node_id : `int`
            The ID that the node has been stored under.
        """
        k = f"{sab}||{code}"
        try:
            node_id = self.nodes[k][1]
            self.nodes[k][0].add(term)
        except KeyError:
            self.nodes[k] = [set([term]), self.ID, label]
            node_id = self.ID
            self.__class__.ID += 1
        return node_id

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_node_id(self, sab, code):
        """Retrieve a node Id from the tracker.

        Parameters
        ----------
        sab : `str`
            The source abbreviation of the node.
        code : `str`
            The code (source ID) of the node.

        Returns
        -------
        node_id : `int`
            The ID that the node has been stored under.
        """
        return self.nodes[f"{sab}||{code}"][1]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def fetch_nodes(self):
        """Yield all the nodes as rows that can be written to file.

        Yields
        ------
        row : `list`
            The node list that can be written to file.
            ``[0]`` node_id, ``[1]`` code, ``[2]`` term ``[3]`` label.
        """
        for k, v in self.nodes.items():
            k = re.split(r"||", k)
            term_str = _INTERNAL_DELIMITER.join(sorted(v[0]))
            yield [
                str(v[1]),
                quote(k[1]),
                quote(term_str),
                f"{k[0]}{_INTERNAL_DELIMITER}{v[2]}",
            ]


# This can be used to change the type of tracker class used.
TRACKER_CLASS = IdTracker
"""The class of Id tracker, the shelve tracker will use less memory but will
 be slower (`class` of `umls_tool.neo4j.extract.IdTracker` or
 `umls_tool.neo4j.extract.IdTrackerShelve`).
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point. For API access see
    ``umls_tools.neo4j.extract.build_files``
    """
    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:
        errors = build_files(
            sm(),
            args.outdir,
            tmpdir=args.tmpdir,
            verbose=verbose,
            sabs=args.sabs,
        )

        # Output any potential mapping errors, these will be relationships
        # where the node can't be found against a valid SAB
        if len(errors) > 0:
            logger.info("the following key errors were raised: ")
            for i in errors:
                logger.info("{5} ({4}): {0} ({1})->{2} ({3}) ".format(*i))

        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(
        "-T",
        "--tmpdir",
        type=str,
        help="The temp directory to use for the node ID cache. Defaults to the"
        " system temp location",
    )
    parser.add_argument(
        "--sabs",
        type=str,
        nargs="+",
        default=None,
        help="Only extract from these SABs",
    )
    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, init_id=1, tmpdir=None, sabs=None, 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.
    init_id : `int`, optional, default: `1`
        The starting ID for the nodes being added.
    tmpdir : `str`, optional, default: `NoneType`
        The temp location for the node ID store, this defaults to the system
        temp location.
    sabs : `list` of `umls_tools.neo4j.extract.Sab`, optional, \
    default: `NoneType`
        Specific source vocabularies to extract from. If not set this will
        default to all source vocabularies.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    logger = log.retrieve_logger(_PROG_NAME, verbose=verbose)

    # Initialise the class wide ID value
    TRACKER_CLASS.ID = init_id

    # Will eventually hold the tracker objects
    source_nodes = None
    cui_nodes = None

    # Will eventually hold the output file objects
    source_node_file = None
    cui_node_file = None
    source_map_rel = None
    source_cui_rel = None
    source_rel = None
    cui_rel = None

    try:
        # Open all the output files we will use
        source_node_file = init_source_nodes(outdir)
        cui_node_file = init_cui_nodes(outdir)
        source_map_rel = init_rel_file(outdir, SOURCE_MAPPING_BASENAME)
        source_cui_rel = init_rel_file(outdir, SOURCE_CUI_REL_BASENAME)
        source_rel = init_rel_file(outdir, SOURCE_REL_BASENAME)
        cui_rel = init_rel_file(outdir, CUI_REL_BASENAME)

        # We create a dummy UMLS root node that all hierarchical vocabs will
        # link to
        umls_root_id = create_umls_root_node(source_node_file, tmpdir)
        sabs = get_sabs(session, required_sabs=sabs)

        tqdm_kwargs = dict(
            unit=" SABs",
            desc="[info] extracting nodes SABs",
            disable=not verbose,
            leave=False,
            total=len(sabs),
        )

        # Set up the trackers for the source and concept nodes
        cui_nodes = TRACKER_CLASS(utils.get_temp_file(dir=tmpdir)).open()
        source_nodes = TRACKER_CLASS(utils.get_temp_file(dir=tmpdir)).open()

        logger.info("extracting nodes...")

        # Will hold errors node/relationship mapping errors
        key_errors = []
        all_source_root_id = []
        all_cui_root_id = []
        # First extract the nodes for each SAB, so we have IDs for all expected
        # nodes
        for i in tqdm(sabs, **tqdm_kwargs):
            # Now start building, first generate source root nodes for the
            # SAB
            source_no_root, source_root_id = extract_source_root_nodes(
                session, source_nodes, i
            )
            all_source_root_id.append(source_root_id)

            # Now all the other non-source nodes
            extract_source_nodes(session, source_nodes, i, verbose=verbose)

            # CUI root nodes for the SAB
            cui_no_root, cui_root_id = extract_cui_root_nodes(
                session, cui_nodes, i
            )
            all_cui_root_id.append(cui_root_id)

            # Now all the other non-source nodes
            extract_cui_nodes(session, cui_nodes, i, verbose=verbose)

        logger.info("writing nodes...")

        # Write all the source nodes to file
        for row in source_nodes.fetch_nodes():
            source_node_file.write("{0}\n".format(_DELIMITER.join(row)))

        # Write all the concept nodes to file
        for row in cui_nodes.fetch_nodes():
            cui_node_file.write("{0}\n".format(_DELIMITER.join(row)))

        logger.info("building relationships...")

        tqdm_kwargs = dict(
            unit=" SABs",
            desc="[info] extracting relationship SABs",
            disable=not verbose,
            leave=False,
            total=len(sabs),
        )

        # Now build all the relationships
        for i, sid, cid in tqdm(
            zip(sabs, all_source_root_id, all_cui_root_id), **tqdm_kwargs
        ):
            # Build the source->cui relationships
            build_source_cui_rel(
                session,
                source_cui_rel,
                source_nodes,
                cui_nodes,
                i,
                verbose=verbose,
            )

            # Build the hierarchical relationship between source nodes
            error = build_source_rel(
                session,
                source_rel,
                source_nodes,
                i,
                sid,
                umls_root_id,
                sabs,
                verbose=verbose,
            )
            key_errors.extend(error)

            # Build the hierarchical relationship between concept nodes
            error = build_cui_rel(
                session,
                cui_rel,
                cui_nodes,
                i,
                cid,
                umls_root_id,
                sabs,
                verbose=verbose,
            )
            key_errors.extend(error)

        logger.info("building cross-source mappings...")

        # Get the source mapping data, all SABs are processed at once since
        # the mappings can cut across SABs
        build_source_mappings(
            session, source_map_rel, source_nodes, sabs=sabs, verbose=verbose
        )
    except Exception:
        # Close/delete all files if we error out
        for i in [
            source_node_file,
            cui_node_file,
            source_map_rel,
            source_cui_rel,
            source_rel,
            cui_rel,
        ]:
            try:
                i.close()
                os.unlink(i.name)
            except (FileNotFoundError, AttributeError):
                # Attribute error is for the IdTracker
                pass
        raise
    finally:
        # Make sure our IdTrackers are closed and deleted when finished
        for i in [source_nodes, cui_nodes]:
            try:
                i.close()
                os.unlink(i.dbfile)
            except (FileNotFoundError, AttributeError):
                # Attribute error is for the IdTracker
                pass
    return key_errors


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def init_source_nodes(outdir):
    """Initialise a source node file. This opens and writes the header.

    Parameters
    ----------
    outdir : `str`
        The output directory path.

    Returns
    -------
    source_node_file : `File`
        A source node file to write to.
    """
    nodes_file = open(os.path.join(outdir, SOURCE_NODE_BASENAME), "wt")
    nodes_file.write("{0}\n".format(_DELIMITER.join(SOURCE_NODE_HEADER)))
    return nodes_file


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def init_cui_nodes(outdir):
    """Initialise a CUI node file. This opens and writes the header.

    Parameters
    ----------
    outdir : `str`
        The output directory path.

    Returns
    -------
    cui_node_file : `File`
        A cui node file to write to.
    """
    nodes_file = open(os.path.join(outdir, CUI_NODE_BASENAME), "wt")
    nodes_file.write("{0}\n".format(_DELIMITER.join(CUI_NODE_HEADER)))
    return nodes_file


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def init_rel_file(outdir, outfile):
    """Initialise general relationship file. This opens and writes the header.

    Parameters
    ----------
    outdir : `str`
        The output directory path.
    outfile : `str`
        The base name of the output file.

    Returns
    -------
    relationship_file : `File`
        A relationship file to write to.
    """
    rel_file = open(os.path.join(outdir, outfile), "wt")
    rel_file.write("{0}\n".format(_DELIMITER.join(GENERAL_REL_HEADER)))
    return rel_file


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_umls_root_node(node_file, tmpdir):
    """Create a UMLS root node for all the hierarchical vocabs.

    Parameters
    ----------
    node_file : `File`
        A node file to write the root UMLS node to.
    tmpdir : `str`
        A temp location to use for a temporary IdTracker.

    Returns
    -------
    umls_root_id : `int`
        The node :ID for the UMLS root node.
    """
    try:
        umls_root = TRACKER_CLASS(utils.get_temp_file(dir=tmpdir)).open()
        umls_root_id = umls_root.add_node(
            UMLS_SAB, UMLS_ROOT_LABEL, UMLS_ROOT_LABEL, UMLS_ROOT_LABEL
        )
        for row in umls_root.fetch_nodes():
            node_file.write("{0}\n".format(_DELIMITER.join(row)))
    finally:
        umls_root.close()
        try:
            os.unlink(umls_root.dbfile)
        except AttributeError:
            # Attribute error is for the IdTracker
            pass

    # Add a dummy node that will act as the root of all hierarchies
    return umls_root_id


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def extract_source_root_nodes(session, idm, sab):
    """Extract the source root nodes for the source vocab.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session object to issue queries.
    idm : `umls_tool.neo4j.extract.IdTracker` or \
    `umls_tool.neo4j.extract.IdTrackerShelve`
        The cache for the node IDs.
    sab : `umls_tools.neo4j.extract.Sab`
        The source vocabulary abbreviation that should be extracted.

    Returns
    -------
    no_root : `bool`
        Does the SAB have a root node, not all SABs are hierarchical, so some
        won't have a root node.
    root_node_id : `int` or `NoneType`
        The root node ID, or ``NoneType``, if the SAB has no root nodes.
    """
    sql = session.query(o.MrConso.CUI, o.MrConso.CODE, o.MrConso.STR).filter(
        and_(o.MrConso.CODE == f"V-{sab.name}", o.MrConso.TTY == ROOT_TTY)
    )

    no_root = False
    try:
        row = sql.one()
        source_root_id = idm.add_node(
            sab.name, row.CODE, row.STR, SOURCE_ROOT_LABEL
        )
    except NoResultFound:
        source_root_id = None
        no_root = True
    except MultipleResultsFound:
        raise
    return no_root, source_root_id  # , umls_source_root_id


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def extract_cui_root_nodes(session, idm, sab):
    """Extract the root nodes for the umls concepts.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session object to issue queries.
    idm : `umls_tool.neo4j.extract.IdTracker` or \
    `umls_tool.neo4j.extract.IdTrackerShelve`
        The cache for the node IDs.
    sab : `umls_tools.neo4j.extract.Sab`
        The source vocabulary abbreviation that should be extracted.

    Returns
    -------
    no_root : `bool`
        Does the SAB have a root node, not all SABs are hierarchical, so some
        won't have a root node.
    root_node_id : `int` or `NoneType`
        The root node ID, or ``NoneType``, if the SAB has no root nodes.
    """
    sql = session.query(o.MrConso.CUI, o.MrConso.STR).filter(
        and_(o.MrConso.CODE == f"V-{sab.name}", o.MrConso.TTY == ROOT_TTY)
    )

    no_root = False
    try:
        row = sql.one()
        root_id = idm.add_node(sab.name, row.CUI, row.STR, CONCEPT_ROOT_LABEL)
    except NoResultFound:
        root_id = None
        no_root = True
    except MultipleResultsFound:
        raise
    return no_root, root_id


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def extract_source_nodes(session, idm, sab, verbose=False):
    """Extract UMLS source vocabulary nodes.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session object to issue queries.
    idm : `umls_tool.neo4j.extract.IdTracker` or \
    `umls_tool.neo4j.extract.IdTrackerShelve`
        The cache for the node IDs.
    sab : `umls_tools.neo4j.extract.Sab`
        The source vocabulary abbreviation that should be extracted.
    verbose : `bool`, optional, default: `False`
        Should a progress monitor be displayed.
    """
    sql = (
        session.query(o.MrConso.CODE, o.MrConso.SAB, o.MrConso.STR)
        .filter(
            and_(
                o.MrConso.LAT == ENGLISH_LANGUAGE,
                o.MrConso.CODE != NOCODE,
                o.MrConso.TTY != ROOT_TTY,
                o.MrConso.SAB == sab.name,
            )
        )
        .order_by(o.MrConso.CODE)
    )

    tqdm_kwargs = dict(
        unit=" rows",
        desc=f"[info] processing {sab.name}",
        disable=not verbose,
        leave=False,
        # total=count
    )

    for idx, row in enumerate(tqdm(sql, **tqdm_kwargs), 1):
        idm.add_node(row.SAB, row.CODE, row.STR, SOURCE_LABEL)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def extract_cui_nodes(session, idm, sab, verbose=False):
    """Extract UMLS concept nodes.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The SQLAlchemy session object to issue queries.
    idm : `umls_tool.neo4j.extract.IdTracker` or \
    `umls_tool.neo4j.extract.IdTrackerShelve`
        The cache for the node IDs.
    sab : `umls_tools.neo4j.extract.Sab`
        The source vocabulary abbreviation that should be extracted.
    verbose : `bool`, optional, default: `False`
        Should a progress monitor be displayed.
    """
    # Get the top ranked CUI string
    sql = (
        session.query(
            o.MrConso.CUI,
            o.UniqueUmlsCui.STR,  # o.MrConso.STR
        )
        .join(o.UniqueUmlsCui, o.MrConso.CUI == o.UniqueUmlsCui.CUI)
        .filter(
            and_(
                # Deliberately not using no code here as the same CUI can have
                # a code and NO_CODE for different AUIs
                o.MrConso.LAT == ENGLISH_LANGUAGE,
                o.MrConso.TTY != ROOT_TTY,
                o.MrConso.SAB == sab.name,
            )
        )
        .group_by(o.MrConso.CUI)
    )
    # count = sql.count()

    tqdm_kwargs = dict(
        unit=" rows",
        desc=f"[info] processing {sab.name}",
        disable=not verbose,
        leave=False,
        # total=count
    )

    for idx, row in enumerate(tqdm(sql, **tqdm_kwargs), 1):
        idm.add_node(CUI_TYPE, row.CUI, row.STR, CONCEPT_LABEL)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_source_mappings(
    session, outfile, source_nodes, sabs=None, verbose=False
):
    """Build a flat file from the UMLS MRCONSO table detailing the
    source->source mappings within the requested SABs.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outfile : `File`
        The output file to write to.
    source_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between source node codes/sabs and their IDs
    sabs : `list` of `umls_tools.neo4j.extract.Sab`, optional, \
    default: `NoneType`
        Any required SABs, the default is for all of them.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    if sabs is None:
        sabs = get_sabs(session)

    # Extract the SAB names
    sab_names = [i.name for i in sabs]

    sql = (
        session.query(o.MrConso.CUI, o.MrConso.CODE, o.MrConso.SAB)
        .filter(
            and_(
                o.MrConso.LAT == ENGLISH_LANGUAGE,
                o.MrConso.CODE != NOCODE,
                o.MrConso.TTY != ROOT_TTY,
            )
        )
        .order_by(o.MrConso.CUI)
    )

    tqdm_kwargs = dict(
        unit=" unfiltered rows",
        disable=not verbose,
        desc="[info] processing",
        leave=False,
    )

    curr_cui = None
    map_buffer = set([])
    for row in tqdm(sql, **tqdm_kwargs):
        # Moved onto another cui
        if curr_cui != row.CUI:
            if len(map_buffer) > 0:
                for left, right in _get_combinations(map_buffer):
                    outfile.write(
                        "{0}\n".format(
                            _DELIMITER.join(
                                [left, right, SOURCE_MAPPING_REL_NAME]
                            )
                        )
                    )
            map_buffer = set([])
            curr_cui = row.CUI
        if row.SAB in sab_names:
            map_buffer.add(source_nodes.get_node_id(row.SAB, row.CODE))

    # If there is some stuff still in the buffer
    if len(map_buffer) > 0:
        for left, right in _get_combinations(map_buffer):
            outfile.write(
                "{0}\n".format(
                    _DELIMITER.join([left, right, SOURCE_MAPPING_REL_NAME])
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_source_cui_rel(
    session, outfile, source_nodes, cui_nodes, sab, verbose=False
):
    """Build a mapping (edge) between the source/sab nodes and the concept
    nodes. To indicate that the source maps to a concept.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outfile : `File`
        The output file to write to.
    source_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between source node codes/sabs and their IDs
    cui_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between cui node codes/sabs and their IDs
    sab : `umls_tools.neo4j.extract.Sab`
        The required SAB.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    sql = (
        session.query(o.MrConso.CUI, o.MrConso.CODE, o.MrConso.SAB)
        .filter(
            and_(
                o.MrConso.SAB == sab.name,
                o.MrConso.LAT == ENGLISH_LANGUAGE,
                o.MrConso.CODE != NOCODE,
                o.MrConso.TTY != ROOT_TTY,
            )
        )
        .group_by(o.MrConso.CODE, o.MrConso.CUI)
    )

    tqdm_kwargs = dict(
        unit=" rows",
        disable=not verbose,
        desc=f"[info] processing {sab.name}",
        leave=False,
    )

    for row in tqdm(sql, **tqdm_kwargs):
        sid = str(source_nodes.get_node_id(row.SAB, row.CODE))
        cid = str(cui_nodes.get_node_id(CUI_TYPE, row.CUI))
        outfile.write(
            "{0}\n".format(_DELIMITER.join([sid, cid, SOURCE_CUI_REL_NAME]))
        )
        outfile.write(
            "{0}\n".format(_DELIMITER.join([cid, sid, CUI_SOURCE_REL_NAME]))
        )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_source_rel(
    session,
    outfile,
    source_nodes,
    sab,
    source_root_id,
    umls_root_id,
    all_sabs,
    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.
    outfile : `File`
        The output file to write to.
    source_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between source node codes/sabs and their IDs
    cui_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between cui node codes/sabs and their IDs
    sab : `umls_tools.neo4j.extract.Sab`
        The required SAB.
    source_root_id : `int`
        The node :ID for the source root node.
    umls_root_id : `int`
        The node :ID for the UMLS root node.
    all_sabs : `list` of `umls_tools.neo4j.extract.Sab`
        All the SABs that are being extracted, this is so that we can cross
        check nodes that are not found in the node store to see if we should
        have them or if they are errors.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    mrcons1 = aliased(o.MrConso)
    mrcons2 = aliased(o.MrConso)

    sql = (
        session.query(
            mrcons1.CODE.label("m1code"),
            mrcons2.CODE.label("m2code"),
            o.MrRel.REL,
            mrcons1.SAB.label("m1sab"),
            mrcons2.SAB.label("m2sab"),
            mrcons1.TTY.label("m1tty"),
            mrcons2.TTY.label("m2tty"),
        )
        .join(mrcons1, mrcons1.AUI == o.MrRel.AUI1)
        .join(mrcons2, mrcons2.AUI == o.MrRel.AUI2)
        .filter(
            and_(
                o.MrRel.SAB == sab.name,
                mrcons1.LAT == ENGLISH_LANGUAGE,
                mrcons2.LAT == ENGLISH_LANGUAGE,
            )
        )
        .group_by(mrcons1.CODE, mrcons2.CODE, o.MrRel.REL)
    )

    tqdm_kwargs = dict(
        unit=" rows",
        disable=not verbose,
        leave=False,
        desc=f"[info] extract source->source rels ({sab.name})",
    )

    sab_names = [i.name for i in all_sabs]
    root_codes = dict([(i.code, i.name) for i in all_sabs])

    if source_root_id is not None:
        outfile.write(
            "{0}\n".format(
                _DELIMITER.join(
                    [str(umls_root_id), str(source_root_id), PAR_REL]
                )
            )
        )
        outfile.write(
            "{0}\n".format(
                _DELIMITER.join(
                    [str(source_root_id), str(umls_root_id), CHD_REL]
                )
            )
        )
    key_errors = []
    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
        ):
            try:
                outfile.write(
                    "{0}\n".format(
                        _DELIMITER.join(
                            [
                                str(
                                    source_nodes.get_node_id(
                                        row.m1sab, row.m1code
                                    )
                                ),
                                str(
                                    source_nodes.get_node_id(
                                        row.m2sab, row.m2code
                                    )
                                ),
                                row.REL,
                            ]
                        )
                    )
                )
            except KeyError:
                # For some reason we could not get a node ID, this could be a
                # legit reason, if the SAB for the relationship has not been
                # requested or if the code is a root code, these do not have
                # the SABs of the SAB that they are the root for, so we will
                # have to adapt to the correct SAB, I might change this in
                # future

                # First test to see what SAB/CODE combination did not match in
                # the relationship mapping
                m1_ok = (row.m1sab, row.m1code) in source_nodes.nodes
                m2_ok = (row.m2sab, row.m2code) in source_nodes.nodes

                # If it is a true error then this will eventually be appended
                # to the the errors
                key_error_info = [
                    row.m1code,
                    m1_ok,
                    row.m2code,
                    m2_ok,
                    sab.name,
                    "SOURCE",
                ]

                # If neither has matched, then it is a true error, so we log it
                # and move on
                if m1_ok is False and m2_ok is False:
                    key_errors.append(key_error_info)
                    continue

                # If the first (left hand in the relationship) code did not
                # match
                if m1_ok is False:
                    # Is the left hand code a root code?
                    if row.m1tty == ROOT_TTY and row.m1code in root_codes:
                        # Map the SAB to the SAB that the code is root for
                        left_sab = root_codes[row.m1code]
                        right_sab = row.m2code
                    elif row.m1sab not in sab_names:
                        # Not a root but also not a requested SAB so we would
                        # not expect to have it so we can relax and move on.
                        continue
                elif m2_ok is False:
                    # Right hand code did not match, is it a root?
                    if row.m2tty == ROOT_TTY and row.m2code in root_codes:
                        left_sab = row.m1code
                        right_sab = root_codes[row.m2code]
                    elif row.m2sab not in sab_names:
                        continue

                try:
                    # If we get here we want to re-query our root node
                    outfile.write(
                        "{0}\n".format(
                            _DELIMITER.join(
                                [
                                    str(
                                        source_nodes.get_node_id(
                                            left_sab, row.m1code
                                        )
                                    ),
                                    str(
                                        source_nodes.get_node_id(
                                            right_sab, row.m2code
                                        )
                                    ),
                                    row.REL,
                                ]
                            )
                        )
                    )
                except KeyError:
                    # If we get here, then we have an error
                    key_errors.append(key_error_info)
    return key_errors


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_cui_rel(
    session,
    outfile,
    cui_nodes,
    sab,
    cui_root_id,
    umls_root_id,
    all_sabs,
    verbose=False,
):
    """Build a flat file from the UMLS cui->cui relationships. This may
    take some time to run.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    outfile : `File`
        The output file to write to.
    source_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between source node codes/sabs and their IDs
    cui_nodes : `umls_tools.neo4j.extract.IdTrackerShelve`
        The mappings between cui node codes/sabs and their IDs
    sab : `umls_tools.neo4j.extract.Sab`
        The required SAB.
    cui_root_id : `int`
        The node :ID for the CUI root node.
    umls_root_id : `int`
        The node :ID for the UMLS root node.
    all_sabs : `list` of `umls_tools.neo4j.extract.Sab`
        All the SABs that are being extracted, this is so that we can cross
        check nodes that are not found in the node store to see if we should
        have them or if they are errors.
    verbose : `bool`, optional, default: `False`
        Show a progress monitor.
    """
    sql = (
        session.query(o.MrRel.CUI1, o.MrRel.CUI2, o.MrRel.REL)
        .filter(o.MrRel.SAB == sab.name)
        .group_by(o.MrRel.CUI1, o.MrRel.CUI2, o.MrRel.REL)
    )

    tqdm_kwargs = dict(
        unit=" rows",
        disable=not verbose,
        leave=False,
        desc=f"[info] extract cui->cui rels ({sab.name})",
    )

    if cui_root_id is not None:
        outfile.write(
            "{0}\n".format(
                _DELIMITER.join([str(umls_root_id), str(cui_root_id), PAR_REL])
            )
        )
        outfile.write(
            "{0}\n".format(
                _DELIMITER.join([str(cui_root_id), str(umls_root_id), CHD_REL])
            )
        )

    sab_names = [i.name for i in all_sabs]
    # root_codes = dict([(i.code, i.name) for i in all_sabs])

    # This will hold rows that may be errors but will be rechecked after the
    # loop below has run
    provisional_errors = []

    # This holds things that are certainly errors
    key_errors = []
    for row in tqdm(sql, **tqdm_kwargs):
        if (
            row.CUI1 != row.CUI2
            and row.CUI1 is not None
            and row.CUI2 is not None
        ):
            try:
                outfile.write(
                    "{0}\n".format(
                        _DELIMITER.join(
                            [
                                str(cui_nodes.get_node_id(CUI_TYPE, row.CUI1)),
                                str(cui_nodes.get_node_id(CUI_TYPE, row.CUI2)),
                                row.REL,
                            ]
                        )
                    )
                )
            except KeyError:
                cui1_ok = (CUI_TYPE, row.CUI1) in cui_nodes.nodes
                cui2_ok = (CUI_TYPE, row.CUI2) in cui_nodes.nodes

                # If both are false it is a certain error
                if cui1_ok is False and cui2_ok is False:
                    key_errors.append(
                        [row.CUI1, cui1_ok, row.CUI2, cui2_ok, sab.name, "CUI"]
                    )
                else:
                    provisional_errors.append((row, cui1_ok, cui2_ok))

    # Now we loop through the provisional errors and test against MRSONSO
    for row, cui1_ok, cui2_ok in provisional_errors:
        error = False
        if cui1_ok is False:
            error = _query_cui(session, row.CUI1, sab_names)
        elif cui2_ok is False:
            error = _query_cui(session, row.CUI2, sab_names)

        if error is True:
            key_errors.append(
                [row.CUI1, cui1_ok, row.CUI2, cui2_ok, sab.name, "CUI"]
            )

    return key_errors


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _query_cui(session, cui, sab_names):
    """Query a CUI against MRCONSO, to see if any of it's SABs are required,
    if so then that is an error.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session to issue queries with.
    cui : `str`
        The UMLS concept identifier to test.
    sab_names : `list` of `str`
        All the requested SAB names.

    Returns
    -------
    is_error : `bool`
        ``True`` is yes the SAB names associated with the concept are found,
        ``False`` if not.
    """
    sql = session.query(o.MrConso.SAB).filter(o.MrConso.CUI == cui)
    for i in sql:
        if i.SAB in sab_names:
            return True
    return False


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _get_combinations(mappings):
    """Get the mapping pairs (both sides) if the nodes that have a CUI mapping.

    Parameters
    ----------
    mappings : `list` of `int`
        The mappings to split into pairwise combinations.

    Yields
    ------
    left_side : `str`
        The left side of the combination map.
    right_side : `str`
        The right side of the combination map.
    """
    for i, j in combinations(mappings, 2):
        yield str(i), str(j)
        yield str(j), str(i)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_sabs(session, required_sabs=None):
    """Get all the SABs that are required to be extracted.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session to execute queries.
    required_sabs : `list` of `str`, optional, default: `NoneType`
        The SABs that are required to be extracted. If this is ``NoneType``
        then all of them will be extracted.

    Returns
    -------
    sabs : `list` of `umls_tools.neo4j.extract.Sab`
        A named tuple for each matched SAB.
    """
    all_sabs = []
    seen_id_name = set()
    seen_file_suffix = set()
    sql = (
        session.query(o.MrSab)
        .filter(o.MrSab.LAT == ENGLISH_LANGUAGE)
        .group_by(o.MrSab.RSAB)
    )

    for row in sql:
        if required_sabs is not None and row.RSAB not in required_sabs:
            continue

        sab = re.sub(r"[\._-]", "", row.RSAB)
        sab = sab.lower()
        sab = sab[0].upper() + sab[1:]

        # This is the code that the SAB will be known by in the UMLS
        code = "V-{row.RSAB}"
        s = Sab(name=row.RSAB, id_name=sab, file_suffix=sab.lower(), code=code)

        if s.id_name in seen_id_name:
            raise ValueError(f"ID name already used: {s.id_name}")
        if s.file_suffix in seen_file_suffix:
            raise ValueError(f"file suffix already used: {s.file_suffix}")

        seen_file_suffix.add(s.file_suffix)
        seen_id_name.add(s.id_name)
        all_sabs.append(s)
    return all_sabs


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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()
