"""
Classes to interact with the UMLS Neo4j graph
"""

from py2neo import Node
import pprint
pp = pprint.PrettyPrinter()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class UmlsInterface(object):
    """
    An interface for doing queries against the UMLS neo4j graph, the queries
    and thier results provide the backbone data for the similarity measures
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, graph, sab=[], filter_longest_term=False):
        """
        Initialise

        Parameters
        ----------

        graph : :obj:`Graph`
            A py2neo Graph object providing access to the UMLS graph
        sab : :obj:`list`
            The SAB (ontologies) that the nodes must belong to. Note that
            this is an OR not an AND
        filter_longest_term : :obj:`bool`
            There are many terms associated with each Concept node, this
            filters then to return the one with the longest string length
        """
        self.graph = graph
        self.sab = sab
        self._sab_query_str = ",".join(["'{}'".format(s) for s in self.sab])

        # Determine the filter method based on filter_longest_term
        self._filter_method = self._filter_none
        if filter_longest_term is True:
            self._filter_method = self._filter_all_longest_terms

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_root_concepts(self):
        """
        Return the root nodes for all, or selected SABs
        """

        # Rather than building an SAB specific query, I will grab them all and
        # filter in python
        cql = "MATCH (cui1:root_node) RETURN cui1"

        c = self.graph.run(cql)

        if self.sab is None or len(self.sab) == 0:
            for i in c:
                yield self._filter_method(i[0])
        else:
            sab_matches = [i[0] for i in c for j in i[0].labels()
                           if j in self.sab]
            for i in sab_matches:
                yield self._filter_method(i)
            # pp.pprint(i)
            # print(i[0].labels())

            # for l in i[0].labels():
            #     print(l)
            # yield self._filter_method(i)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_child_concepts(self, cui, dist_start=1, dist_end=1):
        """
        When given a CUI get any child concepts of the CUI

        Parameters
        ----------

        cui : :obj:`str`
            The UMLS Concept identifier (CUI) to search for
        dist_start : :obj:`int` or `None`
            return Nodes that are at least this edge distance away from ciu.
             Set to `None` for no distance limits
        dist_end : :obj:`int` or `None`
            return nodes that are up to this distance away from cui. Set to
            `None` for no distance limits
        """

        cql = """MATCH (cui1:Concept)-[r:CHD*{LIMITS}]->(cui2:Concept)
              WHERE cui1.cui = $cui {{SAB1}} RETURN cui1, r, cui2"""

        # Add the relationship distance limits
        cql = self._add_rel_dist(cql, dist_start, dist_end)

        # Add any SABs to the query
        cql = self._add_sabs(cql, add=True, node_name=("cui1", "SAB1"))
        cql = self._add_sabs(cql, node_name=("cui2", "SAB1"))

        # Run the query and get the cursor
        c = self.graph.run(cql,
                           parameters={'cui': cui})
        for i in c:
            yield self._filter_method(i)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_shortest_path(self, cui1, cui2):
        """
        Find a single shortest path between 2 cuis. The SABs are applied at
        an interface level

        Parameters
        ----------

        cui1 : :obj:`str`
            The first concept to match
        cui2 : :obj:`str`
            The second concept to match
        """
        cql = """
        MATCH p=shortestPath((cui1:Concept)-[:CHD|:PAR*]->(cui2:Concept))
        WHERE cui1.cui = $cui1 AND cui2.cui = $cui2
        """

        if self.sab is not None and len(self.sab) > 0:
            sab = """AND ALL(x IN nodes(p)
                           WHERE ANY(item in x.sab
                                     WHERE item IN [{}]))""".format(
                                         self._sab_query_str)
            cql = "{} {}".format(cql, sab)
        cql = "{}  RETURN nodes(p)".format(cql)

        # Run the query and get the cursor
        c = self.graph.run(cql,
                           parameters={'cui1': cui1, 'cui2': cui2})

        results = c.next()

        self._filter_method(results[0])
        return results[0]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_shortest_par_path(self, cui1, cui2):
        """
        Find a single shortest path between 2 cuis. The SABs are applied at
        an interface level

        Parameters
        ----------

        cui1 : :obj:`str`
            The first concept to match
        cui2 : :obj:`str`
            The second concept to match
        """
        cql = """
        MATCH p=shortestPath((cui1:Concept)-[:PAR*]->(cui2:Concept))
        WHERE cui1.cui = $cui1 AND cui2.cui = $cui2
        """

        if self.sab is not None and len(self.sab) > 0:
            sab = """AND ALL(x IN nodes(p)
                           WHERE ANY(item in x.sab
                                     WHERE item IN [{}]))""".format(
                                         self._sab_query_str)
            cql = "{} {}".format(cql, sab)
        cql = "{}  RETURN nodes(p)".format(cql)

        # Run the query and get the cursor
        c = self.graph.run(cql,
                           parameters={'cui1': cui1, 'cui2': cui2})

        results = c.next()

        self._filter_method(results[0])
        return results[0]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_path_to_root(self, cui):
        """
        Return the shortest path to root, note depending on the SABs used this
        may be different
        """

        all_root_paths = {}
        # Get the root concepts for the available ontologies
        for i in self.get_root_concepts():
            root_path = []
            for depth, concept in enumerate(
                    self.get_shortest_par_path(cui, i['cui'])):
                try:
                    root_type = self.get_root_label(concept)
                except KeyError:
                    pass
                root_path.append((concept['cui'], depth))
            # all_root_paths.setdefault(root_type, [])
            all_root_paths[root_type] = root_path

        # Return the root paths subset by ontology
        return all_root_paths

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_lcs(self, cui1, cui2):
        """
        Find a least common subsummer between 2 cuis. The SABs are applied at
        an interface level

        Parameters
        ----------

        cui1 : :obj:`str`
            The first concept to match
        cui2 : :obj:`str`
            The second concept to match
        """
        # cql = """

        # rpsets = {}
        # rporder = {}

        cui1_root_paths = self.get_path_to_root(cui1)
        cui2_root_paths = self.get_path_to_root(cui2)

        root_paths = {}
        for k, v in cui1_root_paths.items():
            root_paths.setdefault(k, [])
            root_paths[k].append(v)

        for k, v in cui2_root_paths.items():
            root_paths.setdefault(k, [])
            root_paths[k].append(v)

        # Convert to sets
        root_sets = {}
        for k, v in root_paths.items():
            root_sets.setdefault(k, [])
            for l in v:
                root_sets[k].append(set([e[0] for e in l]))

        common = {}
        for k, v in root_sets.items():
            common[k] = set.intersection(*v)

        distances = {}
        # Loop through the different ontologies
        for k, v in root_paths.items():
            distances.setdefault(k, [])
            # The paths in v reflect the path to root for
            # each of the supplied concepts
            for conceptno, path in enumerate(v):
                for concept in path:
                    if concept[0] in common[k]:
                        lcs_dist_from_concept = concept[1]
                        lcs_dist_from_root = len(path) - concept[1]
                        concept_dist_from_root = len(path)
                        distances[k].append((concept[0],
                                             lcs_dist_from_root,
                                             lcs_dist_from_concept,
                                             concept_dist_from_root))
                        break

        return distances

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_shortest_path_len(self, cui1, cui2):
        """
        Find a the length of the shortest path between 2 cuis.
        The SABs are applied at an interface level

        Parameters
        ----------

        cui1 : :obj:`str`
            The first concept to match
        cui2 : :obj:`str`
            The second concept to match
        """
        return len(self.get_shortest_path(cui1, cui2))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_root_label(self, node):
        """
        Get the root node label
        """

        labels = node.labels()

        # Make sure it is a concept
        if "Concept" not in labels or "root_node" not in labels:
            raise KeyError("node is not a root_node")

        return tuple([i for i in labels if i not in ['Concept', 'root_node']])
        # for i in node.labels():
        #     print(i)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _add_rel_dist(self, cql, dist_start, dist_end):
        """
        Add relationship distances to the cypher query
        """

        dist_str = ""
        if dist_start > dist_end:
            raise ValueError("relationship end distance > start distance")
        elif dist_start is not None and dist_end is not None:
            dist_str = "{}..{}".format(dist_start, dist_end)

        return cql.format(LIMITS=dist_str)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _add_sabs(self, cql, add=False, node_name=("cui1", "SAB")):
        """
        Adds SAB filters to the query
        """

        sab = ""
        # ANY(item IN {}.sab WHERE item IN [{}])
        if self.sab is not None and len(self.sab) > 0:
            sab = " AND ANY(item IN {}.sab WHERE item IN [{}])".format(
                node_name[0],
                ",".join(["'{}'".format(s) for s in self.sab]))
        # print(sab)
        if add is True:
            sab = sab + "{SAB1}"

        kwargs = {node_name[1]: sab}
        return cql.format(**kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _filter_all_longest_terms(self, record):
        """
        Get the longest term from a Concept node
        """
        # print(record.__class__.__name__)

        if isinstance(record, Node):
            return self._filter_longest_term(record)

        for i in record:
            # print(i.__class__.__name__)
            # print(i.labels)
            i = self._filter_longest_term(i)
            return record
        # else:
        #     return self._filter_longest_term(record)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _filter_longest_term(self, node):
        """
        Get the longest term from a Concept node
        """
        try:
            max_len = 0
            long_term = ""
            for t in node['term']:
                if len(t) > max_len:
                    long_term = t
                    max_len = len(t)
                node['term'] = [long_term]
        except TypeError:
            pass

        return node

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _filter_none(self, record):
        """
        A null method that does not do any filtering
        """
        return record


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class UmlsSimilarity(object):
    """
    A class for doing similarity measures on a UMLS Neo4j graph
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, interface):
        """
        Initialise

        Parameters
        ----------

        interface : :obj:`UmlsInterface`
            An interface object to do the queries against the graph
        """
        self.interface = interface

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_path_sim(self, cui1, cui2):
        """
        Determine the similarity between 2 concepts using the PATH method. This
        is the recipricol of the number node count
        """
        return 1.0/self.interface.get_shortest_path_len(cui1, cui2)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_wup_sim(self, cui1, cui2):
        """
        Determine the similarity between 2 concepts using the WU and Palmer
        method. The Wu & Palmer measure calculates relatedness by considering
        the depths of the two concepts in the UMLS, along with the depth of the
        LCS. The formula is score = 2*depth(lcs) / (depth(s1) + depth(s2)).
        This means that 0 < score <= 1. The score can never be zero because the
        depth of the LCS is never zero (the depth of the root of a taxonomy is
        one). The score is one if the two input concepts are the same.
        """

        # Now calculate the similarity, this is done per ontology
        # for k, v in distances:
        distances = self.interface.get_lcs(cui1, cui2)

        scores = []
        for ont, lcs_values in distances.items():
            lcs = max(lcs_values[0][1], lcs_values[1][1])
            score = (2 * float(lcs)) / (float(lcs_values[0][3]) +
                                        float(lcs_values[1][3]))
            scores.append((ont, cui1, cui2, score))
        return scores
