"""Functions and classes to query the UMLS."""

from umls_tools import orm as o
import re


NO_MATCH = "N"
"""The code for no matching (`str`)
"""
EXACT_MATCH = "E"
"""The code for exact matching (`str`)
"""
SUBSTRING_MATCH = "S"
"""The code substring matching matching (`str`)
"""
TOKEN_STRING_MATCH = "T"
"""The code for tokenised word matching (`str`)
"""
METAMAP_MATCH = "M"
"""The code for a metamap match (`str`)
"""


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CuiMapper(object):
    """Attempt to map CUIs in short text strings.

    This is not designed as a named entity tagger, rather to look for whole
    term matches with a fall map to Metamap.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session object to issue queries.
    lats : `list` of `str`, optional, default: `NoneType`
        The restrict to AUIs with any of the specific languages. If ``None``
        this defaults to ``['ENG']``.
    metamap_server : `umls_tools.metamap.MetamapServer`, optional, \
    default: `NoneType`
        A server object to issue queries against. If ``NoneType`` then  Metamap
        will not be used. It is assumed that the metamap arguments have already
        been passed when the object was initialised.
    like_format : `str`, optional, default: `%{0}%`
        The format of the SQL LIKE query, this must be a format string i.e.
        ``{0}``. This defaults to any characters before or after the string.
    token_delim : `str`, optional, default: `[-._]`
        Remove character sets from the search string and replace with
        white-space (which is then used for tokenising).

    Raises
    ------
    ValueError
        If the ``metamap_server`` is set but the ``metamap_args`` are not set.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(
        self,
        session,
        lats=None,
        metamap_server=None,
        like_format="%{0}%",
        token_delim=r"[-._]",
    ):
        self.token_delim = token_delim
        self.like_format = like_format
        if lats is None:
            lats = ["ENG"]
        self.lats = lats
        self.metamap_server = metamap_server
        self.session = session

        self.base_match = session.query(o.MrConso)

        # if lats is not None and len(lats) > 0:
        #     if len(lats) == 1:
        #         self.base_match = self.base_match.filter(
        #             o.MrConso.LAT == self.lats[0]
        #         )
        #     else:
        #         self.base_match = self.base_match.filter(
        #             o.MrConso.LAT.in_(self.lats)
        #         )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def map_term(self, search_term):
        """Perform a CUI mapping sequence on a ``search_term``.

        Parameters
        ----------
        search_term : `str`
            The string to search for.

        Returns
        -------
        cuis : `list` of `tuple`
            Matching cuis, each tuple has ``[0]`` ``umls_tools.orm.MrConso``
            (the cui match). ``[1]`` the start position of the match in the
            ``search_term``. ``[2]`` the end position of the match in the
            search term.
        match_type : `str`
            The match type, so exact (``E``), substring (``S``),
            tokenised (``T``) or Metamap (``M``). If start/end are -1 this is
            an error start where the substring can't be identified.
        """
        results = self.match_term(search_term)
        if len(results) > 0:
            return results, EXACT_MATCH

        results = self.search_term(search_term)
        if len(results) > 0:
            return results, SUBSTRING_MATCH

        if self.metamap_server is not None:
            results = self.match_metamap(search_term)
            if len(results) > 0:
                return results, METAMAP_MATCH
        else:
            results = self.match_tokens(search_term)
            if len(results) > 0:
                return results, TOKEN_STRING_MATCH

        return [], NO_MATCH

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_term(self, search_str):
        """Find an exact match to a search string.

        Parameters
        ----------
        search_str : `str`
            The string to search for.

        Returns
        -------
        cuis : `list` of `tuple`
            Matching cuis, each tuple has ``[0]`` ``umls_tools.orm.MrConso``
            (the cui match). ``[1]`` the start position of the match in the
            ``search_term``. ``[2]`` the end position of the match in the
            search term.
        """
        # Otherwise only extract the concept Ids
        return [
            (i, 0, len(search_str) - 1)
            for i in self.base_match.filter(o.MrConso.STR == search_str)
            if i.LAT in self.lats
        ]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def search_term(self, search_str):
        """Find a non-exact (like) match to a search string, so the string can
        have leading and trailing characters.

        Parameters
        ----------
        search_str : `str`
            The string to search for.

        Returns
        -------
        cuis : `list` of `tuple`
            Matching cuis, each tuple has ``[0]`` ``umls_tools.orm.MrConso``
            (the cui match). ``[1]`` the start position of the match in the
            ``search_term``. ``[2]`` the end position of the match in the
            search term. If start/end are -1 this is an error start where the
            substring can't be identified.
        """
        results = self.base_match.filter(
            o.MrConso.STR.like(self.like_format.format(search_str))
        ).all()

        search_str = search_str.lower()
        return_mapped = []
        for i in results:
            if i.LAT not in self.lats:
                continue
            try:
                start = i.STR.lower().index(search_str)
                end = start + len(search_str) - 1
            except ValueError:
                # for some reason the substring could not be ID'd, so we
                # indicate an error state
                start = -1
                end = -1

            return_mapped.append((i, start, end))
        return return_mapped

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_tokens(self, search_str):
        """Find substring matches to the search string.

        Parameters
        ----------
        search_str : `str`
            The string to search for.

        Returns
        -------
        cuis : `list` of `tuple`
            Matching cuis, each tuple has ``[0]`` ``umls_tools.orm.MrConso``
            (the cui match). ``[1]`` the start position of the match in the
            ``search_term``. ``[2]`` the end position of the match in the
            search term.

        Notes
        -----
        This searches the normalised words table. Delimiter characters are
        normalised to whitespace before tokenisation, then tokenised with
        ``str.split()`` so leading/trailing whitespace and repeated
        delimiters do not produce empty tokens. It will only return matches
        where all words have matched (but the order may be different), so it
        is conservative.
        """
        if self.token_delim is not None:
            search_str = re.sub(self.token_delim, " ", search_str)
        search_list = search_str.split()
        matching_cui = {}

        for i in search_list:
            q = self.session.query(o.MrXnwEng).filter(o.MrXnwEng.NWD == i)
            for j in q.all():
                try:
                    matching_cui[j.CUI].add(i)
                except KeyError:
                    matching_cui[j.CUI] = set([i])

        return_cuis = []
        for k, v in matching_cui.items():
            if len(v) == len(search_list):
                for i in self.base_match.filter(o.MrConso.CUI == k):
                    return_cuis.append((i, 0, len(search_str) - 1))
        return return_cuis

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_metamap(self, search_str):
        """Find Metamap matches matches to the search string.

        If you are using this you will need to set the ``metamap_server`` and
        ``metamap_args`` parameters.

        Parameters
        ----------
        search_str : `str`
            The string to search for.

        Returns
        -------
        cuis : `list` of `tuple`
            Matching cuis, each tuple has ``[0]`` ``umls_tools.orm.MrConso``
            (the cui match). ``[1]`` the start position of the match in the
            ``search_term``. ``[2]`` the end position of the match in the
            search term.

        Raises
        ------
        RuntimeError
            If Metamap is not available.
        """
        return_results = []
        try:
            for r in self.metamap_server.run(search_str):
                for u in r:
                    for mc in u.get_mapping_candidates():
                        q = (
                            self.session.query(o.MrConso)
                            .filter(o.MrConso.CUI == mc.cui.cui)
                            .filter(o.MrConso.STR == mc.cui.term)
                        )
                        for cui in q:
                            return_results.append((cui, mc.start, mc.end))
        except (ValueError, RuntimeError):
            # ValueError: Parse error: 'classical ) Hodgkin
            # lymphoma' is not in '(classical) Hodgkin lymphoma
            # RuntimeError: results is length='2', from
            # json.decoder.JSONDecodeError: Expecting value:
            # line 1 column 646 (char 645)
            pass
        except AttributeError as e:
            if self.metamap_server is None:
                raise RuntimeError("metamap not set") from e
            raise
        return return_results


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def match_code(session, code):
    """Get codes that match a clinical code. This searches the STR field in
    the ``MRCONSO`` table.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session object to perform the query.
    code : `str`
        The clinical ontology code to search for.

    Returns
    -------
    codes : `list` of `umls_tools.orm.MrConso`
        Matching codes.
    """
    return session.query(o.MrConso).filter(o.MrConso.CODE == code)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_all_sty(session):
    """Get all the unique semantic types from MRSTY.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session object to perform the query.

    Returns
    -------
    unique_sty : `umls_tools.orm.UniqueUmlsCui`
        The top ranked AUI for the CUI.

    Notes
    -----
    This requires that the ``UNIQUE_CUI``
    `table <https://cfinan.gitlab.io/umls-tools/scripts.html#umls-unique-cui>`_
    is in place.

    See also
    --------
    umls_tools.admin.unique_cui.create_unique_cui
    """
    return session.query(o.MrSty.TUI, o.MrSty.STY).group_by(o.MrSty.TUI)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_cui(session, cui_str):
    """Get all the concepts that match a CUI.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session object to perform the query.
    cui_str : `str`
        The UMLS concept ID value to search for.

    Returns
    -------
    """
    return session.query(o.MrConso).filter(o.MrConso.CUI == cui_str).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_unique_cui(session, cui_str):
    """Get all the top ranked AUI for a CUI. This queries the `UNIQUE_CUI`
    table.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session object to perform the query.
    cui_str : `str`
        The UMLS concept ID value to search for.

    Returns
    -------
    unique_cui : `umls_tools.orm.UniqueUmlsCui`
        The top ranked AUI for the CUI.

    Notes
    -----
    This requires that the ``UNIQUE_CUI``
    `table <https://cfinan.gitlab.io/umls-tools/scripts.html#umls-unique-cui>`_
    is in place.

    See also
    --------
    umls_tools.admin.unique_cui.create_unique_cui
    """
    return (
        session.query(o.UniqueUmlsCui)
        .filter(o.UniqueUmlsCui.CUI == cui_str)
        .one()
    )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_ranks(session):
    """Get a cache of the data in the MRRANK table.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        The session object to perform the query.

    Returns
    -------
    ranks : `dict`
        The keys of the dictionary are tuples of (``SAB``, ``TTY``) and the
        values are tuples of (``RANK``, ``SUPPRESS``)
    """
    ranks = dict()
    for row in session.query(o.MrRank):
        ranks[(row.SAB, row.TTY)] = (row.MRRANK_RANK, row.SUPPRESS)
    return ranks


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class UmlsQuery(object):
    """Query the UMLS when available as a relational database, this uses an
    SQLAlchemy interface.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        An SQLAlchemy session to complete the queries.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, session):
        self.session = session

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_code(self, code, sab=None):
        """Match on a code and optionally an SAB.

        Parameters
        ----------
        code : `str`
            The code to match on.
        sab : `list` of `str`, optional, default: `NoneType`
            Restrict the match to a specific sources.

        Returns
        -------
        matches : `list` of `umls_tools.orm.MrConso`
            Any code matches.
        """
        # put together the basics of the query
        q = self.session.query(o.MrConso).filter(o.MrConso.CODE == code)

        # If we are restricting to an SAB
        if sab is not None:
            q = q.filter(o.MrConso.SAB.in_(sab))

        return q.all()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_term(self, term, sab=None):
        """Attempt to exact match a term string in the UMLS.

        Parameters
        ----------
        term : `str`
            The term to match on.
        sab : `list` of `str`, optional, default: `NoneType`
            Restrict the match to a specific sources.

        Returns
        -------
        matches : `list` of `umls_tools.orm.MrConso`
            Any term matches.
        """
        # put together the basics of the query
        q = self.session.query(o.MrConso).filter(o.MrConso.STR == term)

        # If we are restricting to an SAB
        if sab is not None:
            q = q.filter(o.MrConso.SAB == sab)

        return q.all()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def match_uniq_term_from_term(self, term, sab=None):
        """This selects an entry from the unique CUI table based on an entry in
        MRCONSO that has been matched based on STR.

        Parameters
        ----------
        term : `str`
            The term to match on.
        sab : `list` of `str`, optional, default: `NoneType`
            Restrict the match to a specific sources.

        Returns
        -------
        matches : `list` of `umls_tools.orm.MrConso`
            Any term matches.

        Notes
        -----
        Use this if we are looking for an exact STR match in MRCONSO but want
        the unique term that represents it.
        """
        q = (
            self.session.query(o.UniqueUmlsCui)
            .join(o.MrConso, o.MrConso.CUI == o.UniqueUmlsCui.CUI)
            .filter(o.MrConso.STR == term)
        )

        # If we are restricting to an SAB
        if sab is not None:
            q = q.filter(o.MrConso.SAB == sab)

        return q.all()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def count_aui(self, cuis=None):
        """Get a total AUI count for the whole of the MRCONSO table or
        specific CUIs.

        Parameters
        ----------
        cuis : `list` of `str`
            The concepts to restrict to.

        Returns
        -------
        aui_count : `int`
            The number of AUIs.
        """
        if cuis is None:
            cuis = []
        q = self.session.query(o.MrConso.AUI)

        # If the user wants to restrict to certain cuis and not count the
        # whole table
        if len(cuis) > 0:
            q = q.filter(o.MrConso.CUI.in_(cuis))

        return q.count()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def count_cui(self):
        """Get the total CUI count.

        Returns
        -------
        cui_count : `int`
            The number of unique CUIs.
        """
        q = self.session.query(o.MrConso.CUI).distinct().count()
        return q
